Repository navigation
[Triton/Gluon] GFX950 sparse mla optimizations - #6159
Merged
Merged
Conversation
Dispatch deals workgroups to the 8 XCDs round-robin by linear id, so grid-order ids put a request's rows, a prefill row's neighbours and a row's head blocks on different L2s. With XCD_REMAP the linear program id goes through remap_xcd (_xcd_work), so each XCD gets one contiguous run of work items, query-major with the head block fastest. The programs that gather the same KV rows then run on one L2 at the same time. The driver remaps every launch except SWA-only and 16-head prefill, and launches padded past their split count (the TP=4 draft: 3 splits launched as 4), where the remap would spread the immediately-exiting pad programs into the first round. On the DSv4.1 AgentX shapes (MI355X, same-session pairs against the previous commit, sparse MLA per step): * TP=2: decode -4.6 %, tail -5.7 %, chunk -2.1 %; * TP=4: decode -7.4 %, tail -2.9 %, chunk +0.7 % (within noise). Without speculative decode (one row per request), decode of 1-32 requests stays within 1 %, and the 31- and 32-request launches at TP=2 run 5.5-6.8 % faster. The one exception is a TP=4 64-row launch built from the 32 requests taken twice (+15 %): grid order kept a request's two copies on one XCD, and distinct requests share no KV. The output is bit-identical.
Decode and tail rows re-read each other's KV rows too: a request's verify and draft rows, and a row's head blocks. With the XCD-aware order they also run on one L2. Packed fp8 launches below prefill size therefore cache their gathers instead of loading them .cg, as prefill already did. On the DSv4.1 AgentX shapes (MI355X, same-session pairs against the previous commit, sparse MLA per step): * TP=2: decode -1.9 %, tail -5.2 %, chunk +0.2 %; * TP=4: decode -2.0 %, tail -6.0 %, chunk -0.4 %. The output is bit-identical.
Every fp8_dsv4_mla segment now takes _pipe_segment. The kernel picks the walk from the segment's format, and the other formats keep _process_segment. Against _process_segment's fp8_dsv4_mla path the walk: * reads a tile's slot ids in the row, rope and score-column layouts in one round trip, a tile ahead, where the old path waits for two dependent round trips per tile; the sentinel mask then needs no cross-warp exchange; * gathers a row with 16 lanes x 32 B instead of 32 x 16 B, so each lane resolves half as many 64-bit row addresses and slot ids; * carries the exponent bytes as the dwords that hold them (loop-carried bytes were repacked with the payload at every back edge); * takes the sentinel target from the first tile's slot ids, and searches only when that tile is all sentinels; * exchanges the softmax row max and row sum across warps through LDS buffers of its own: one barrier each instead of two through the shared reduction scratch. The partial sums add in gl.reduce's order. The dequant stays automatic. It is the scaled upcast where Triton has cdna4.scaled_upcast; otherwise it is the same instruction as inline asm (_deq_asm, unchanged) on the int16 view of the tile, on every launch. Main used the generic exp2 dequant for launches of more than one workgroup per CU, which includes every launch over 64-bit gathers. What this leaves unused goes: the fp8_dsv4_mla branches of _gather_full and _decode_tile, and the slot-id prefetch (IDX_PREFETCH, _next_slots). test_pa_decode_sparse_two_loop_vllm_pool covers vLLM's pool layout: * page-pitched packed caches with 32- and 128-token blocks; * ragged lengths on every tile edge; * sentinels in every other entry or in whole leading tiles; * split-K off and on; * 8 to 64 heads. Its asm variant hides the upcast from the driver and checks the asm dequant bit-identical to it. On the DSv4.1 AgentX shapes (MI355X, same-session pairs against the previous commit, sparse MLA per step): * TP=2: decode -6.3 %, tail -6.6 %, chunk -0.1 %; * TP=4: decode -13.9 %, tail -7.4 %, chunk -1.6 %; * without speculative decode, decode of 1-16 requests is 2.5-4.1 % faster at both TPs, and 31-32 requests are within 1.2 % at TP=2. Without the upcast the steps are 10-30 % faster than main's own fallback: * TP=2: decode -20.6 %, tail -23.1 %, chunk -10.5 %; * TP=4: decode -30.4 %, tail -23.0 %, chunk -10.2 %. The output is bit-identical with either dequant.
…anks At 32 heads (DSv4 at TP=2) a row ran as two 16-head programs, and each gathered, dequantized and staged the same key rows. The 32-head program runs BLOCK_M 32 on 8 warps: the MFMA layouts put 2 warps on M (16 heads each) and 4 on N. Every wave keeps one 16-head Q and accumulator slice, and the program gathers each key row once for both head blocks. One such program fills a CU, so the split count now counts the programs a CU holds at once (wg_per_cu). The driver runs it at prefill, and below prefill size where the 16-head grid would overfill the CUs with at least two tiles per program. A few single-row requests split over many programs (non-speculative decode) are faster as twice the 16-head programs; with the 32-head program they ran 12-13 % slower. On the DSv4.1 AgentX shapes at TP=2 (MI355X, same-session pairs against the previous commit, sparse MLA per step): decode -9.9 %, tail -1.5 %, chunk -4.0 %. Non-speculative decode of 1-32 requests is unchanged (within 0.3 %), and of 64 requests 3.5 % faster. TP=4 is unchanged. The output is bit-identical.
…grams With PREFETCH, _pipe_segment stages tile t before it issues tile t+1's gathers and tile t+2's slot ids. The gathers then overlap the previous tile's dots, and the loop-carried tile needs no copy at the back edge (a copy is a vmcnt(0)). That pays where one program fills a CU: the 8-warp 32-head program, below prefill size. With two 4-warp programs per CU, or at prefill, the other program hides the gather latency, and the loop-carried tile costs more than the prefetch saves (TP=4 decode 7 % slower with it). The driver turns it on for 8-warp launches below prefill size. On the DSv4.1 AgentX shapes at TP=2 (MI355X, same-session pairs against the previous commit, sparse MLA per step): decode -4.2 %, tail -1.6 %, chunk -0.2 %. TP=4 is unchanged. Against main, with the whole series: * TP=2: decode -24.4 %, tail -19.1 %, chunk -6.1 %; * TP=4: decode -21.8 %, tail -15.6 %, chunk -1.4 %. The output is bit-identical.
…k launches At 64 heads (DSv4 at TP=1, or a DP-attention rank) a row ran as four 16-head programs, two per CU. The 32-head program (BLOCK_M 32 on 8 warps, one per CU) gathers each key row once for 32 heads. On launches with a top-k segment it is faster whenever its last round of programs is more than half full. When that round is at most half full, the 16-head grid runs it one program per CU, which is faster than two, and the 16-head grid wins until four other rounds carry the 32-head program's gain. The 192-row spec decode (1.5 rounds) is 5 % slower with the 32-head program. SWA-only launches are 5-16 % slower with it at the AgentX shapes. The driver now runs the 32-head program at 64 heads on launches with a top-k segment when its last round is full or more than half full, or after four full rounds. Both programs were timed on spec decodes of 8-32 requests, non-spec decodes of 10-56 rows and tail prefixes of 260-1,344 rows. On every top-k launch where they differ by more than 1 %, the rule picks the faster one. Both grids get the same split count there, so the output is bit-identical. The tile-ahead gather and the XCD remap follow the existing rules. On the DSv4.1 AgentX shapes at TP=1 (MI355X, same-session pairs against the previous commit, sparse MLA per step): * spec: decode unchanged, tail -4.7 %, chunk -4.4 %; * non-spec: decode of 1-16 requests within 0.3 %, of 31 / 32 / 64 requests -5.5 / -5.3 / -7.0 %; tail -5.6 %, chunk -4.4 %. Against main with the whole series at TP=1: spec decode -15.0 %, tail -21.7 %, chunk -6.7 %. TP=2 and TP=4 launch the same programs as before.
…logue Before wo_a, vLLM's DeepSeek-V4.1 path rotates the trailing 64 lanes of each attention output row back (inverse RoPE) and, when wo_a reads MXFP8, quantizes the rows to e4m3 with one E8M0 scale per 32 lanes. With the aiter kernel both run as a separate pass over the bf16 rows (rocm_inverse_rope_mxfp8_rows, or rocm_inverse_rope_rows_ for a bf16 output); vLLM's other decode backend does them in its reduce epilogue. sparse_mla_fwd and pa_decode_sparse now take inv_rope_positions with inv_rope_cos_sin_cache, and out_mxfp8 = (data, scale) in place of out. The store that writes the output applies them: the attention kernel's without split-K, the reduce's with it. The rows are rounded to bf16 first and the arithmetic is that of vLLM's pass, each product of the rotation rounded on its own, so the output is bit-identical to running the pass after the kernel. * Without split-K the tile moves once through LDS into a layout with 32 consecutive lanes per thread, so the rotation pairs and the scale groups stay in registers. The e4m3 store writes half the bytes of the bf16 one. * The reduce works in its split-sum tile, where every split lane holds a copy of the row. The copies can differ by a rounding, and like the plain store the epilogue's stores keep split 0's. Launches without the epilogue compile to the same code as before. test_pa_decode_sparse_inv_rope_mxfp8_epilogue checks both outputs bit for bit against vLLM's arithmetic on the kernel's own bf16 rows: 8 to 64 heads, 1 to 2,100 rows, no split, 3 and 8 splits. On the DSv4.1 AgentX shapes (MI355X, the fused kernel against the kernel plus vLLM's pass, same session and GPU), sparse MLA with its epilogue per step: * TP=2: decode -12.2 %, tail -8.7 %, chunk -14.9 %; * TP=4: decode -11.9 %, tail -9.6 %, chunk -12.2 %; * TP=1: decode -10.6 %, tail -9.6 %, chunk -15.5 %; * non-speculative decode of 1-64 requests: -9.1 to -21.5 % at every TP. The output is bit-identical to the unfused path on every call type.
…ilogue The epilogue rounded each product of the inverse RoPE on its own (inline asm v_mul_f32) and kept the reduce's work in its split-sum tile, both only to stay bit-identical to vLLM's pass. Drop both: * The rotation is plain Gluon arithmetic, which the compiler contracts into FMAs. * The cos/sin come in with masked buffer loads, and the position is read before the move into the epilogue layout, so their latency overlaps the move. Masked gl.loads branched on exec and, with the FMA code, waited inside the branches. * The reduce spreads the summed row over all its lanes first, which is faster at 8 splits. test_pa_decode_sparse_inv_rope_mxfp8_epilogue now checks against the f32 reference rotated back, the MXFP8 output dequantized. On the AgentX calls the outputs differ from vLLM's pass only by rounding: at most one step, in about 1 in 10^7 e4m3 values and 1 in 4*10^5 bf16 values. Per step on the DSv4.1 AgentX shapes against the previous version: within 1.2 % on spec decode, tail and chunk, and up to 3.3 % faster on non-spec decode, at TP=1, 2 and 4. Launches without the epilogue compile to the same code.
The launcher had one rule per head count for when an 8-warp program covers 32 heads. At 32 heads: at prefill, or when the 16-head grid overfills the CUs and every program gets two tiles. At 64 heads, top-k launches only: when the 32-head grid's last round is more than half full, or past four rounds. Each sized a 16-head split count to decide, and the split count was then sized again for the 32-head grid. Size the split count once, on the 16-head grid: a 32-head grid gets the same count, with half the programs at one per CU instead of two. Then one rule for both head counts: more than one tile per program and, for grids of up to four rounds, a last round more than half full. Bigger grids only need programs of four tiles or more. On the DSv4.1 AgentX shapes this moves three launches: the TP=2 tail and chunk SWA-only calls go to 16 heads (-6.1 % and -4.5 %), and so do the TP=1 non-spec 16-request top-k calls (+0.6-1.1 %). Per step, every suite at TP=1 and TP=2 stays within 0.7 %. It also fixes shapes the old rules were slow on: TP=2 decode of 43-64 requests is 4.6-9 % faster, of 36-40 rows 6 %. TP=4 has 16 heads and is unchanged.
The output epilogue moved the attention kernel's and the reduce's plain stores into the epilogue helper. That changed the code around them enough for LLVM to reschedule the tile loop of the separated-rope formats: GLM-5.2's rope slot ids moved behind the row gathers, and the wait for them became a vmcnt(0) on the gathers, 5-8 % per prefill step. The plain stores are back as they were, and the helper only runs with the epilogue. Launches without it compile to the same code as before the epilogue again (checked on GLM-5.2 and GLM-5.3 fp8 KV, and DSv4.1 at 16 and 32 heads, with and without split-K). Launches with it are unchanged.
…l head blocks The XCD remap gives each XCD a contiguous range of work ids: a row's head blocks and splits, then the rows next to it. A row's head blocks read the same KV rows, so they gain from sharing an L2. With one head block per row (up to 16 heads, or 32 heads on a 32-head program) the remap only groups a row's splits and its neighbours. That made decode of 64 and 128 rows 8 % slower at 16 heads, 0.90x against main at TP=4 and 64 requests. Remap only rows of two or more head blocks. On the DSv4.1 AgentX shapes (MI355X, same-GPU pairs against the previous commit, sparse MLA per step): * TP=4 non-spec: decode of 64 requests -11.6 %, of 1-32 requests -0.9 to -1.5 %, mixed decode unchanged; * TP=4 spec 5+1: decode +1.1 %, tail -0.4 %, chunk -0.6 %; * TP=2 non-spec: decode of 64 requests -1.4 %, the rest within 0.5 %; * TP=2 spec 5+1: decode +0.4 %, tail and chunk within 0.2 %. On synthetic DSv4.1 decodes at 16 heads: 64 and 128 rows -8.4 %, 5+1 decode of 8 requests -7.3 %, of 32-256 requests +0.8 to +3.4 %. At 32 heads every size stays within 1 % except 5+1 decode of 1,024 requests (+1.9 %).
The XCD remap decoded work ids row-major, so each XCD ran one row's splits back to back, then the next row's. On full grids that made the attention kernel 10-14 % slower than the same rows in split-major order: 64 and 128 decode rows on 8 and 4 splits, and 8 spec requests of 6 rows on 8 splits. Rotating the split index by the row, which keeps the heavier first splits apart, did not help. With split-K, each XCD still gets a block of ceil(rows / 8) rows, now run split-major with the head block fastest, so a row's head blocks still run together. Without split-K the two orders are the same, and the code stays as it was. Which launches remap does not change, and nor does the output. On the DSv4.1 AgentX shapes (MI355X, same-GPU pairs against the previous commit, sparse MLA per step): * TP=2 non-spec decode of 32 requests -8.5 %, mixed decode -8.3 %; both run 8 splits on two 16-head blocks; * TP=1 non-spec decode of 16 requests -9.6 % (8 splits on four 16-head blocks); * every other suite within the noise (0.6 %). TP=4 runs the same code: with one head block per row it does not remap. On synthetic DSv4.1 decodes at 32 heads: 32 requests -8.0 %, the other sizes within 0.3 %.
…t lengths _vllm_pool_case now takes each row's SWA and top-k lengths, and sizes its pools to fit the longest row. The tile-edge pattern the tests used moves to _edge_lens(T), so their inputs are unchanged. The asm-dequant variant of the vLLM-pool test goes: it forced the inline asm fallback by hiding cdna4.scaled_upcast and compared it bit for bit.
…nk conflicts The fp8_dsv4_mla walk ran its bf16 dots on v_mfma_f32_16x16x16_bf16, half of gfx950's bf16 rate: on the TP=1 chunk call the MFMA units were busy 22 % of the time for 11 % of the peak. On 16x16x32 the MFMA time halves, but those operands read the KV tile from LDS with more bank conflicts, and with the 8-element row pad the tail call got slower. The tile's store was the larger problem: each lane writes a 64 B run of its row, so lanes 4 apart hit the same banks, 60 % of the tile's conflict cycles. Padding 8 elements (16 B, so 128-bit accesses stay aligned) after every 128 spreads them. With the 16x16x32 operands, the tile's conflict cycles per wave drop from 3,584 (pad 8 per row) and 2,560 (pad 16, prefill) to 1,536, measured on the kernel's own accesses under rocprofv3. A 4-element pad removed more conflicts, but broke the 16 B alignment, and the unaligned stalls made the kernel 45 % slower. The fp8_dsv4_mla walk now runs its bf16 dots on 16x16x32 and pads its KV tile by 8 after every 128 elements, so the DSv4 launcher no longer passes KV_LDS_PAD. Other cache formats compile to the same code (checked on GLM-5.2 and GLM-5.3 fp8 KV, 11 builds each). On the DSv4.1 AgentX shapes (MI355X, same-GPU pairs against the previous commit, sparse MLA per step): * spec 5+1: TP=4 decode -3.2 %, tail -14.3 %, chunk -12.3 %; TP=2 -10.7 / -12.3 / -12.7 %; TP=1 -12.0 / -12.8 / -12.3 %; * non-spec decode of 1-64 requests: TP=4 -3.3 to -4.8 %, TP=2 -3.2 to -5.9 %, TP=1 -3.9 to -7.3 %. The padding alone gives about half of this. The output differs from the previous commit's by accumulation order; its error against the f32 reference is unchanged on every recorded call type checked.
… Q in LDS At 64 heads a row ran as two 32-head programs, each gathering, dequantizing and staging the same key tiles: on the TP=1 chunk call that staging is about a third of each tile's instructions, and the 64 heads gave no reuse (TP=1 took 1.96x TP=2's time for twice the heads). A 64-head program does not fit with Q in registers: 256 VGPRs plus spills, 2.15x slower. For 64-head programs Q now goes to LDS, once per program, and the QK reads it 128 dims at a time with the matching K slice. The 8 warps split 64 heads x 64 keys as 4 x 2, so each key warp owns two 16-key MFMA columns, which the softmax row exchange now handles. Q and the KV tile take 68 KB of LDS each, one program per CU. It pays only with enough rows. At 64 heads a ratio-1 top-k call loses under one round of programs (64 / 128 rows +32 / +46 %) and at 1.5 rounds (384 rows +5 %), and wins from one full round on (256 rows -15 %, 512-1,024 -17 to -19 %, 1,536-12,288 -20 to -27 %); SWA-only calls (two tiles) lose throughout. So the launcher picks it at 64 heads for top-k launches without split-K that have at least one full round of programs and, up to four rounds, a last round more than half full. These programs skip the tile-ahead gather, with which they spill. On the DSv4.1 AgentX shapes at TP=1 (MI355X, same-GPU pairs against the previous commit, sparse MLA per step): spec 5+1 tail -16.1 %, chunk -23.6 %; non-spec tail -15.6 %, chunk -23.5 %. Decode keeps its programs. TP=2, TP=4 and the other cache formats compile to the same code.
…r fp8 caches sparse_mla_fwd sent every two-loop call to pa_decode_sparse, which only reads fp8_dsv4_mla or bf16 block caches, so a DeepSeek-V4 model could not keep its SWA and compressed caches as plain per-tensor fp8 rows the way GLM-5 does. The gluon kernel already decodes fp8_scalar per segment, each with its own scale; route the fp8_scalar two-loop to the MLA launcher and give it the extra segment. New extra_kv_scale (defaults to kv_scale). The launcher also keeps a [nb, block, R] pool paged when its page pitch is padded: reshape(-1, R) would have copied the whole cache. Dots stay bf16 with an extra segment (FP8_MFMA folds the V scale after the loop, which needs one segment); an fp8 q is widened in the kernel.
…o-loop FP8_MFMA accumulates P*V on raw fp8 code points and takes the V scale off in the epilogue, so it was limited to one segment. Rebase the main segment's sum onto the extra segment's scale before the extra loop (acc *= s_main / s_extra) and apply s_extra at the end; the online-softmax rescales in between are linear, so this is exact. Two-loop against the fp32 reference, random data: relative L2 0.0201 with equal segment scales and 0.0202 with 0.75 / 1.25 (fit slope 0.9997); the one-segment fp8 path is at 0.0167 and bf16 dots at 0.0021. The gap is P rounded to e4m3.
…dots for bf16 and per-tensor fp8 caches bf16 and per-tensor fp8 caches with the rope inside the row (DeepSeek-V4, GLM-5.3) still ran as 16-head, 4-warp programs on 16x16x16 bf16 dots. Their walk gathers each key tile into registers, dequantizes it (fp8) and stores it to LDS, about half of the tile loop, so at 32 or 64 heads each tile was staged two or four times per row. * Heads per program (_staged_head_block, both launchers): 64 heads (8 warps, Q in LDS, which every walk now allows) without split-K once the rows fill half the CUs; else 32 heads when each program has two or more tiles and the grid's last round is more than half full (past four rounds, four or more tiles); else 16. The split count stays the 16-head grid's. Forced on every call, bigger programs lost on small grids (non-spec decode of 1-16 requests: +8 to +55 %); the rule keeps those at 16 heads. * A bf16 row's one 32- or 64-head program is XCD-remapped too when there is no split-K: a request's draft rows and a prefill's rows read the same KV rows (TP=2 spec decode +3.6 % without it, -5.1 % with it). Per-tensor fp8 tiles are half the bytes and gained nothing from it, so the MLA launcher does not remap. * bf16 dots on 16x16x32, with an LDS pad fitted to these walks' store (32 B per lane, 32 lanes per row): 8 elements after every 256 and 64 after every 8,192. The fp8_dsv4_mla walk's pad does not carry over. Conflict cycles per wave over 8 tiles: 2,048 before, 3,072 for 16x16x32 on the old pad (per-tensor fp8 decode and tail 2-5 % slower), 1,536 now. On the DSv4.1 AgentX calls rebuilt on per-tensor fp8 / bf16 caches (MI355X, same-GPU pairs against the previous commit, sparse MLA per step; spec 5+1 decode / tail / chunk): * per-tensor fp8, bf16 dots: TP=4 -4.2 / -5.8 / -4.0 %, TP=2 -19.8 / -13.0 / -15.8 %, TP=1 -26.7 / -32.0 / -38.5 %; non-spec decode of 1-64 requests -1.7 to -9.2 %; * per-tensor fp8, fp8 dots: TP=1 -16.7 / -23.2 / -31.8 %, TP=2 -6.5 / +2.8 / -3.5 %, TP=4 within 1 %; non-spec -0.9 to +3.6 %; * bf16: TP=1 -31.9 / -33.4 / -34.6 %, TP=2 -5.1 / -7.8 / -8.4 %, TP=4 +5.5 / -1.1 / -6.4 %; non-spec -7.6 to +2.5 %. The TP=4 decode loss comes from the new pad at 144-192 rows only (other row counts within 1 %): no resource stalls and fewer wave cycles, but a longer tail. Per-tensor fp8 with bf16 dots now runs the spec calls in 0.76-0.99x the time of fp8_ds_mla caches. GLM-5.3-Flash (PR table shapes, C = 1-16,384): 0.90-0.97x the time at 16 heads, 0.74-1.01x at 32, 0.56-0.95x at 64. fp8_dsv4_mla and separated-rope launches (GLM-5.2) compile to the same code. Outputs: 32-head programs equal 16-head ones; 64-head programs and 16x16x32 change the rounding, with the error against an f64 reference unchanged. The two-loop test adds 96 / 192 rows and 64 heads, which reach the new programs, and checks the bf16 two-loop against the reference rather than against the fp8_scalar output: on one new case the two sat 0.023 apart before this change too, each within 0.015 of the reference.
…und of 16-head ones _staged_head_block took 64-head programs once 2 x rows >= CUs. At exactly half the CUs (128 rows on 256) the 16-head grid is one full round of two programs per CU, and the 64-head grid leaves half the CUs idle: TP=1 decode of 128 requests became 18-25 % slower than on 16-head programs (GLM-5.3-Flash 0.85x and DeepSeek-V4.1 on bf16 caches 0.95x against main). Take 64-head programs only past that round; 128 rows now get 32-head programs. TP=1 decode at 128 rows (MI355X, 32K context, same-GPU pairs, time against 16-head programs): GLM-5.3 1.25x -> 0.83x, DeepSeek-V4.1 bf16 1.18x -> 0.91x, per-tensor fp8 1.19x -> 0.86x. 96-112 and 144-192 rows keep their programs.
Drop the per-tensor fp8 two-loop test file, and run the vLLM-pool and epilogue tests once per kernel path instead of across rows, heads and forced split counts: decode with split-K (6 rows, 16 heads), 32-head programs (192 rows) and prefill on 64-head programs (2,100 rows); the epilogue in the reduce's store and in the attention kernel's. The launcher picks every config. The branch now adds 10 test cases to main's instead of 421.
…the staged-walk optimizations GLM-5.2 keeps per-tensor fp8 rows of 576: the 512-wide latent and a 64-wide rope part, which the walk gathers into its own LDS buffer and dots after the latent. 9def36f left those rows out, so GLM-5.2 still ran 16-head programs on 16x16x16 dots. * 16x16x32 dots and the latent tile's pad (8 per 256, 64 per 8,192) cover per-tensor fp8 rows with the rope appended too. * 64-head programs keep Q's latent in LDS and its rope part in registers, and the QK adds the rope dot on both Q paths; an fp8 q's rope part is widened to bf16 there as well. * The MLA launcher's head-block rule takes those rows, and now runs only at 32 and 64 heads: below 16 heads 9def36f had replaced its 8-head programs with 16-head ones (the stores mask the heads, so the output was right, at up to twice the work). GLM-5.2 on the PR table shapes (MI355X, same-GPU pairs against the previous commit; time, has_invalid off as GLM-5.2 runs): 16 heads 0.92-1.01x (geometric mean 0.95-0.97), 32 heads 0.74-1.04x (0.87-0.94), 64 heads 0.71-1.08x (0.83-0.92). With has_invalid on the walk sits at 256 VGPRs and spills, and launches that stay on 16-head programs get up to 1.20x slower. Other cache formats compile to the same code. A 64-head test covers the latent-in-LDS path with and without a rope part.
…unch kwargs where they apply pa_decode_sparse picked fp8_dsv4_mla's 32- and 64-head programs in two inline blocks and the staged walks' through _staged_head_block. Both rules are now helpers (_packed_head_block next to _staged_head_block, sharing _last_round_fills), and one place calls them at 32 or 64 heads. Two launch arguments only go where the kernel reads them: PIPE_PREFETCH on fp8_dsv4_mla's 32-head programs (bf16 32-head launches got True for a walk that ignores it), and KV_LDS_PAD on the walks without a tile pad of their own (the staged walks got 16 at prefill and ignored it). Every launch over config_dump.py's 29,464 shapes keeps its program size, grid and arguments; only those two arguments change, and the kernels compile to the same code.
The fp8_dsv4_mla launcher runs 8 heads (TP8 of a 64-head model) as a 16-head program with the other 8 masked. c013fe4 left the vLLM-pool test without a case for it; 37 rows at 8 heads (decode with split-K) put one back. Checked at 8 heads beyond the suite: the two-loop at 6-2,100 rows with both sentinel kinds, the epilogue in both stores, and bf16 and per-tensor fp8 caches against an f64 reference.
Comments and docstrings the branch added, cut to what the code does, without measurements or history. No code changes.
…ntinels up to 16 heads With has_invalid on, GLM-5.2's walk (per-tensor fp8, rope appended) sits at 256 VGPRs in programs of up to 16 heads, and 16x16x32 made it slower than main there. Those programs keep 16x16x16 and main's tile pad (KV_LDS_PAD at prefill for rows with the rope appended); 32- and 64-head programs keep 16x16x32. GLM-5.2 with has_invalid on, PR table shapes, time against main (geometric mean, range): 8 heads 0.98 (0.93-1.07) -> 1.00 (0.99-1.01), 16 heads 1.04 (0.95-1.21) -> 1.00 (0.98-1.01), 32 heads 0.93 -> 0.91 (0.77-1.01), 64 heads 0.86 -> 0.85 (0.70-1.02). has_invalid off is unchanged.
Main's test files, plus one test of the inverse-RoPE / MXFP8 store on the SWA + top-k two-loop: split-K (the reduce writes the output) and one 32-head program per row (the attention kernel writes it), with and without MXFP8, against the reference rotated back.
…clearer names No launch changes and the kernels compile to the same code. * The walks are named for what they read: _dsv4_segment (fp8_dsv4_mla, was _pipe_segment), _staged_segment (gathered into registers and staged in LDS, was _process_segment) and _async_segment; their helpers follow (_dsv4_slot_ids, _dsv4_gather, _dsv4_scales, _split_slot_pow2). PIPE_PREFETCH becomes DSV4_PREFETCH. * _dsv4_segment runs one tile loop with and without PREFETCH. * _segment picks the walk for both segments; _first_valid_slot and _with_target are the sentinel-target search and Seg rebuild both walks shared. * _row_reduce_lds (was _xwarp_rowred), _scaled_upcast_asm (was _deq_asm), _widen_scales; _read_slots goes through _read_slots_at. * _dsv4_scales widens the exponents, so _stage no longer infers that from the shape; the exp2 dequant branch fp8_dsv4_mla no longer reached goes.
… round rule _packed_head_block and _staged_head_block become _dsv4_block_m and _staged_block_m, and the round test they both ran is _rounds_suit_32. Every launch keeps its program size, grid and arguments.
…ring The module docstring lists the three tile walks and where the output epilogue runs; a few comments follow the renames. No code change.
vgokhale
previously approved these changes
Oct 5, 2026
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The new DSV4 gather misaddresses accepted non-power-of-two pages, and epilogue validation and coverage have gaps.
Review effort: Balanced
Findings: 3
Open (8)
What changed in this PR
Optimizes GFX950 sparse MLA attention for DeepSeek and GLM workloads.
Changes:
- Adds fused inverse-RoPE and MXFP8 output handling.
- Introduces XCD remapping, wider head tiles, LDS reductions, and improved staging.
- Extends cache handling and adds epilogue tests.
| File | Description |
|---|---|
aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py |
Implements kernel and epilogue optimizations. |
aiter/ops/triton/attention/pa_decode_sparse.py |
Adds dispatch, tuning, and epilogue APIs. |
aiter/ops/triton/attention/sparse_mla.py |
Extends cache formats and launcher behavior. |
op_tests/triton_tests/attention/test_pa_decode_sparse.py |
Tests fused inverse-RoPE and MXFP8 output. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
… and XCD remap on gfx950
…s, 1-D cos/sin caches and off-device MXFP8 buffers
Contributor
Author
|
Addressed copilot reviews and handled conflicts to keep gfx942 path working |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The LDS Q path can omit trailing dimensions, two epilogue arguments can be silently ignored, and scalar-FP8 two-loop coverage is missing.
Review effort: Balanced
Findings: 1
Open (1)
Resolved since last review (8)
…eject non-power-of-two kv_lora_rank / qk_rope_head_dim
vgokhale
approved these changes
Oct 7, 2026
michaelzhang-ai
added a commit
to michaelzhang-ai/sglang
that referenced
this pull request
Oct 15, 2026
…e commit The previous commit dropped the ROCm/aiter#6042 cherry-pick on the wrong premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the sgl-project#6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of sgl-project#6159 and applies cleanly to the pin, so cherry-pick that instead. It is on main, so the PR-ref fetch is no longer needed. sgl-project#5967 and sgl-project#6007 are in the pin, so their patches stay removed.
michaelzhang-ai
added a commit
to sgl-project/sglang
that referenced
this pull request
Oct 15, 2026
The previous commit dropped the ROCm/aiter#6042 cherry-pick on the wrong premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the #6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of #6159 and applies cleanly to the pin, so cherry-pick that instead. It is on main, so the PR-ref fetch is no longer needed. #5967 and #6007 are in the pin, so their patches stay removed.
michaelzhang-ai
added a commit
to michaelzhang-ai/sglang
that referenced
this pull request
Oct 15, 2026
… commit sgl-project#44534 dropped the sgl-project#6042 cherry-pick on the premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the sgl-project#6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of sgl-project#6159 and applies cleanly to the pin, and it is on main, so the PR-ref fetch is no longer needed.
michaelzhang-ai
added a commit
to sgl-project/sglang
that referenced
this pull request
Oct 15, 2026
… commit #44534 dropped the #6042 cherry-pick on the premise that the 2bb549102 pin already contains it. It does not: the pin (Oct 10) predates the #6042 merge (Oct 13). The PR-branch commit 042a28918 conflicts because ROCm/aiter#6159 rewrote pa_decode_sparse.py between the old and new pins. The squash-merge commit 57e2c6b42 sits on top of #6159 and applies cleanly to the pin, and it is on main, so the PR-ref fetch is no longer needed.
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
Summary of the added optimizations:
The main target is dsv4 like sparse MLA, while enabling certain optimizations for GLM where it helps.
Technical Details
Comparison against main, ignoring the fusion path as that is orthogonal optimization. This isolates the benefit of the kernel/config changes.
DeepSeek-V4.1-Flash · fp8_dsv4_mla KV
40 layers: 2 SWA-only, 18 c2, 20 c1 (top-k 512, SWA 128).
DeepSeek-V4-Flash · fp8_dsv4_mla KV
43 layers: 2 SWA-only, 21 c4 (top-k 512), 20 c128 (every compressed row).
GLM-5.2 · per-tensor fp8 KV
78 layers, 576-wide rows (rope appended), 2,048 indices per row. has_invalid off.
GLM-5.3-Flash · per-tensor fp8 KV
rope-free 512-wide rows, 2,048 indices per row (one layer).
DeepSeek-V4.1-Flash · bf16 KV
DeepSeek-V4.1-Flash's layers on bf16 rows.