- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration
- remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly
- drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections
- add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines
- attention, rotary_emb, and fp8_ops kernels are unchanged
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md
Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard
- drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper
- add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style
- rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs
- Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].
- deepen common-shape BF16 GEMV tuning with warp-row tiling for LLaMA/Qwen2/GPT-NeoX/OPT decode projections
- add fused BF16 up/gate SwiGLU CUDA primitive with ASTRAI_SWIGLU=0/1/auto dispatch
- keep the unfused linear backend as the default path; auto enables no shape until per-architecture checkpoint gates pass
- fall back to the linear/torch chain when kernels are absent, on CPU, in training, or outside supported M/K/dtype shapes
- add gemv/swiglu benchmark scripts, dispatch and parity tests, and kernel documentation
Benchmark: NVIDIA L20 (sm_89), CUDA 12.8, PyTorch 2.11.0+cu128, idle GPU. AstrAI 1B config (24 layers, hidden 1536, vocab 100000), BF16, prompt 128, 32 greedy decode tokens, CUDA graphs enabled, A/B in separate interleaved processes (3 rounds, 8 trials each, medians). Default vs ASTRAI_SWIGLU=1 per generate call: batch 1 134.8->129.1 ms (+4.44%), batch 2 136.2->130.9 ms (+4.06%), batch 4 145.5->140.3 ms (+3.66%). Greedy output identical at batch 1, differs at batch 2/4, so auto stays unfused by default; kernelless fallback verified bit-identical greedy.
- Replace per-element loads with 128-bit uint4 vectorized loads (8 halves per access), improving every measured shape: q/k/v at M=2 from 6.0us to 5.4us, q_proj speedup 2.28-2.45x, mlp_down at M=4 2.76x, lm_head at M=1 +6-8%
- Extend kernel M support from {1,2,4,8} to all M in 1-8 via new BLOCK_M cases 3,5,6,7, since cuBLAS wmma templates pad small M to 8/16 rows and waste compute
- Keep the auto-dispatch allowlist unchanged: a 64-step greedy-walk probe on the real decode path showed mlp_down (K=6912) divergence at step 1 and argmax flips for every candidate odd-M band, the same noise class already present in the merged M=2/4 entries, so no entry has the stability evidence the gate requires
- Rejected alternatives with measurements: split-K accumulation (k/v shapes regress 6.0us to 9.2us, code removed) and MMA tiles (small M is DRAM-bound at ~1 FLOP/byte vs the ~138 needed)
- Update test_gemv M-rejection case to M=9 and test_linear_dispatch multirow fallback to M=9 for the widened range
Benchmark: 8x L20 (sm_89, CUDA 12.8), single-GPU microbench, 200 iters after 20 warmup, weights L2-resident; q(1536x1536) M=3 8.9->5.3us, kv(256x1536) M=3 8.7->3.0us, down(1536x6912) M=3 53.5->10.3us; full gate 691 passed, test_bf16_gemv_uses_current_stream passes in isolation after GPU contention rerun
- add decode-shape benchmark harness
- add bf16 GEMV CUDA primitive with head-dim generic kernel
- dispatch decode-time linear layers to gemv for M=1
- extend gemv coverage to small decode batches
- the kernel's last block folds amax into the history window and publishes the next scale in-kernel (atomicAdd ticket + fences), replacing the host update chain
- quantize bindings split into quantize(transposed) / quantize_dual with fixed arities and a QuantLayout enum; the python adapter becomes a thin attention-style wrapper over pybind (Optional ring_state at the boundary, no torch.library custom_ops)
- tests: in-kernel fold vs host reference (exact), dual/transposed orientation byte-equality
Benchmark: L20 (sm_89), 1.2B model, full train step. Per-linear fixed overhead 28.8us -> 8.8us; fp8 vs bf16: M=512 77.5ms, M=2048 144.5ms (1.15x), M=8192 527.4ms (1.28x); losses bit-identical.
- re-enable kFastLoop on the 128x128 CTA for congruous layouts: the base-pair fragment addressing freed the registers the old offset tables spilled, and the predication-free interior loop now wins across the band (fast body 142 SASS instr with zero predicated fallback vs 719/136 generic; 128 regs, no spill)
- move the big/small CTA dispatch boundary from 3/4 to 5/8 wave: with the fast big-CTA loop the crossover sits between 49 and 63 tiles (63-tile rect +8%, 1024^3 now takes the big CTA)
- fuse the linear bias into the GEMM epilogue: FP8Params.bias_ptr adds in fp32 before the single bf16 rounding, replacing the separate out + bias elementwise pass; guarded loads keep N tails exact and batch broadcast falls out of the row-major layout
- resolve Python None bias in the pybind layer (py::object + cast) so ops/fp8.py and fp8.py pass the argument through untouched; drop the _empty_bias sentinel machinery
- add fused-bias tests covering odd N tails, no-bias parity and batched broadcast
Benchmark: L20 (sm_89), CUDA-graph e2e. Big-CTA fast loop + dispatch: 1024^3 102.6->106.3T, 1152^3 128.5->133.3T, 2048^3 173.8->178.2T, 3072^3 180.2->185.3T, 8192^3 196.2->197.7T. Bias fusion (with-bias GEMM vs unfused out + bias): 1024^3 90.5->106.1T (+17%), 2048^3 162.2->178.3T (+10%), 4096^3 178.2->191.1T (+7%). Fused bias differs from the split path by <=1 bf16 ulp and is closer to the fp64 reference. 596 tests pass.
- mm_fp8 accepts 3D operands through the same signature: grid.z slices by batch strides, size-1 batches broadcast (stride 0), inner .t() views fold into the layout tag at zero copy
- fix _LinearFp8 backward crash on 3D [B,L,d] training inputs (flatten before mm_fp8, reduce grad_b over leading dims)
- expose kRasterGroup/kStreamOut as template knobs; drop the 64x128 mid CTA and staged crosswise-B path from dispatch (direct wins everywhere re-measured, including DRAM-streamed B)
- dispatch thresholds grounded in fresh sweeps: m<=64 -> 64x64 CTA (+27% at 64x8192x2048), small-CTA crossover at SM*14/3 total tiles (+13% at 96 tiles), threshold counts batch x per-matrix tiles (+31% at 64x512^3 bmm, +25% at 8x1024x2048)
- remove scripts/tools/bench_fp8_gemm.py (superseded by csrc/tests/fp8_sweep.cu for kernel-level tuning)
Benchmark: NVIDIA L20, E4M3, NT pre-quantized, median of 100-200 iters
- 64x8192x2048: 29.1 -> 22.8 us (94 TF/s)
- 1024x1536x2048: 67.4 -> 59.6 us (108 TF/s)
- bmm 64x512^3: 139.8 -> 106.7 us; bmm 8x1024x2048: 186 TF/s
- regression-free: 4096^3 192 TF/s, 8192^3 200 TF/s, 512^3 unchanged
- parameterize warp tile (WarpM/WarpN) in Fp8GemmTraits; MMA loops, fragment arrays and epilogue scale with kMt/kNt instead of the fixed 64x32/4x4, enabling cuBLAS-style 64x64 CTAs of 32x32 warps
- dispatch by output tiling (grid-searched via csrc/tests/fp8_sweep.cu): fewer than 48 output tiles take 64x64/32x32 with a lean ring (4 CTAs/SM fill the wave-quantization gap: 512^3 goes 16 -> 64 CTAs); larger shapes keep 128x128 with the kStages+1 ring
- kStages+1 canonic ring rotation drops the post-compute barrier on the congruous path (one __syncthreads per k-tile); LeanRing keeps the kStages ring for the small CTA; direct-crosswise operands always rotate kStages+1 (their prefetch issues right after barrier 1 and would race a lean ring - caught by the pure C layout suite)
- stage the bf16 epilogue through the reclaimed operand smem: swizzled scatter + barrier + coalesced 16B copy-out replaces 8 disjoint 16B per-warp segments (~50% write efficiency before)
- hoist per-lane ldmatrix swizzle offsets out of the mainloop (stage-relative table + ring-base add) so the innermost loop stops recomputing IMAD/LOP3 address chains
- bypass the torch.library dispatch for real CUDA tensors in quantize/mm_fp8 wrappers (~5us/call, ~40% of a 512-wide call's wall time); fake/subclass tensors keep the custom_op route
vs the previous kernel + python path, wall clock on NT squares: 512^3 52 -> 13us (4.0x, 5.2 -> 20.5 TF, now 1.36x cuBLAS _scaled_mm), 1024^3 1.05x, 2048^3 1.02x (46.9 -> 48.2 TF kernel-only); correctness: 4 layouts x 6 shapes pure C suite PASS, 588 pytest PASS
- split quantize into quantize.cuh, templated on input type (bf16/fp16/fp32)
- rename pybind entry quantize_bf16 to quantize; validate the fmt enum
- fix fp8x2 packing: one 32-bit word packs two pairs (halves were dropped)
- drop the dead OutFp8 template param; GEMM output is always bf16
- fp8_state.reset() restores recipe/format defaults too (test state leak)
- rewrite tests for the two-primitive API with fp32-domain amax references
- g/x/w may each be bf16 or pre-quantized fp8 matching fmt; a pre-quantized operand skips its quantize kernel
- snapshot sx/sw/sg before the ring finalize overwrites the aliased scale slot so the gemm dequantizes with the quantize scale
- forward carries its scale to backward so gradients reuse the forward's scale
- grad_input/grad_weight forced bf16; a pre-quantized g dequantizes before the bias-sum
- regression test: two delayed steps with a changing amax must not leak the scale ratio
- Finalize scale rings inside the quantize kernels: a last-block epilogue (threadfence + counter elect) folds amax into hist, reduces the window and publishes the next scale on device, zero extra launches; _ScaleRing packs [hist | scale | counter] into one CUDA buffer.
- Split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments.
- Drop the inference weight-quantization cache; the optimizer bumps the weight version every step, so a cache would miss anyway.
- Zero amax scratch via empty + cudaMemsetAsync instead of torch::zeros, cutting a ~50us fill_ dispatch per quantize.
- Stage crosswise-B operands K-major with cp.async (contract >= 8192) and PRMT-transpose per k_seg region in smem, interleaved with the MMAs; the sync LDG + byte-scatter path it replaces was long-scoreboard bound (ncu 4.6 vs 0.4 stalls/issue).
- Load crosswise-A direct with an in-register PRMT transpose; its operands are typically L2-resident and the staging round trip measured as a net loss.
- Enable grouped rasterization for the congruous NT forward (shared B stripe keeps the weight operand hot in L2) and make the smem budget layout-aware (Fp8GemmSmem) while holding two CTAs per SM.
- Annotate ops/fp8.py return types; drop weight-cache and decorator tests, hoist their imports to module level.
e2e 12L/dim1024/B4xT512 fused AdamW: fp8 137.8ms/step vs bf16 210.3ms, 1.53x. Kernel vs cuBLASLt _scaled_mm: fwd 1.03-1.09x, dX 1.33-1.47x, dW 1.30-1.39x (from 1.10/1.42-1.49/1.52-1.56x), before the pre-transposed copies cuBLASLt needs for dX/dW. fp8 train step vs bf16: 1.34x at 2048 tokens (was 1.25x), 1.08x at 512.
- last-block epilogue (threadfence + counter elect) folds amax into hist[idx], reduces the window and publishes the next scale on device — zero extra launches per linear layer
- _ScaleRing packs [hist | scale | counter] into one CUDA buffer; the eager hist-write / max / scale-copy chain and update() are gone
- split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments
- e2e 12L/dim1024/B4xT512 (fused AdamW): fp8 137.8ms/step vs bf16 210.3ms, 1.53x; fwd 1.82x, bwd 1.50x
- linear_forward_fp8 accepts pre-quantized w8 (matching fmt) and skips the weight quantize; amax_w returns 0 on that path since no bf16 values are seen
- bias is now fused into the GEMM epilogue for both dtypes, replacing the separate torch-level add (one elementwise kernel per linear removed)
- FP8Params.bias becomes void* with a new bias_scale slot: null scale = raw bf16 bias, non-null = fp8 storage dequantized in the epilogue after the operand scaling and before any output quantization
- ops/fp8.py relaxes the w dtype check to bf16-or-fp8 and passes bias_scale through
- regression test covers w8/b8, w8/bf16-bias and the amax_w = 0 contract vs an explicit quantization reference
- loader.py: lazy/cached import; is_available defers the actual load; get_module raises on unavailable
- ops/{attention,rotary,fp8}: use get_module instead of touching private _modules or their own _mod() cache
- package-data: ship astrai.extension.lib *.so in built wheels (non-editable installs previously lost every kernel)
- route dX/dW through the fused 128x64 fast kernel via contiguous transposes
- drop the legacy 64x64 kernel, cutting dX 1.55->0.38 ms and dW 1.28->0.26 ms
- sync all threads after cp.async.wait_group to fix sporadic NaN in large GEMMs
- add fp8_mm_prequant_fp8 custom op for FP8-in/FP8-out GEMM
ViperEkura
0z5a