AstrAI

Commit Graph

Author	SHA1	Message	Date
ViperEkura	e9b03f4fca	perf: apply cp.async, XOR swizzle, pre-scaled Q to decode MMA kernel 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)	2026-07-08 16:15:14 +08:00
ViperEkura	6adc221c10	refactor: extract shared MMA utils into gqa_mma_utils.cuh - Move mma16816, ld2, pk2, pkb, ldmatrix_x4/x2/x2_trans to shared header - gqa_prefill_attn_mma.cuh and gqa_decode_attn_mma.cuh both include it	2026-07-07 23:01:15 +08:00
ViperEkura	9e63cb9ed0	feat: MMA head-packing decode kernel with scalar fallback dispatch - 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	2026-07-07 22:56:02 +08:00

Author

SHA1

Message

Date

ViperEkura

e9b03f4fca

perf: apply cp.async, XOR swizzle, pre-scaled Q to decode MMA kernel

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)

2026-07-08 16:15:14 +08:00

ViperEkura

6adc221c10

refactor: extract shared MMA utils into gqa_mma_utils.cuh

- Move mma16816, ld2, pk2, pkb, ldmatrix_x4/x2/x2_trans to shared header
- gqa_prefill_attn_mma.cuh and gqa_decode_attn_mma.cuh both include it

2026-07-07 23:01:15 +08:00

ViperEkura

9e63cb9ed0

feat: MMA head-packing decode kernel with scalar fallback dispatch

- 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

2026-07-07 22:56:02 +08:00

3 Commits