- add loss_mask parameter to get_logprobs to decouple attention from loss masking
- DPO/GRPO strategies compose key-padding + causal mask before model forward
- prevents prompt tokens from being masked out of attention and missing causal masking
- use fmaf() for V-accumulation in scalar decode paths to reduce rounding
- delay scale multiplication to after dot-product in scalar prefill
- unify __expf/expf across MMA and scalar paths for consistent numerics
- harmonize divide-by-zero guards to 1e-20f
- add both absolute and relative error checks in standalone tests (atol=0.01, rtol=0.01)
- Merge 3 duplicated dispatch blocks into single attn_dispatchers.cuh
- Merge compute_num_splits from attn_utils.cuh into dispatcher header
- All dim3 grid/block declarations and <<<>>> launches are single-line
- Production .cu files (35-42 loc) only handle torch wrapping + pybind11
- Test files include dispatcher header directly, removing all #ifndef ASTRAI_NO_MMA duplication
- Introduce KernelTraits<HEAD_DIM, BC, WARPS, STAGES> compile-time config bundle, replacing scattered <KD, NC8, KT2, ...> template params
- Template all MMA and scalar kernels on IsCausal/HasMask bools to eliminate inner-loop runtime branches
- Dispatch to 4-path IsCausal/HasMask kernel variants at entry points based on p.causal_offset and p.use_mask
- Update standalone test files with new kernel signatures, add causal test cases
- Fix duplicate using bf16 in MMA kernels that include attn_mma_utils.cuh
- Replace .tolist() calls with _total_len in gather(); move _slot_len updates from per-layer write to once-per-step bind_tasks
- Use torch.as_tensor instead of torch.tensor in decode penalty history construction
- encode(str) single-thread, encode(List[str]) Rust parallel encode_batch
- apply_chat_template accepts single Messages or List[Messages] for batch
- add Message/Messages type aliases at module level
- JsonlStore falls back to built-in messages config when no config file found and tokenizer_path is provided
- DatasetFactory.load forwards tokenizer_path to store for SFT/SEQ+jsonl
- assistant turns train, other roles masked, position_ids doc_reset
- RolloutGenerator._instruction_to_messages builds system/user/assistant list (instruction->system, input->user, output->assistant), replacing single-user-turn concatenation
- Remove _iter_samples helper; _prepare_prompts zips parallel list-of-strings fields directly per the collate_fn contract
- Tests adopt a system-aware chat template and pin the three-field role mapping
- Drop unused imports caught by ruff F401 (torch.Tensor in scheduler.py, iter_raw_records in pipeline.py, Tuple in evaluate_rouge.py)
- train_context.py referenced non-existent cfg.batch_size, replaced with cfg.batch_per_device
- default group_size lowered from 8 to 1: without a group concept (DPO), scheduler batch equals batch_per_device; rollout-based DPO can opt in via extra_kwargs['group_size']>=2
- inline expressions (rollout_batch_size, max_seq_len) extracted for readability
- add tests/trainer/test_online_e2e.py: end-to-end online_dpo via Trainer.train, exercising KV-cache-backed rollout path
- SamplingPipeline.sample gains return_logprobs; both greedy and multinomial paths now share a single log_softmax+gather instead of duplicating the sampling logic
- module-level sample() becomes a thin forwarder instead of re-implementing the three-branch logic
- eliminates ~10 lines of duplicated softmax/gather code; no caller-facing API change
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
- New FSDP2Executor registers as 'fsdp2' in ExecutorFactory, using per-module fully_shard() instead of FSDP1 FlatParameter wrapper
- FSDP2 preserves original Parameter objects as DTensors, eliminating use_orig_params=True hack
- FSDP2Executor implements _no_sync via set_requires_gradient_sync, clip_grad_norm via unshard, unwrap_model via DTensor.full_tensor
- Drop **_extra/**_ddp_only_kwargs fallbacks in BaseExecutor/FSDPExecutor/FSDP2Executor, replaced by parallel_mode-aware executor_kwargs dispatch in train.py (ddp-only kwargs only passed when parallel_mode=ddp)
- Export FSDP2Executor in astrai.parallel.__init__
- BaseExecutor.prepare now takes factories and instantiates model via model_fn(), runs before_wrap hook, wraps DDP/FSDP, then builds optimizer/scheduler on the wrapped model
- optimizer/scheduler creation moved into executor.prepare, eliminating the old 'create-then-wrap' hack reliance on use_orig_params=True
- FSDPExecutor/BaseExecutor accept **_extra kwargs to tolerate DDP-only keys (broadcast_buffers, gradient_as_bucket_view) being forwarded via executor_kwargs
- dataloader builds stay external; executor only handles model/optimizer/scheduler
- train_context.py rewritten to load checkpoint state_dict before prepare via a before_wrap closure
- Store owns window_size/stride and __getitem__/__len__/sample_window
- Dataset classes become thin delegators binding a Store to a train-type key mapping
- Drop BaseDataset.get_index and the RecordDataset中间类 (window死代码)
- DatasetFactory forces window_size=0 for record datasets so record semantics never get window-tainted
- token_count/num_records split the legacy len() semantics (raw stream length vs record count)
- Update tests to the new .store/.token_count API and window/record mode switching
- Drop BaseStrategy.model_fn (stored but never read)
- Drop model_fn= passed to StrategyFactory.create in train_context
- Simplify FSDPExecutor.clip_grad_norm None branch to delegate to super()
- Remove DDPExecutor._gather_state_dict override (identical to base)
- Extract shared core (mask building, primary-id extraction, tensorisation, position-id generation) to astrai/preprocessing/core.py; Pipeline and TokenizeTransform both consume it, eliminating ~60% duplicated logic
- Promote BFD _plan to module-level plan_bfd(lengths, max_len) returning pure index bins; BFDPacking.apply and evaluate_ifd._pack_bins both call it, removing the second BFD implementation
- Split Pipeline._flush (49 lines) into _inject_doc_reset_position_ids + _inject_continuous_position_ids + _to_tensors; split Pipeline.run by delegating record iteration to core.iter_raw_records
- Remove dead no-op pop/塞回 in Pipeline.run (L110-111)
- ChatTemplate: defer Jinja2 compilation to cached_property, exclude compiled template from __getstate__ (its dynamic root function has __module__=None and falls back to __main__, breaking pickle)
- AutoTokenizer: bypass __getattr__ for underscore-prefixed attrs to prevent infinite recursion during unpickle when __dict__ is empty
- TrainConfig.max_grad_norm defaults to None
- executor.clip_grad_norm returns grad norm without clipping when None
- train.py --max_grad_norm defaults to None
- dpo_tokenize now uses tokenizer.apply_chat_template to match SFT format
- Prompt rendered with add_generation_prompt=True
- Chosen/rejected appended as assistant turn
- Remove leftover dead code from _extract_text
- Update tests to mock apply_chat_template
- DPO requires sequence-level sum of token logprobs, not per-token mean
- mean reduction made beta*ratio_diff ~0.03 (near-zero gradient)
- loss stalled at 0.6931 because logsigmoid(0.03) has vanishing grad
- sum gives beta*ratio_diff ~10 with meaningful gradients
- load(store=...) binds directly, skipping format detection/processor
- load_path now optional when store is given
- Remove redundant from_store (merged into load)
- Caller can fully control Store construction + processor setup
- StreamStore/RecordStore → Streamable/Recordable (stateless mixins)
- Store is sole base class, no MRO ambiguity
- H5Store/MmapStore/JsonlStore mix in both traits explicitly
- segments_are_records declared per-subclass (H5/Jsonl=True, bin=False)
- Add tests for dpo_tokenize, lazy jsonl, dual-mode H5, stream-only bin
- Remove unused _to_tensor helper
- StreamStore: fetch(begin, end, key) for stream access (SEQ/SFT)
- RecordStore: mixin with fetch_record(i, key) for record access
- H5Store/MmapStore/JsonlStore now dual-inherit both (C3 MRO)
- JsonlStore supports lazy mode via processor= (no TokenizeTransform)
- RecordDataset base class holds processor, DPO/GRPO simplified
- dpo_tokenize pure function for on-the-fly JSONL tokenisation
- DatasetFactory builds processor for jsonl+record datasets
- train.py passes tokenizer_path=param_path uniformly
- progress: len(dataset) returns sample count (stream=windows, record=records)
- json no longer auto-detected as jsonl format
- Extract tokenization/mask/position logic from JsonlStore into TokenizeTransform
- JsonlStore now pure reader: reads JSON records, delegates to transform
- Store no longer imports tokenizer or preprocessing components
- Replace per_record param with segments_are_records class attribute
- Store subclasses declare segment semantics as format-level property
- Store gains fetch_record/num_records alongside stream fetch/__len__
- save_bin/load_bin support per-record offsets via record_keys param
- H5Store/MmapStore/JsonlStore all support dual stream+record access
- DPODataset/GRPODataset use fetch_record, no cross-record concat
- dpo_collate_fn + collate_fn wired through TrainConfig
- fixes attention context leakage in DPO from windowed concatenation
- stream_chat: default empty system_prompt, single-turn mode
- generate_batch: drop hardcoded system role
- generate.py: preserve original fields in messages branch
and use response_key for the output column name
- Few-shot examples now include subject preamble (consistent format)
- Add --seed flag for option permutation (default 0, -1 to disable)
- Shuffles A/B/C/D positions per-question to neutralise positional bias
- Add FrequencyPenaltyStrategy (logit -= penalty * count)
- Per-task rep_window for penalty history lookup
- Wire through engine, task, executor, API layer
- Add --frequency_penalty and --rep_window to stream_chat.py
- 9 unit tests for frequency penalty strategy
- evaluate_mmlu: fix double few-shot injection (build_prompt no longer
adds few-shot, apply_chat handles it once)
- evaluate_humaneval: fix pass@k k-filtering to be per-problem instead
of using first problem's n globally; reuse ProcessPoolExecutor across
problems; fix closure UnboundLocalError in test_one; handle None in
report when k > n
- evaluate_ifd: remove dead code (score_plain/score_messages); add
multi-file/directory input support with --input_path/--output_dir;
add summary.json aggregation and --max_samples; add --dtype flag
- evaluate_ppl: add --device and --dtype flags (was hardcoded to cuda)
- evaluate_ifeval: fix docstring path (scripts/tools -> scripts/eval)
- analyze_weights: add --output JSON export; fix dead code filter
("_norm" not in r was always True)
- BFDSplitPacking splits over-length sequences into chunks before BFD
- All keys (loss_mask, position_ids, ...) split in lockstep for alignment
- No tokens lost vs bfd which truncates over-length sequences
- Tests: token preservation, chunk alignment, short unchanged, vs bfd
- Add --resume bool flag to train.py CLI
- --param_path always loads weights only by default
- --resume restores epoch, consumed_samples, optimizer & scheduler
- Checkpoint.load() now preserves full meta dict
- Update test_early_stopping to use new param_path/resume API
- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing
- Mask prompt tokens to 0 so their logprobs excluded from ratio/KL
- Switch to token-level ratio + PPO clipping via reduction='none'
- Slice response token logprobs from full sequence output
- Replace k3 KL estimator with non-negative k1 estimator
- Fix epsilon from finifo.eps (~1e-38) to 1e-8
- Remove unused 'reduction' param from GRPOStrategy.__init__
- Clarify offline batch semantics in docstring
- Add 11 unit tests for masking, advantage, KL, sync, clipping
- Sync training.md and architecture.md docs
- Remove MIN_BLOCKS template param and __launch_bounds__ attribute
- Profiling shows smem (not registers) is the occupancy bottleneck for D>=64, making the hint a no-op
- D=64 sees 2-4% speedup, D=128 unchanged (smem-capped at 1 block/SM)
- Update comment blocks in kernel header and both dispatch sites
- Verified correctness via standalone CUDA test (max_err ~1e-4)
- Double-buffered KV (STAGES=2) for D<=128: next tile cp.async overlaps current tile MMA compute, hiding global load latency
- Q loaded directly from global into mma A-operand registers, removing sQ staging and prologue syncwarp
- Predicated cp.async unifies full and partial tile paths, eliminating scalar fallback branch
- STAGES=1 fallback for D=256 (double-buffer would exceed smem budget)
- Applied to both contiguous and paged decode MMA kernels
- ~1.27x average speedup on L20 (sm_89), zero precision loss
- Add attn_paged_decode wrapper in ops.py with gather fallback
- Register kernel in loader.py and export from __init__.py
- Extract test_utils.cuh shared by all attention unit tests
- Rename attn_paged_vs_contiguous.cu to attn_paged_decode_test.cu
- Refactor decode/prefill tests to use common bf16 helpers and cpu ref
- Fix k_cache dim check in attn_paged_decode.cu
- PagedAttentionParams merged into attn_common.h
- Scalar variant: warp-per-query-head split-KV, resolves page table per-position for the K/V shared-memory tile load
- MMA variant (sm_80+): tensor-core head-packing with cp.async, single page-table lookup per tile (BC=32 fits within page_size>=32)
- Standalone test: 14 cases across head_dim 32/64/128/256, GQA, multi-batch, both paths verified against CPU reference
- inline decode_use_mma() into dispatch_decode()
- drop G>1 guard, MMA works correctly for G>=1
- decode is memory-bound, tensor cores + cp.async still win at G=1
- Remove tiles_total/8 min-work cap that limited splits for small workloads
- Simplify decode_num_splits to only use base_blocks and tiles_total
- Short sequences now generate more blocks, improving SM utilization
- Convert AttentionParams to a template struct supporting arbitrary types
- Rename attn_common.cuh -> attn_common.h (no CUDA-specific code remains)
- Include standard headers explicitly in each .cuh instead of via attn_common.cuh
- Allow .h files in csrc/ via .gitignore
- Replace bf16 pre-scale Q loading with direct 32-bit aligned bf16x2 reads
- Apply scale in float32 after Q@K^T, before online softmax
- Reduces causal max error from 2^-6 to 2^-8 with zero perf cost
- Rename all csrc/kernels/gqa_*.cuh/cu to attn_*, with _split_q / _split_kv
strategy suffix and optional _mma compute suffix
- Remove non-split MMA decode kernel, keep only split-KV path
- Convert scalar decode fallback to split-KV (o_part/ml_part + combine)
- Move combine kernel to attn_decode_split_kv.cuh (shared by both paths)
- Rename GQAParams to AttentionParams
- Update all C++ #include, PYBIND11, and Python extension references
- read the 8 Q elements each lane needs straight from global into the mma
A-operand layout, pre-scaled, instead of staging through shared sQ
- removes the sQ smem area (20KB->16KB) and the serialized per-warp prologue
with its WARPS __syncthreads barriers
- result vs torch SDPA: prefill 0.70-0.82x -> 0.85-1.00x (matches torch at
seq=128; 2048 1.251->1.159ms), correctness unchanged across head dims
- double-buffer K/V one tile ahead via cp.async to overlap load with tensor-core math (ncu long_scoreboard 2.12->0.53)
- reorder wait->barrier->prefetch so one __syncthreads/tile covers both cross-warp publish and buffer-reuse (was two)
- add predicated cp_async_16_pred (src-size=0 zero-fills OOB) to unify full/partial tiles, dropping the scalar fallback
- halve BC to 16 to keep 3 blocks/SM despite the doubled smem
- pack adjacent bf16 output into one 32-bit STG, removing the uncoalesced scalar-store penalty (14%->5% sectors)
- result vs torch SDPA: prefill 0.61-0.78x -> 0.70-0.82x, spills eliminated
Decode has only batch*kv_head independent tasks, so the grid was tiny (e.g. 16 blocks) leaving most SMs idle (ncu: 0.04 waves/SM, 11% DRAM).
- Partition KV across gridDim.z blocks emitting unnormalised (O, m, l) partials, reduced by a new combine kernel
- Choose split count to fill the device (~2 blocks/SM), capped by tile count and 32; fall back to single-pass direct-write when batch*kv_head already saturates the SMs
- Refactor decode dispatch into named helpers, de-duplicate scalar fallback
Result: now DRAM-bound at 63% (99->543 GB/s), 2.1-2.5x over torch SDPA in the low-parallelism regime, on par at high parallelism
Moving Qa ldmatrix into the tile loop caused warps 0-2 to read
warp 3's Q data from sQ (only the last warp's data survives the
serialized load loop). Reverted to loading Qa during the init phase
and keeping it resident; __launch_bounds__ still forces 128 regs
(33% occupancy) with spill to local memory.
- Move Qa[KD][4] into tile loop (reload from sQ per tile)
cutting ~32 resident registers for HEAD_DIM=128
- Replace extern __shared__ with static template-sized smem
(no cudaFuncSetAttribute or dynamic allocation needed)
- Add __launch_bounds__ with MIN_BLOCKS param, dispatch by HEAD_DIM
(hd=128→4, hd=64→6, hd=32→6)
- Remove dynamic smem from scalar kernel and C test
- Result: hd=128 168→128 regs, 25%→33% occupancy
- Extract SingleOutputMaskBuilder for SFT and pretrain configs
- Extract MultiOutputMaskBuilder for DPO and GRPO configs
- Keep SectionedMaskBuilder as backward-compatible facade
- Register "single" and "multi" names in MaskBuilderFactory
- Add parity and rejection tests for concrete builders
Decode MMA kernel previously used scalar global→shared loads with
LD=HEAD_DIM+8 padding and per-tile scale multiply. This commit brings it
in line with the prefill MMA kernel (which already had these optimizations):
- cp.async K/V loads (bypasses registers, halves load instructions)
- XOR swizzle: LD=HEAD_DIM instead of HEAD_DIM+8 (zero waste smem)
- Pre-scale Q during load (removes per-tile scale multiply in softmax)
- Clean up prefill MMA kernel comments (no code change)
~2x speedup on decode (0.47ms→0.24ms at seq_len=512)
- add case 32 to decode/prefill dispatch switch
- fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory
- restructure decode dispatch to #ifndef/#else/#endif matching prefill
- split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API
- Add gqa_decode_attn_mma.cuh for tensor-core decode path
- Add dispatch_decode<> selecting MMA vs scalar based on G and mask
- Add TORCH_CHECK for unsupported head_dim instead of silent scalar launch
- centralize CXX_FLAGS/NVCC_FLAGS in csrc/build.py as single source
- add --use_fast_math, --ptxas-options=-O3,-v, --extra-device-vectorization
- add -march=native -funroll-loops host flags
- setup.py reads shared cxx_flags/nvcc_flags from registry
- sync pure-C test build commands with new flags
- Add gqa_decode_attn/gqa_prefill_attn dispatch functions
- Internal _available/__modules with underscore prefix
- CUDA kernel path with F.scaled_dot_product_attention fallback
- GQA head expansion in fallback path
- replace scalar fragment loads with ldmatrix.sync.x4/x2
- add smem row-stride padding (LD = HEAD_DIM + 8) to eliminate 8-way bank conflicts from HEAD_DIM being a 32-bank multiple
- switch build flag from positive to negative: -DASTRAI_NO_MMA for pre-sm_80 only; mma is the default path
- vectorize scalar path smem loads with float4 ld8
- fix pure-C test configs for ld8 alignment
- one query row per group of G=8 lanes, each owning HEAD_DIM/G dims of qreg[]/acc[] in registers
- removes full 32-lane warp_reduce_sum; S dot reduces over only G lanes
- templated on <HEAD_DIM,G,ROWS,P_BC>, block=(G,ROWS)=(8,32)
- per-group shuffle mask so causal loop-bound divergence doesn't deadlock the shuffle
- update pure-C test to the templated launch
ViperEkura