- 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
- Split .cuh/.cu for gqa_decode_attn and gqa_prefill_attn
- gqa_prefill_attn: tiled shared-memory K/V, fused load, compute-opt, mask support
- Add pure C tests under csrc/tests/ for fast nvcc-only iteration
- Update .gitignore for build artifacts
- Replaced old KVCage API with PageCache/ContiguousCache
- Added --cache contiguous|paged switch for decoding comparison
- Added argparse for all params (batch/prompt/gen/device/dtype)
- Fixed PageCache decode crash by extending pages for full sequence
- Sentinel-anchored unconditional pass: both branches now predict the same N response tokens
- Single-token responses (rl=1) fully supported
- ctx_len tracked per sample; skip_reason replaces silent None
- --per_token flag for per-token IFD breakdown
- Add KVCache/CacheView abstract base classes in cache.py
- Add ContiguousCache (contiguous per-slot buffer, default) alongside PageCache (paged, renamed from old KVCache)
- Merge make_table_tensor + bind into bind_tasks on KVCache interface
- Remove task_cached/task_record_hashes from base class (PageCache-only)
- Scheduler: decode all position groups instead of just the largest (eliminates 63% group skip rate)
- Scheduler: accept optional cache param for swapping implementations
- Model layer type hints use CacheView base class
- Batch 1-32: 1-7% speedup from eliminating Storage.gather overhead
- All 183 inference tests pass
- Replace SIGALRM+exec() with subprocess.run(timeout=) for test execution
- Add --test_only flag to skip generation and test existing completions
- Add --generate_only flag for generation-only runs
- Derive pass@k values from num_samples (filter k > n)
- Support loading completions from array JSON (not just JSONL)
- Remove separate ValidationCallback, merge into MetricCallback
- Progress bar now tracks optimizer steps instead of micro-steps
- Remove unused log_interval config field and CLI flag
- Fix validation all_reduce: use SUM(loss, count) instead of AVG
- Simplify metric logging: always log every optimizer step
- Add grad_norm display to progress bar
ViperEkura