- 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
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
- 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
ViperEkura