Conversation
…ache buffer gate Three independent, measured improvements to the gfx950 Gluon sparse-MLA decode, each gated so default behaviour is unchanged unless an env override is set. 1. _decode_num_splits: after the cost-model search, collapse to the fewest splits that preserves both the wave count and the per-split BLOCK_K iteration count. The model treats extra splits as ~free while waves==1, so at C=64/H=16 main=128 extra=8 it picked 3 where 2 does the same 2 iterations with 1/3 fewer CTAs and a 2-way instead of 3-way reduce. Idea from vLLM's _decode_gfx950_num_splits (vllm-project/vllm#52212). Measured -8% at that shape (17.41 -> 15.99 us); split choice unchanged at extra 272/1024/2048, so no regression surface elsewhere. 2. CS0_ALIGN: the row base of every cache gather is block_idx*cs0 + pos*576 with block_idx/pos gathered at runtime, so divisibility analysis assumed 1-byte alignment and a contiguous 512-byte row gather lowered to hundreds of global_load_ubyte. 576 is a literal the compiler already reasons about; cs0 was the one opaque term. The driver computes the largest power of two <=16 dividing both page strides on the host, so the assertion can never be false. Byte loads 416 -> 288, runtime -1.8%. 3. MAIN_USE_BUFFER_LOAD / EXTRA_USE_BUFFER_LOAD replace the single USE_BUFFER_LOAD, which was gated on max() of the two cache spans -- one oversized cache disabled the fast path for both gathers. Each cache is now gated on its own span. Measured on the forced-global path: the 8-token extra gather costs +26% without a buffer descriptor while the 128-token main gather costs +1.3%, because EXTRA_BLOCK_SIZE=2 scatters 8 tokens across ~8 pages while MAIN_BLOCK_SIZE=64 keeps 128 tokens in ~2. Env overrides for experiments, all inert by default: AITER_PA_DECODE_BLOCK_K, AITER_PA_DECODE_MFMA_K, AITER_PA_DECODE_FORCE_GLOBAL_LOAD (1|main|extra). Correctness: op_tests/triton_tests/attention/test_pa_decode_sparse.py 114 passed / 10 skipped, and numerics match vLLM's triton reference to 6.2e-3 max relative error across b=1..128, h=1..128 (unchanged tolerance). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two independent changes to the sparse-MLA decode path. 1. Attention kernel: carry the software-pipelined KV tile across the MFMA in raw fp8 instead of dequantized bf16, and expand it into LDS in NOPE_CHUNK-wide pieces. A [BLOCK_K=64, HEAD_SIZE=512] tile is 32 VGPRs/lane as u8, 64 as bf16 and 128 as the f32 the dequant goes through, and the loop keeps two tiles in flight, so the one-shot dequant was what pinned the kernel at 1 wave/SIMD. Splitting first makes the dependence chain explicit (piece c's converts feed piece c's ds_writes and die), which drops unified VGPRs 306 -> 264; waves_per_eu=2 then closes the last 8 to exactly 256. The NoPE gather's warps move from dim 1 to dim 0 so the column direction is a pure per-lane register repeat and the piece split is a register rename (assert_trivial). It also shrinks the slot vector, which lives in SliceLayout(1, gather_l): 88 -> 66 SGPRs, 63 -> 0 accvgpr copies. 2. Reduce kernel: size its head tile for workgroup count instead of inheriting the attention kernel's BLOCK_M. The combine is pure bandwidth, but BLOCK_M=num_heads gave one workgroup per query -- 64 on a 256-CU part, so 3/4 of the GPU idle and one wave to hide every partial load behind. One head per workgroup gives num_queries*num_heads. Reduce at C128A: 6.5 -> 2.6 us. Kernel-only device time, C=64 H=16 fp8 OCP packed 2-loop, vs the in-tree triton kernel (vLLM PR #52212): C128A 18.1 -> 14.1 us (triton 14.5) C4A 22.5 -> 21.9 us (triton 22.0) op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
vLLM's DSv4 dispatch already probes for this parameter and falls back to out.copy_(result) when it is missing, which costs a full [T, H, D] device copy per call -- 3.4 us at C=64 H=16 D=512, on a ~14 us kernel. Default None keeps the allocate-and-return behaviour. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
waves_per_eu=2 caps the allocator at 256 unified VGPRs. The chunked dequant reaches that with one spilled VGPR when both caches gather through buffer_load, but a cache past buffer_load's 2 GB offset gathers through 64-bit addresses and lands at ~321 -- there the cap buys nothing and costs ~150 scratch stores (C4A forced-global: 26.9 -> 37.7 us). Gate it on the fast path being active. Also force NOPE_CHUNK back to HEAD_SIZE when the gather's warps tile dim 1: the piece split is only a register rename (and only legal) with warps on dim 0. Forced-global C=64, new vs pre-series: C128A 20.5 -> 18.9, C4A 27.4 -> 25.9 us. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…split count Two coupled changes to how the decode launch is split. MAIN_SPLITS (kernel): the SWA and top-k segments have different shapes, so one split count has to be wrong for one of them. Splitting a segment past its BLOCK_K tile count does not divide the work, it multiplies it -- every split then owns a partial tile, which costs a full masked gather for a fraction of the tokens. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole tiles while extra keeps all the programs; split_id >= MAIN_SPLITS gets an empty main range and contributes an extra-only partial, which the reduce already handles. _decode_num_splits_occ (driver): pick the count from occupancy and tile count instead of the old cost model. The kernel is now 256 unified VGPRs / 68,608 B LDS = 2 workgroups per CU, so aim for ~2*CU programs, capped by the larger segment's tile count. Once base_wg alone fills the machine an extra split only buys the second occupancy slot, which is worth a round trip of the [queries, heads, D] f32 partials plus a reduce launch only when each split still owns >=4 tiles. attn (kernel + reduce) us vs the old policy, C=64 H=16 unless noted: extra 272 512 1024 2048 C=128: 272 512 1024 2048 old 21.3 22.0 29.8 45.6 26.1 30.8 47.0 81.2 new 19.6 20.9 26.9 38.7 23.4 26.6 37.4 60.2 C=256 extra=8 keeps NUM_SPLITS=1 (the guard); forcing a split there cost 31%. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…oftmax Gather width (GATHER_TW1): spend all 32 lanes of the warp's dim-1 slice on the head dim so one instruction requests a whole 512 B token row, instead of 8 lanes requesting a 128 B quarter and needing four instructions per row. For a scattered top-k gather that is one memory request per token instead of four. The cost is a longer per-lane slot vector (the row index lives in SliceLayout(1, gather_l), which now has 8 dim-0 slots per lane instead of 2), so it loses ~1% where the segment is tiny and wins where it is not: extra 8 128 272 512 128 B/req 11.50 13.12 15.28 16.27 (main-kernel us, C=64 H=16) 512 B/req 11.63 11.79 14.92 15.13 Because dim 1 is now all threads, the chunked dequant splits rows instead of columns; _split2_dim0 is the dim-0 counterpart of _split2 and CHUNK_AXIS picks whichever axis still carries per-lane register repeats. NaN-propagating max: Triton's default max ignores NaN, which on AMD costs a v_max_f32 x, x, x canonicalize per operand -- 60 of this kernel's 96 v_max were those no-ops. Nothing here produces NaN (masked lanes are -inf, the all-masked row is guarded explicitly), so propagate instead, as the gfx1250 unified attention kernel already does. v_max_f32_e32: 96 -> 0, total instructions 4683 -> 4634. Softmax in the base-2 exponent domain: fold qk_scale*log2(e) into S once out of the MFMA and carry m_i already scaled, so the sink combine and the partial store stop converting. Measured neutral (the compiler was already contracting S*qk_scale - m_new_s into v_fma_f32); kept because m then lives in exactly one space end to end. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… the cap binds Once the ~2*CU program target caps the split count, the only remaining question is whether each split's last tile is partial, and S*ceil(tiles/S) is the number of tiles actually executed. At 18 tiles (the 1-loop seq=1152 case) S=8 executes 24 and S=6 executes 18, which was a 5% regression against the previous cost model. Pick the largest S in range that executes the fewest. 1-loop packed fp8 C=64 H=16 (us), vs c6f0489: seq 136 400 1152 before 13.14 18.38 27.78 after 11.68 16.49 26.2 (was 29.25 before this fix) Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The host sizes the launch from the batch's AVERAGE segment lengths, which is all it can get without a device sync. Recompute the useful count per query in-kernel from that query's own indptr entries and let the surplus programs write a neutral partial (m = -inf, l = 0) and leave. The reduce masks its accumulator load on m > -inf, so a program that bowed out never has to write its [BLOCK_M, HEAD_SIZE] f32 partial at all. Free on a uniform batch (nothing to give back) and a large safety net when the launch over-splits, which is what a mixed-length serving batch causes: per-query top-k length is min(ctx // compress_ratio, 1024) while the SWA segment is always exactly min(ctx, 128), so all the raggedness lands in the segment the split count is sized for. attn us, C=64 H=16, 16x-ragged batch: extra 272 1024 static, 16 splits 31.45 34.26 adaptive 20.45 32.55 Split policy: keep min(2*CU/base_wg, tiles) and do NOT refine it toward counts that divide the tile count evenly. Minimizing executed tiles S*ceil(tiles/S) is worth ~5% on the uniform 1-loop seq=1152 shape, but `tiles` is a batch average and wall clock is set by the longest query -- on a 16x-ragged batch that picks 5 splits where 8 is 22% faster. Also measured and rejected: bounding the reduce's split loops by the same per-query count. It turns them into dynamic loops, and losing the static unroll costs 5-7% on a uniform batch, more than the ~0.9 us it saves when over-split. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…issue order PC sampling (stochastic, gfx950, C=64 extra=272) says the kernel is latency bound, not throughput bound: MFMA 2.3%, VALU 6.9%, LDS 6.4% busy, nothing saturated, waves waiting 46.5% of the time with WAITCNT the top stall at 31.7% and s_waitcnt vmcnt(0) -- a *full* VMEM drain -- alone at 8.7%. Two of those drains were avoidable. Prologue: the indptr loads gate the whole KV chain (indptr -> segment range -> indices gather -> cache addresses -> cache gather, three dependent round trips), but they were issued after Q's load AND after Q's through-LDS layout conversion with its two barriers. Q is independent of all of it. Issuing the indptr loads first not only overlaps them with Q, it lets the uniformity analysis scalarize them: 2 global_load_dwordx2 + 2 s_waitcnt vmcnt(0) + 4 v_readfirstlane become 2 s_load_dwordx2 through the constant cache, waited on lgkmcnt, which does not touch the vector memory pipe at all. Gather order: issue the UE8M0 scales before the bulk fp8. vmcnt is a single in-order FIFO, so a wait can only say "at most N outstanding", never name a specific load -- with the scales issued last, the first dequant piece waited behind every data load too. Issued first, they ride the wait the first piece's own data needs anyway. Main-kernel us, C=64 H=16, median of 3 alternating reps on one GPU: extra 8 272 1024 before 11.94 15.40 22.35 after 11.74 15.04 22.44 Prologue hoist alone was codegen-positive but timing-neutral (~0.5%, under the noise floor); the two are committed together because they were measured together. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…he path A cache past buffer_load's 2 GB offset has to gather through 64-bit addresses, and that made gluon 1.21-1.30x SLOWER than the in-tree triton kernel on a real 2.09 GiB compressed cache, reversing a 1.32x win. Comparing the two kernels' TTGIR shows why, and it is not the data gather: triton scales: tensor<32x2x!tt.ptr<i8>> data: tensor<32x128x...> gluon scales: tensor<64x512x!tt.ptr<i8>> data: tensor<64x512x...> We were building a full-width pointer tensor for a value with only HEAD_SIZE/64 distinct entries per row. With buffer_load the identical 32-bit offsets CSE away and it costs 48 byte loads, so it never showed; through 64-bit addresses they do not, and it costs 288 plus the address arithmetic and exec-mask branching that comes with them. Gather NG = HEAD_SIZE/64 wide instead and broadcast back out in registers. scl_l is the dim-2 slice of a 3-D layout picked so that reshaping [BLOCK_K, NG, 64] to [BLOCK_K, HEAD_SIZE] reproduces gather_l exactly -- column c = g*64 + j maps to thread 4g + j//16 -- so the broadcast is a register rename (assert_trivial). Gated on the 64-bit path being active. With buffer_load the wide form already CSEs to the same 48 loads, so the narrow one only adds a layout convert per tile: +3.5% at extra=1024 ungated, restored to parity when gated. USE_BUFFER_LOAD is already a constexpr, so the gate costs no extra kernel variants. Global-path codegen (C=64, extra=272): instructions 7848 -> 5931, global_load_ubyte 288 -> 48, s_and_saveexec_b64 332 -> 88, VGPRs 321 -> 283. Real 2.09 GiB compressed cache, attn us, gluon vs the in-tree triton kernel: C4A (extra=272) 25.36 / 20.84 (1.215x) -> 21.24 / 20.95 (1.014x) C128A (extra=8) 16.88 / 13.35 (1.264x) -> 13.24 / 13.28 (0.998x) Also restructures _cache_load to take the per-token row offset and the in-row column offset separately rather than a pre-summed [BLOCK_K, W] tensor. That is the shape the triton kernel uses (token_data_ptr[:, None] + offsets[None, :]); on its own it is codegen-neutral here because MLIR folds the addptr chain back together, but it is what makes the narrow scale gather natural to express. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nged) kv_smem's row padding sets which LDS banks a transposed K read lands on: bank = (row * pitch_dwords) mod 32, so at PAD=8 the 32 lanes of a column walk share 8 banks and the measured bank-conflict share is 46%. Exposing it as a knob turns that into a testable claim -- and the test says it is NOT causal: pad 4 8 16 32 64 (predicted conflict degree conflict 2-way 4-way 8-way 16-way 32-way rises left to right) extra=272 19.63 14.85 14.87 15.34 15.99 main-kernel us extra=1024 29.49 23.34 23.29 23.37 24.36 Time is not monotone in conflict degree -- pad=4 has the fewest conflicts and is the slowest by 32% -- and pad 8/16/32 are identical across a 4x conflict range. So the LDS cost this kernel pays is latency and synchronization (17.1% of stalls on lgkmcnt, 5.9% on s_barrier, 7 barriers per tile), not bank conflicts. Default stays 8; the knob is kept so the ablation is a 30-second re-run. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nged) Completes the launch-config knob set so the whole BLOCK_M / BLOCK_K / MFMA_K / waves_per_eu space is sweepable without editing the driver. Swept at C=64 H=16 against the current 16/64/16/2; every alternative is worse, so the defaults stand (attn us, extra=8 / 272 / 1024): BLOCK_K=64 MFMA_K=16 (current) 14.54 18.35 26.85 BLOCK_K=32 MFMA_K=16 16.31 25.73 35.67 BLOCK_K=16 MFMA_K=16 (nw=1) 23.03 40.35 79.22 BLOCK_K=64 MFMA_K=32 18.12 22.77 34.37 BLOCK_K=32 MFMA_K=32 23.00 42.55 70.43 BLOCK_K=16 is the interesting one: num_warps = BLOCK_K/16 = 1 removes every cross-warp reduction and barrier, and it is 1.6-3.0x SLOWER. So the 4-of-7 per-tile barriers that come from the softmax reductions are not what dominates, despite being 35% of ATT stall cycles. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ATT put 8.5% of all stall cycles on just 68 buffer_store_dwordx4 instructions (~154 stall cycles each) -- the split-K accumulator write-out -- and counters put the partials at ~31% of the kernel's HBM traffic (10.5 MB of 33.4 at 5 splits). Storing them in bf16 halves both the bytes and the store instruction count. The dtype is taken from part_acc itself (acc.to(part_acc_ptr.dtype.element_ty) on the way out, .to(float32) on the way back in), so the kernel needs no flag. The skip_reduce path keeps f32 because it hands part_acc to the caller. Interleaved A/B, main + reduce us, C=64 H=16: extra 8 272 512 1024 2048 f32 13.71 16.44 19.36 26.38 37.46 bf16 13.55 15.39 14.93 23.90 34.69 ratio 0.989 0.936 0.771 0.906 0.926 Numerics: the mantissa bits this costs are far inside the error the fp8 KV format already carries. Against an fp64 reference built from the pre-quantization KV, max|delta|/max|ref| (test_part_bf16_acc.py): geom f32 partials bf16 partials bf16 vs f32 C128A extra=8 3.517e-2 3.567e-2 5.9e-3 C4A extra=272 2.343e-2 2.343e-2 4.1e-3 extra=1024 2.859e-2 2.885e-2 4.1e-3 extra=2048 2.807e-2 2.811e-2 5.4e-3 i.e. total error moves 3.517e-2 -> 3.567e-2 at worst, and the bf16-vs-f32 delta is ~6x smaller than the error already present. AITER_PA_DECODE_PART_BF16=0 restores f32 partials. Also adds AITER_PA_DECODE_PART_ST / _PART_LD cache-modifier knobs, measured and left at the default: no modifier is within noise of .cg/.wb, and .cs is 4% worse because it evicts partials the reduce is about to read -- i.e. the default policy already keeps them resident, so cache policy was not the lever. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
9459334 gated waves_per_eu=2 on buffer_load because the >2 GB path needed ~321 unified VGPRs, where the 256 cap bought nothing and cost ~150 scratch stores (+40% at C4A). Narrowing the scale gather (e9cb7b3) moved that path to 283, and at 283 the cap lands cleanly: 256 VGPRs with 17 spill slots instead of 148, so 2 waves/SIMD instead of 1. Forced-global, attn us, C=64: extra 272 1024 waves_per_eu 0 20.98 29.11 waves_per_eu 2 15.63 24.25 (0.745x, 0.833x) The buffer path is unaffected (it already asked for 2): 1.001x / 1.005x / 0.999x / 1.003x at extra 8 / 272 / 1024 / 2048. Net effect on a real 2.09 GiB compressed cache, vs the in-tree triton kernel: extra 8 272 1024 gluon 13.04 15.15 24.87 triton 13.34 20.87 31.95 ratio 0.976x 0.725x 0.778x which is now within noise of the same shapes on a small cache (0.995x / 0.709x / 0.773x), i.e. crossing buffer_load's 2 GB limit costs ~0-6% instead of +27-36%. The lesson is about the gate, not the number: it was measured correctly and became stale the moment the register picture moved. The comment now says to re-measure it whenever that happens. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The index list is one int32 per gathered token -- orders of magnitude under buffer_load's 2 GB offset limit even for a full batch -- so it can keep the fast path however big the KV caches are. It was also the largest single source of exec-mask branching in the kernel: 20 of the 55 s_and_saveexec_b64 on the buffer path, all from the masked tail, because a masked gl.load predicates while a masked buffer_load folds the mask into the offset. Codegen: s_and_saveexec_b64 55 -> 36, global_load_dword 60 -> 0 replaced by 60 buffer_load_dword, at the cost of 15 more spill slots (14 -> 29). Interleaved A/B over 9 rounds, main + reduce us, C=64 H=16: extra 8 272 512 1024 2048 gl.load 14.68 16.07 15.99 24.30 34.51 buffer 14.61 15.63 16.02 24.50 34.60 ratio 0.995 0.972 1.002 1.008 1.003 This is the third time re-measuring after the register picture moved changed a verdict. The same change was tried and REVERTED earlier in this series (see RESULTS.md 7.4), when it measured 0.967x at extra=272 but 1.037x / 1.047x at extra=1024 / 2048 -- a wash. bf16 partials and the ungated waves_per_eu freed enough pressure that the long-extra regression is now within noise (1.008x / 1.003x) while the extra=272 win holds, so it is a net positive. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…pilogues An f32 divide lowers to a ~10-instruction sequence (v_div_scale / v_rcp / v_div_fmas / v_div_fixup), so dividing the whole [BLOCK_M, HEAD_SIZE] accumulator by l pays that per element. Take one reciprocal per row and broadcast-multiply. Applied to both epilogues. The NUM_SPLITS == 1 path in the attention kernel is constexpr-eliminated whenever the launch splits, which is both production configs (C128A S=2, C4A S=5), so the win is entirely in the reduce kernel: kernel instrs divide sequence (rcp/scale/fmas/fixup) _pa_decode_sparse 5604 -> 5604 2 -> 2 (path eliminated) _pa_decode_sparse_reduce 336 -> 264 40 -> 5 (8/16/8/8 -> 1/2/1/1) Reduce-kernel time, 5 alternating reps on one GPU: extra 8 272 1024 2048 divide 2.82 2.42 3.56 3.67 rcp*mul 2.70 2.19 3.57 3.65 ratio 0.957 0.905 1.003 0.995 4-10% where the split count is low; neutral at 8 splits, where the reduce is bandwidth-bound over more partials and the ALU saving hides. Whole-attn effect 0.990-0.996x. Numerics: one extra rounding in f32, far below the bf16 output's resolution. Reported by mehmet.kaymak@amd.com, who spotted it in the NUM_SPLITS == 1 epilogue; extended here to the reduce kernel where the divides actually execute. op_tests/triton_tests/attention/test_pa_decode_sparse.py: 114 passed, 10 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… unchanged) The cs0 alignment hint is now codegen-neutral. Re-measured on the current kernel it is worth 0.989-1.005x on both the buffer_load and the 64-bit gather path, and the emitted load mix is byte-for-byte the same footprint with it on and off across six builds per path. The reason is the row/column gather split: the in-row column offsets are compile-time arange constants now, so load vectorization no longer depends on proving cs0's divisibility and gl.multiple_of(cs0, 16) tells the compiler nothing it does not already use. The hint and its host-side stride scan could be deleted; this only adds the override so the check is one command. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ers it needs
Short top-k was the one shape losing to the in-tree triton kernel (1.07-1.11x in
the ISL-8192 traces, all C128A layers). Three changes, each gated on the launch
shape, take it to parity or better without touching long top-k at all.
1. Fused dequant via inline asm. `v_cvt_scalef32_pk_bf16_fp8` does fp8 x scale ->
bf16 in one instruction; the backend emits three (convert with the scale
hardwired to 1.0, then v_pk_mul_f32, then v_cvt_pk_bf16_f32) and neither fold
is implemented -- see /home/mekaymak/amd-fp8-dequant-lowering. Reaching it from
Gluon needs the fp8 passed as int16 (the AMDGPU backend cannot allocate a
sub-32-bit inline-asm *input*, but <2 x i16> is one VGPR out of one
buffer_load_dword) and the E8M0 byte shifted into the operand's exponent field
rather than exp2'd. Bit-identical: the scale is a power of two, and fp8 -> bf16
is exact. 5605 -> 4625 instructions.
Two details are load-bearing for register pressure, which is what decides
whether any of this pays:
- The scale operand is int16, not f32. At pack=2 Triton gives an f32 input
TWO registers and the asm only ever names the first. The hardware reads
bits [30:23], i.e. bits [14:7] of the high half, so `e << 7` in an i16
lands the exponent in the same place for one register.
- Only $0 needs early-clobber. The second convert re-reads $2/$3 after $0 is
written, so $0 must not alias an input; $1 is written last and may.
Together: 5 VGPRs per invocation -> 3, and spills 33 st/50 ld -> 21/28, below
the f32 chain's 29/29. Skipping either costs more than the fold gains (at 5
VGPRs/invocation the fused form is 1.20-1.31x SLOWER at long top-k).
The int16 row offset is `nope_row >> 1`, not a second `block_idx*(cs0//2)`
product -- both address the same byte, and sharing `block_idx*cs0` with the
scale gather drops 96 v_mad_u64_u32 and 48 v_add_lshl_u32.
2. waves_per_eu follows the launch. num_warps is 4, one wave per SIMD per
workgroup, so the second occupancy slot exists only if a CU gets a second
workgroup. At <= CU workgroups it cannot, and the 2-waves/EU cap is then pure
loss -- it spills ~29 dwords across the per-tile softmax row-max reduction and
buys no overlap. Note the headroom alone is worth exactly nothing (1.000x);
it is what the fused dequant does with it that pays.
3. BLOCK_M 16 -> 8 when the launch is small AND top-k is short (<= 2 KV tiles).
Doubling heads_blocks is real parallelism, unlike a finer split-K, which hands
every split a partial tile -- forcing splits to reach 2 wg/CU measured
1.14-1.36x worse. Refused when it would cost a split (at C=128 the policy
answers the doubled heads_blocks with 1 split instead of 2, which is 1.11x).
Main kernel, C=4/16/64/128 at top-k 64, all arms one process one GPU, median of
13 rounds, duplicated control arm within 0.3%:
C 4 16 64 128
+occ1 1.000 1.000 0.995 1.001 <- headroom alone: nothing
+asm 0.961 0.958 0.960 0.955
+bm8 0.931 0.933 0.955 0.954
vs in-tree triton (attn = kernel + reduce), baseline -> new:
C=4 0.830 -> 0.773 C=16 1.077 -> 1.005 (the losing cell, now parity)
C=64 0.983 -> 0.939 C=128 0.756 -> 0.721
Long top-k (272 and 1024, C=64 and C=128) is unchanged: 0.997-1.003x against the
old default, inside the duplicated control's own 0.993-1.000x.
BLOCK_K=32 was tried for the same register relief and is a clean loss everywhere
(1.06x at top-k 64, 1.31-1.51x beyond) -- the halved tile doubles the per-tile
index/scale/barrier cost and buys nothing.
AITER_PA_DECODE_ASM_DEQ / _WPEU / _BLOCK_M still force any of the three either
way; the asm path is gfx950 + packed OCP fp8 only.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ath; add GRID_ORDER Follow-up to 19f5977, which measured only the buffer_load path. **The >2GB cache path is gather-bound, not instruction-bound.** The fused convert works there -- it still removes 959 instructions and, at 1 wave/EU, all 25 spills (vgpr 286 -> 289, scratch 25/25 -> 0/0) -- but buys ~0, and asking for 1 wave/EU is a reproducible LOSS: 1.016x at C=16 top-k 1024 (spread 0.5-0.9%, control 0.999x). 256 workgroups round-robined over 8 XCDs is not one per CU in practice, and that path needs the second occupancy slot more than it needs the registers. So one_wg_per_cu now requires use_buffer_load, which also switches the ASM_DEQ default off there. Verified back to 1.000x on that cell, buffer path unaffected. BLOCK_M=8 stays on both paths (neutral to 1.1% better on the 64-bit one). **AITER_PA_DECODE_GRID_ORDER** permutes the launch grid; default "qsh" is the existing (queries, splits, heads). Grid dim 0 varies fastest and XCD assignment is round-robin over the linear workgroup id, so which axis sits in dim 0 decides what shares an XCD's L2. All six permutations are correct (57 passed each); the default is already the best of them, and getting it wrong is expensive: order C=64 tk64 C=64 tk1024 C=128 tk64 C=128 tk1024 qsh (dflt) 1.000 1.000 1.000 1.000 qhs 0.998 0.999 1.003 0.999 sqh 1.031 1.109 1.044 1.151 shq 1.044 1.114 1.047 1.156 hqs 1.050 1.003 1.001 1.002 hsq 1.046 1.116 1.048 1.154 Any order with splits in dim 0 is 3-16% worse: consecutive workgroups are then the splits of one query, which round-robin across 8 XCDs, so one query's Q vector and KV rows get pulled into 8 separate L2s. With queries in dim 0 and num_queries a multiple of 8, query q always lands on XCD q%8 and all its splits and heads share one L2 -- and the split-K partials stay in the L2 that the reduce (also queries in dim 0) will read them from, which is why the reduce moves too (2.20 -> 2.55 us). hqs is the exception that confirms it: equal to the default except at C=64 top-k 64, the one cell where BLOCK_M=8 makes heads_blocks=2, so heads in dim 0 spreads each query over 2 XCDs and strides queries to only 4 of the 8. This corrects an earlier conclusion in RESULTS.md that XCD locality was worth nothing here. It is worth up to 1.156x -- the earlier experiment perturbed a swizzle that happened to preserve queries-in-dim-0, so it measured the wrong thing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… calls
From ASM_DEQ_FINDINGS.md §4.1, and it is the largest single win of this series.
A single-instruction inline asm reads all its sources before writing dst, so no
output can clobber an input and neither needs early clobber. The two-output blob
form does need it -- without `=&v` on the first output the allocator may hand it the
scale's register and the second convert then reads a clobbered scale, wrong in
exactly the elements it produces -- and that forced simultaneous liveness is what
was spilling: both bf16 outputs are live before any input dies, per 4-fp8 group,
across the whole unrolled dequant, where the f32 chain's intermediates die pairwise.
Peak, not average, is what the allocator spills against, and the 256-VGPR cap has
zero headroom.
Splitting it costs +184 instructions (each call re-reads its sources) and buys back
most of the spills:
form instructions scratch st/ld
f32 chain 5604 29 / 29
fused, two-output blob 4624 21 / 28
fused, two single-output calls 4808 11 / 12
Strictly better than the blob at every buffer-path shape measured (3-arm A/B, one
GPU, control arm within 0.4%):
C=16 top-k 8 0.985 -> 0.973 C=64 top-k 64 0.983 -> 0.970
C=16 top-k 64 0.986 -> 0.973 C=64 top-k 512 1.321 -> 1.081
C=128 top-k 8 1.017 -> 0.977 C=128 top-k 64 1.011 -> 0.977
C=128 short top-k -- the production C128A shape the occupancy gate enables -- goes
from a 1.01-1.02x LOSS to a 2.3% win, so the blob was costing us there. The reduce
knock-on also nearly disappears (3.43 -> 2.34 us at top-k 512), confirming
FINDINGS §3: the reduce was never slow, the blob's scratch traffic was evicting the
split-K partials from L2 before the reduce read them.
Default vs the pre-session default, main kernel, same protocol (control 0.986-1.007x):
C top-k 8 top-k 64 top-k 272 top-k 1024
4 0.918 0.918 0.947 0.932
16 0.892 0.906 0.944 0.934
64 0.924 0.921 0.994 1.001
128 0.939 0.933 0.998 1.000
vs in-tree triton (attn = kernel + reduce): C=16 top-k 8 goes 1.114x -> 1.014x and
top-k 64 1.128x -> 1.038x; C=64 1.007x -> 0.945x; C=128 0.753x -> 0.716x. Long
top-k at C=64/128 is untouched (0.994-1.001x), which is the gate working.
The blob form is removed rather than kept behind a knob -- it is dominated
everywhere and a second inline-asm spelling in a hot file is a liability.
Also tried and rejected: NOPE_CHUNK=64 (FINDINGS §6). Un-chunking is sound in
principle -- chunking exists to bound the f32 intermediate's live range and the
fused convert deletes that intermediate, so it drops spills to 16/17 and it is
catastrophic WITHOUT the fold (2.0-2.7x) -- but in an arm-count-matched A/B it
helps only in the window where the gate already disables the fold (top-k 512:
1.292 -> 1.163, still a loss) and HURTS where the fold is enabled (top-k 64:
0.980 -> 1.019). FINDINGS §6's favourable reading came from a run whose extra arms
inflated it; see RESULTS.md §16.1 on arm-count sensitivity.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The fp8 loop peeled the partial last tile into a separate masked instantiation.
Gluon inlines, so that second call site was a second full copy of the
gather+dequant+MFMA body -- ~1000 instructions and, on the 64-bit path, 24
s_and_saveexec + 19 s_cbranch of predication. Merely *emitting* it cost 0.900-0.907x
(buffer) and 0.943-0.972x (64-bit) at >=2 tiles/split, because the registers it
demanded spilled the tile loop.
The merge works because the full-tile path already handled invalid slots correctly:
it clamps the row in-bounds and lets the score mask do the rest. Masking only ever
existed to bound the index-list read. So instead of masking that load, clamp k_pos
into the segment -- an out-of-range lane then gathers a REAL row whose score is -inf,
exactly how -1 sentinels are already handled -- and fold `k_pos < hi` into the mask.
The partial tile becomes the last iteration of the loop; no peeled copy.
Masking is on the MFMA's OUTPUT, not the gather: S is [BLOCK_M=16, BLOCK_K=64] f32,
16 values per lane, against the [64, 512] tile the old path masked. And the mask is
built directly in the MFMA accumulator layout --
col_mask = (k_start + gl.arange(0, BLOCK_K, layout=gl.SliceLayout(0, qk_layout))
< seg_hi)[None, :]
S = gl.where(col_mask, S, -inf)
-- so applying it is a per-lane v_cndmask with no data movement. Deriving it from the
slot-layout `valid` vector via convert_layout instead is a cross-lane change per tile
and at 16 tiles ate the entire saving: that version measured 1.072x at C=32 x1024
where the native mask measures 1.016x, and every cell improved when it was replaced.
The PV MFMA needs no mask at all, since P = exp2(S - m) is already exactly 0 there.
Codegen (C=64 x272, f32-chain build, buffer path):
UNI_TILE=0 5605 instructions 29/29 scratch 36 saveexec 112 ds_write_b128
UNI_TILE=1 3931 instructions 16/16 scratch 23 saveexec 76 ds_write_b128
Main kernel, ratio vs UNI_TILE=0, control arm within 0.2%:
x64 x272 x1024
buf C=32 0.960 0.941 1.016
buf C=64 0.962 0.987 0.973
buf C=128 0.957 1.028 0.984
glb C=64 0.922 0.814 0.991
10 of 12 cells faster, two mild losses (1.6% and 2.8%) with no predictor clean enough
to gate on without overfitting. Default ON; AITER_PA_DECODE_UNI_TILE=0 restores the
peeled tail. The bf16 (non-fp8) path keeps the old structure.
Correctness: 60/60 with UNI_TILE on and off, on the 64-bit path, and with the fused
dequant disabled. Plus a 12-case differential against UNI_TILE=0 covering lengths that
are not multiples of BLOCK_K and 3x/4x/8x/16x-ragged batches -- bit-exact on both
address paths. Ragged is the case that matters and it caught a real bug: the
loop-internal _gd_fp8 call sits one indent level deeper than the peeled one, so a
first patch missed it and left the prefetched partial tile with no range test. The
signature was that queries with one tile were exact and queries with a full tile
followed by a partial one were wrong -- and only against triton, not against
UNI_TILE=0, did it show as 0.85 max relative error rather than plausible noise.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…teps were dead Once the score mask is built natively in the MFMA layout (b585afc), the UNI_TILE branch of _slots was doing three things it did not need to: 1. `gl.maximum(hi - 1, 0)` -- _gd_fp8 only runs under `if n_full > 0`, i.e. hi > lo >= 0, so hi >= 1 and the clamp-to-zero can never fire. 2. `slot < num_rows` -- that check belonged to the MASKED path, where `other=-1` injected a sentinel and a predicated-off lane's slot was undefined. A clamped read returns a genuine in-segment index, exactly as trustworthy as the full-tile path's, which never checked it. This also removes the only use of `num_rows` outside MASKED, so _gd_fp8 stops threading it. 3. `gl.where(valid, slot, 0)` when not HAS_INVALID -- the clamped read is already in bounds, so the select is pure overhead. The full-tile path already gated it on HAS_INVALID; UNI_TILE now shares that gate. With those gone the UNI_TILE and full-tile cases differ only in the load offset and one extra term in `valid`, so they fold into one body and the buffer/global load dance is written twice instead of three times. Not just tidier -- it is faster, because all three were per-tile work: instructions 3931 -> 3848 scratch 16/16 -> 14/15 (UNI_TILE=0: 5605, 29/29) Main kernel vs UNI_TILE=0, control within 0.2%, before -> after this commit: x64 x272 x1024 buf C=32 0.960 -> 0.950 0.941 -> 0.931 1.016 -> 1.028 buf C=64 0.962 -> 0.946 0.987 -> 0.990 0.973 -> 0.938 buf C=128 0.957 -> 0.950 1.028 -> 1.003 0.984 -> 0.951 glb C=64 0.922 -> 0.918 0.814 -> 0.835 0.991 -> 0.977 11 of 12 cells now faster than the peeled tail; the one loss is C=32 top-k 1024 at 1.028x. Still bit-exact against UNI_TILE=0 across the 12-case differential (non-multiple -of-BLOCK_K lengths, 3x/4x/8x/16x-ragged batches) on both address paths, and 60/60 with UNI_TILE on, off, and on the 64-bit path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ment is one FMA
max commutes with a positive scale (qk_scale = scale * log2(e), and any real attention
scale is > 0), so scale the row max -- one value per lane -- instead of every element
of S. That leaves `S * qk_scale - m_new` as the only use of the scaled S, and it lowers
to v_fma_f32 where the split form needed a separate multiply across the tile:
v_fma_f32 0 -> 16 total instructions 3848 -> 3842
**Performance-neutral**, and recorded as such so it is not re-litigated: 0.989-1.006x
across C=32/64/128 x top-k 64/272/1024 plus the 64-bit path, entirely inside the control
arm's own 0.996-1.003x. The reason is size -- S is [BLOCK_M=16, BLOCK_K=64], four f32
per lane per tile, which cannot move a 10 us kernel however it is spelled. Sizing the
tensor first would have predicted this. The multiply that *would* matter is
`acc * alpha_pv` at 128 f32 per lane, and MFMA has no accumulator-scale operand, so it
cannot be fused at all.
Kept rather than reverted because it is strictly fewer instructions and one rounding
instead of two; applied unconditionally rather than behind a knob so it costs no extra
constexpr. Numerics: the row max is the same product bit for bit, only the fused
subtract differs -- 3 of 4 differential shapes are bit-identical and the fourth differs
by 0.23 of a bf16 output ulp, in the more accurate direction. -inf * positive = -inf,
so masked columns stay masked. 60/60.
Note for the next person: an earlier revision of this patch put the exp2 edit in
_decode_tile and the max edit in _qkpv_fp8, leaving one path double-scaled and the other
unscaled. It showed up as 27-44% relative error AND as "12 multiplies removed, no FMA
added" -- the codegen counters diagnosed it before the numbers did. Two copies of the
same softmax exist (_decode_tile for the bf16/legacy-tail path, _qkpv_fp8 for the
pipelined fp8 path); a text patch must name which.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…gregates Replace the ~40-parameter helper signatures with three aggregates (the unified_attention_2d.py pattern) and name-derive the packed-row offset constants. DSv4/uniform configurations compile to byte-identical .amdgcn against the pre-restructure kernel (11/11 configs incl. both address paths, splits=1, UNI_TILE=0, BLOCK_M=8, fnuz); the DSv4 wrapper only switches to the new MAIN_FMT/EXTRA_FMT constexpr names and passes the two new (None-elided) scale-pointer arguments, keeping the kernarg layout unchanged. Adds the separated-rope (ROPE_SEPARATE) MLA geometry under the same kernel: second pow-2 LDS plane for the K-only rope, chained QK MFMA over kv_lora_rank + rope, V = plane 0; new cache formats "tensor" (flat per-tensor fp8, scale folded into qk_scale / p — no dequant work in the tile loop) and "dsmla" (vLLM fp8_ds_mla 656 B rows). Also fixes the bf16 _decode_tile call arity (2 stray args since the UNI_TILE commit made every bf16 config fail to trace). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… MLA decode
New API for MLA-lineage models (GLM-5 style: kv_lora_rank latent + appended
decoupled rope, token-granular top-k gather, no sink) over the shared gfx950
gluon kernel. Follows the aiter.mla.mla_decode_fwd calling convention as vLLM's
ROCMAiterMLASparseImpl uses it, minus the head padding (H=8 runs natively at
BLOCK_M=8) and the q quant (q stays bf16). Cache format inferred: bf16, flat
per-tensor fp8 (+ scalar kv_scale), or vLLM fp8_ds_mla.
Measured on MI355X vs the asm path at true GLM shapes, C=1..256 x topk {512,
2048} x H {8, 16}, both address paths: this kernel wins short top-k -- 0.75-0.99x
of mla_decode_fwd at topk=512 for C<=128 (fp8) and 0.57-0.90x to C=64 (bf16) --
and runs at parity to 1.36x slower at topk=2048, with the crossover around
C=64-128 where the asm kernel's persistent-CTA scheduling starts to scale
better. Accuracy is equal or better throughout: at fp8 it is more accurate than
the asm kernel at the asm kernel's own format, 3.0e-2 vs 4.2e-2, because asm
runs fp8 end-to-end and re-quantizes P for the PV dot.
(The perf figures in this message's original form -- "0.38-0.55x" -- were wrong.
bench_head2head.py summed self_device_time over every profiler row, and on ROCm
aiter's C++ ops carry attributed device time alongside their kernels, so the asm
arm was counted twice while the Triton arms were counted once. Corrected and
re-swept 2026-08-20; the bench now filters DeviceType.CUDA. The gluon-vs-gluon
and gluon-vs-triton ratios elsewhere on this branch are unaffected -- every arm
there is a Triton kernel with no op wrapper.)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… tests Kernel: comments trimmed to the constraints the code cannot show (comment-only; the DSv4/uniform configurations still compile byte-identical, 11/11). sparse_mla_decode_fwd: drop all AITER_PA_DECODE_* env overrides in favor of fixed tuned constants and computed policy; drop the short-top-k BLOCK_M gate (never reachable at MLA top-k sizes). Tests slimmed to the used paths: bf16 + per-tensor fp8 matrix, one fp8_ds_mla smoke, split-K, bitwise scale-folding, cache shapes, and the 64-bit path via a MAX_BYTES monkeypatch (no env var). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_pa_decode_sparse_gfx950_gluon: all 19 AITER_PA_DECODE_* overrides replaced by the tuned constants and computed policy they already defaulted to, dead locals removed, and the comments cut back to the constraints the code cannot show. Behaviour is pinned by capturing every launch parameter across 76 driver configurations / 122 kernel launches (2D pool and 3D packed caches, one and two segments, split-K, skip_reduce, sinks, head alignment, fnuz, sentinels): grid, positional arguments and constexprs are identical to before, PART_LOAD_CACHE aside. _pa_decode_sparse_reduce: load the attn sink once in pass 1 and reuse its base-2 form for the l_final fold, take .cg on the partial loads directly, and drop the PART_LOAD_CACHE constexpr -- both launchers stop passing it. Kernel-side comments are otherwise comment-only: the DSv4/uniform configurations still compile byte-identical to the frozen reference, 11/11. op_tests test_pa_decode_sparse 114 passed / 10 skipped, test_sparse_mla_decode 14 passed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Optional path for the flat per-tensor OCP fp8 cache (sparse_mla_decode_fwd's
fp8_mfma=, default off). The gather's own code points go straight into LDS with
no dequant, both dots run on the fp8 matrix core
(v_mfma_f32_16x16x32_fp8_fp8 -- K=32 is the only shape the backend offers for
plain fp8, and it divides kv_lora_rank, the rope width and BLOCK_K alike), and
the per-tensor scale still folds outside the tile loop: K-side into qk_scale,
V-side onto the accumulator once in the epilogue. q keeps arriving bf16 and is
quantized to e4m3 inside the kernel from a per-program amax, so its scale folds
into that program's qk_scale and m/l stay in the true-score domain -- split-K
programs remain comparable in the reduce.
Codegen at the GLM shape: 256 -> 192 VGPRs, 5 -> 0 spills, 24 -> 0 B scratch,
136 -> 68 MFMA, 144 -> 0 dequant converts, LDS halved, tile loop 504 -> 337
instructions.
LDS_PAD has to follow the element size. The pad shifts which banks a transposed
K read lands on, which is a byte property, so the bf16-tuned 8 elements is only
8 bytes on an fp8 plane and lands on the wrong stride: carrying it over costs
1.083x, while 16 / 32 / 64 all land at 0.86-0.89x. _lds_pad_for() doubles it for
this path. Counters found it -- nothing is saturated (MFMA 7%, VALU 19%, HBM
traffic unchanged) but LDS busy was 0.564, the highest of any unit, with the
wait fraction up 0.365 -> 0.470.
Measured vs aiter's asm mla_decode_fwd, C=64 topk=2048 (the GLM shape), against
the bf16-staging path in parentheses:
H=16 asm 26.43 bf16 28.89 (1.093x) fp8 25.59 (0.968x)
H=8 asm 26.68 bf16 27.26 (1.022x) fp8 23.86 (0.894x)
so this recovers the long-top-k regime the bf16 path was losing. Across
C=16..256 the fp8 path is 0.75-0.96x of the bf16 path. The cost is accuracy: it
lands at the asm kernel's level (4.1e-2 vs asm 3.9e-2) where the bf16 path is
2.5e-2, because q and p both go through e4m3 -- a deliberate trade, hence a flag.
Off by default the path is inert: the DSv4/uniform configurations still compile
byte-identical to the frozen reference, 11/11. op_tests 128 passed / 10 skipped.
Tried and rejected: BLOCK_K 128/256 (1.3-1.5x worse; 64 stays optimal),
kv_splits 16/32, waves_per_eu 3 and 4 (both spill -- 192 VGPRs does not reach
the <=170 needed for a third wave, so occupancy stays pinned at 2/SIMD),
GATHER_TW1 8/16/32 (within 1%), and a power-of-two pre-scale on p before the
e4m3 cast (1 / 16 / 256 agree to three digits; 4096 saturates).
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…onvention
sparse_mla_decode_fwd now takes q as float8_e4m3fn plus the scalar f32 q_scale
it was quantized with, which is how aiter's asm mla_decode_fwd is called and how
vLLM's sparse path already prepares q: one scaled_fp8_quant over [C, H*d_qk]
with the layer's calibrated _q_scale. The scale folds into qk_scale exactly the
way the cache's k_scale does, so nothing per-tile changes. bf16 q keeps working
and is unchanged.
The scale is per-tensor, not per-token, to match what the asm kernel takes.
This composes with either dot path, and both combinations are useful:
fp8 q + bf16 dots accepts production's already-quantized q without giving up
the accurate dots. fp8 -> bf16 is exact (3 mantissa bits
into 8), so this costs only q's own quantization:
5.4e-3-1.2e-2 rel-err vs 3.3-4.6e-3, and it is free in
time (29.01 vs 29.27 us at C=64 H=16 topk=2048).
fp8 q + fp8 dots the full production convention, and now the fastest arm.
Dropping the in-kernel amax that bf16 q needed saves
1.3-1.8 us, and accuracy is unchanged from deriving the
scale in-kernel (2.69e-2 vs 2.81e-2) -- a per-program amax
buys nothing over one calibrated per-tensor scale.
vs aiter's asm mla_decode_fwd at GLM shapes, topk=2048, arms interleaved,
control 0.96-1.03x (tensor = bf16 q + bf16 dots, tensor8 = fp8 q + fp8 dots):
C=16 H=16 asm 18.53 tensor 18.29 (0.987x) tensor8 16.94 (0.914x)
C=16 H=8 asm 18.43 tensor 17.35 (0.941x) tensor8 15.76 (0.855x)
C=64 H=16 asm 27.44 tensor 29.46 (1.074x) tensor8 24.36 (0.888x)
C=64 H=8 asm 27.07 tensor 27.66 (1.022x) tensor8 22.87 (0.845x)
C=128 H=16 asm 37.55 tensor 44.79 (1.193x) tensor8 38.16 (1.016x)
C=128 H=8 asm 37.29 tensor 42.38 (1.137x) tensor8 36.50 (0.979x)
C=256 H=16 asm 64.08 tensor 83.48 (1.303x) tensor8 66.62 (1.040x)
so the long-top-k regime the bf16 path lost from C=64 up is now a win to C=128,
and C=256 is within 4% instead of 30%.
Both new arguments default off and are elided when unset, so the DSv4 and
uniform-pool launchers are untouched: identity gate still 11/11 byte-identical,
op_tests 128 passed / 10 skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@stefanskiasan I added support for glm 5.3 flash, you should be able to set qk_rope_head_dim to 0 to deactivate rope part. |
|
Verified your GLM-5.3 fix on 8× MI350X (gfx950) — it works, and it matches what I had patched locally. Pulled
One data point that may interest you, from the vLLM side: the merged ROCm implementation for GLM-5.3 (vllm-project/vllm#53943 → ZJY0516/vllm#5) does not route to this kernel — it still reaches Thanks for the quick turnaround — six hours from report to fix. |
This was referenced Aug 27, 2026
Open
This was referenced Aug 29, 2026
5 tasks
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Critical and moderate correctness, cache-format, stride, and specialization issues remain unresolved.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds gfx950 Gluon sparse MLA support and extends sparse paged attention for FP8/DSV32 cache formats, with tests and benchmarks.
Changes:
- Adds
sparse_mla_fwdwrapper and cache-format handling. - Extends kernels for split-K, LSE, reduction, and output handling.
- Adds correctness tests and performance benchmarks.
File summaries
| File | Findings |
|---|---|
op_tests/triton_tests/attention/test_sparse_mla.py |
Nit, 1 vote: Add correctness coverage for pre-quantized FP8 Q with q_scale. |
op_tests/op_benchmarks/triton/bench_sparse_mla.py |
Moderate, 1 vote: Account for Q rereads across active splits. Moderate, 2 votes: Compute heads_blocks using the wrapper’s BLOCK_M rule. |
aiter/ops/triton/attention/sparse_mla.py |
Moderate, 3 votes: Validate contiguous Q last-dimension layout. Moderate, 2 votes: Require row-contiguous flat caches. Moderate, 1 vote: Handle packed fnuz encoding explicitly or require uint8. Nit, 3 votes: Move tuning values to JSON configuration. Nit, 1 vote: Cover pre-quantized Q and express both return tuple shapes. |
aiter/ops/triton/attention/pa_decode_sparse.py |
Moderate, 1 vote: Require contiguous provided output tensors. Nit, 2 votes: Move hard-coded launch and tuning values to JSON configuration. |
aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py |
Critical, 3 votes: Include all specialization and tuning keys in the kernel representation. Critical, 1 vote: Guard or reject the DSV32 rope-store aliasing case. Moderate, 3 votes: Guard zero l_final reciprocals. Moderate, 2 votes: Mask invalid split-K maxima before exponentiation. Nit, 1 vote: Include reduction branch keys in the reduction kernel representation. |
Review details
Suppressed comments (6)
aiter/ops/triton/attention/pa_decode_sparse.py:55
- The Gluon output stores also use
offs_dwithout an output last-dimension stride, but this newout=path only validates shape, dtype, and device. A non-contiguous caller-provided output will therefore be written to incorrect addresses; requireout.stride(-1) == 1(as the MLA wrapper already does).
assert out.shape == q.shape, f"out shape {tuple(out.shape)} != q {tuple(q.shape)}"
assert out.dtype == dtype, f"out dtype {out.dtype} != {dtype}"
assert out.device == q.device
return out
aiter/ops/triton/attention/sparse_mla.py:260
- [verified] The new tests never call
sparse_mla_fwdwith an already-quantized FP8 Q andq_scale; alldots="fp8"cases pass BF16 Q and exercise only internal quantization. The separateQ_FP8path and scale folding can fail without detection. Add a correctness case using a pre-quantized Q and its scale.
if q_is_fp8:
# Caller-quantized q, the asm calling convention: one scaled_fp8_quant
# over [C, H*d_qk] with the layer's scale. The kernel folds the scale
# into qk_scale, so nothing per-tile changes.
if q_scale is None:
aiter/ops/triton/attention/sparse_mla.py:87
- [verified] This packed-cache branch accepts any 1-byte dtype, including
torch.float8_e4m3fnuz, but then forcesfp8_fnuz=Falseand interprets the mixed record as OCP e4m3. Passing a packed cache with fnuz storage would therefore silently dequantize the payload incorrectly. Requiretorch.uint8for this mixed-format record or carry and handle the encoding explicitly.
kv.ndim == 3
and kv.element_size() == 1
and kv.shape[2]
== kv_lora_rank + 4 * (kv_lora_rank // 128) + 2 * qk_rope_head_dim
aiter/ops/triton/attention/sparse_mla.py:196
- The declared return type only describes
(out, lse), butskip_reduce=Truewith more than one split returns(part_acc, part_m, part_l)(three tensors). This makes the public API misleading to type checkers and callers; express the two possible tuple shapes (or use overloads keyed onskip_reduce).
) -> tuple[torch.Tensor, torch.Tensor | None]:
op_tests/op_benchmarks/triton/bench_sparse_mla.py:76
- [verified]
sparse_mla_fwdrereads Q for every split, but this Q term is counted once regardless ofnum_splits. For bandwidth runs with split-K active,bytes_movedunderreports traffic and overstates GB/s. Multiply the Q term bynum_splits(or the actual active split count).
q = num_tokens * num_heads * D_QK * kv_elem_bytes
op_tests/triton_tests/attention/test_sparse_mla.py:124
- All correctness cases construct Q in BF16, so the production pre-quantized-Q branch (
q.dtype == float8_e4m3fn, including its requiredq_scale) is never validated; the benchmark only times it without an oracle. Add a correctness case that passes FP8 Q plus its scale and compares against the independent reference.
out, _ = sparse_mla_fwd(
q, cache, ptr, idx, sm, kv_scale=ks, qk_rope_head_dim=rope, **kwargs
)
- Files reviewed: 4/5 changed files
- Comments generated: 9
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This kernel extends the existing sparse paged attention kernel (initially designed for dsv4 decode) to cover GLM5 like sparse decode/prefill as well. Also, optimizes it further for dsv4 as well.
The target is vllm, hence it consumes the default kv cache format for those models (
fp8 with tensor scalefor glm5, andfp8 ds mlaformat for DSV4.TODO:
Technical Details
Performance tuning for dsv4 + and sparse mla support for GLM like LLMs.
Test Plan
New tests for the new path (sparse mla).
Test Result
Performance evaluation
Execution times are in us and it includes the reduce time if it exists.
GLM 5.2
fp8 cache with a per-tensor scale,
topk = 2048, q pre-quantized.shippedis thedefault configuration the launcher picks; no flags.
ASM is the existing aiter/ASM path from VLLM.
Warm: No L2 cache flushing, cold is with cache flushing to remove cache effect between runs.
H=8 is TP8, H=16 is TP4. Numbers are from MI355 + Triton 3.7.1
DSV4
VLLM currently uses 2 loops, 1 for SWA and 1 for compressed KV part. Hence, 2-loop is the actual vllm path.
Triton is the gfx950 specific triton implementation from vllm.
The two DSv4 layer types differ enough to report separately:
H=16, TP8 setting.
index_topkcap)