From c6f04892eb4fe2a5862bc7d06c8d701b7d7a274f Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 06:03:56 +0000 Subject: [PATCH 01/51] [GLUON][GFX950][DSV4] Split-count collapse, cs0 alignment hint, per-cache 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) --- .../gfx950/attention/pa_decode_sparse.py | 23 ++++++- .../ops/triton/attention/pa_decode_sparse.py | 66 ++++++++++++++++++- 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 54d526ca51b..ee56afbed1c 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -658,6 +658,7 @@ def _pa_decode_sparse( EXTRA_IS_FP8: gl.constexpr, MAIN_BLOCK_SIZE: gl.constexpr, EXTRA_BLOCK_SIZE: gl.constexpr, + CS0_ALIGN: gl.constexpr, NOPE_DIM: gl.constexpr, ROPE_DIM: gl.constexpr, HEAD_SIZE: gl.constexpr, @@ -667,7 +668,13 @@ def _pa_decode_sparse( HEAD_ALIGNED: gl.constexpr, MFMA_K: gl.constexpr, UNIFORM: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, + # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache + # whose span exceeds that must gather via 64-bit gl.load -- but the two caches + # are sized independently (SWA window vs full compressed history), so gating + # them together would drop the fast path on a small cache just because its + # partner is large. At tiny top-k the main/SWA gather is ~94% of the tokens. + MAIN_USE_BUFFER_LOAD: gl.constexpr, + EXTRA_USE_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -676,6 +683,16 @@ def _pa_decode_sparse( un-normalized partials for the reduce kernel. HAS_INVALID gates -1-sentinel handling (clamp + score mask) on the full-tile fast paths.""" NUM_WARPS: gl.constexpr = gl.num_warps() + # Page stride alignment hint. The row base of every cache gather is + # ``block_idx*cs0 + pos*576`` (packed) / ``... + pos*HEAD_SIZE`` (uniform) with + # block_idx/pos gathered at runtime, so the divisibility analysis has to assume + # 1-byte alignment and a contiguous 512-byte row gather lowers to hundreds of + # global_load_ubyte instead of global_load_dwordx4. 576 and HEAD_SIZE are + # literals the compiler can already reason about; cs0 is the one opaque term. + # The driver sets CS0_ALIGN>1 only after checking the strides on the host. + if CS0_ALIGN > 1: + main_cs0 = gl.multiple_of(main_cs0, CS0_ALIGN) + extra_cs0 = gl.multiple_of(extra_cs0, CS0_ALIGN) query_idx = gl.program_id(0) split_id = gl.program_id(1) pid_h = gl.program_id(2) @@ -800,7 +817,7 @@ def _pa_decode_sparse( BLOCK_K, HEAD_ALIGNED, UNIFORM, - USE_BUFFER_LOAD, + MAIN_USE_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -845,7 +862,7 @@ def _pa_decode_sparse( BLOCK_K, HEAD_ALIGNED, UNIFORM, - USE_BUFFER_LOAD, + EXTRA_USE_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index c6fafb181e6..0f95b0fd79f 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -452,6 +452,30 @@ def _decode_num_splits( cost = waves * (m_it + e_it) + GAMMA * splits + DELTA * fill if best_cost is None or cost < best_cost - 1e-9: best_splits, best_cost = splits, cost + + # Collapse to the FEWEST splits that keeps both the wave count and the + # per-split BLOCK_K iteration count. The cost model above treats extra + # splits as nearly free while ``waves`` stays 1 (GAMMA is small and it has no + # reduce term), so at small per-query work it over-splits: at C=64/H=16 with + # main=128, extra=8 it picked 3, where 2 does the same 2 iterations in the + # same single wave but launches 1/3 fewer CTAs and reduces over 2 partials + # instead of 3 (measured 16.2 -> 14.4 us on gfx950/MI355X). + # + # Ported from vLLM's _decode_gfx950_num_splits (PR #52212), which fixed the + # same over-splitting in the in-tree triton decode. + def _iters(s): + m = math.ceil(math.ceil(avg_main / s) / block_k) if avg_main > 0 else 0 + e = math.ceil(math.ceil(avg_extra / s) / block_k) if avg_extra > 0 else 0 + return m + e + + def _waves(s): + return (base * s + cu - 1) // cu + + if best_splits > 1: + target_waves, target_iters = _waves(best_splits), _iters(best_splits) + for splits in range(1, best_splits): + if _waves(splits) == target_waves and _iters(splits) == target_iters: + return splits return best_splits @@ -485,7 +509,18 @@ def _pa_decode_sparse_gfx950_gluon( # Tuned launch config (gfx950 / MI355), inlined. BLOCK_M = heads per MFMA M-tile; # BLOCK_K = KV tile; num_warps = BLOCK_K // 16 (warps tile the dot-N, MFMA N=16). BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 - num_warps = BLOCK_K // 16 + # AITER_PA_DECODE_BLOCK_K: experiment override for the KV tile width. + # num_warps stays BLOCK_K//16 (warps tile the dot-N, MFMA N=16). + import os as _os + _bk = _os.environ.get("AITER_PA_DECODE_BLOCK_K") + if _bk: + BLOCK_K = int(_bk) + # AITER_PA_DECODE_MFMA_K: 32 selects CDNA4's double-rate v_mfma_f32_16x16x32_bf16 + # (16 does 16x16x16, i.e. half the K per instruction). + _mk = _os.environ.get("AITER_PA_DECODE_MFMA_K") + if _mk: + MFMA_K = int(_mk) + num_warps = max(1, BLOCK_K // 16) NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -551,7 +586,30 @@ def _pa_decode_sparse_gfx950_gluon( avg_main = indices.numel() / max(1, num_queries) avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 - use_buffer_load = cache_bytes < MAX_BYTES + # Largest power-of-2 (<=16) dividing BOTH page strides. Lets the kernel assert + # 16B-aligned row bases so a contiguous row gather can vectorize to dwordx4. + _s0 = int(cache.stride(0)); _s1 = int(extra_cache.stride(0)) + cs0_align = 1 + for _a in (16, 8, 4, 2): + if _s0 % _a == 0 and _s1 % _a == 0: + cs0_align = _a + break + # Gate each cache on its OWN span: one oversized cache must not disable the + # fast path for the other. (cache_bytes, the max of the two, is kept for the + # arch/format asserts below.) + main_use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES + extra_use_buffer_load = max_addressable_bytes(extra_cache) < MAX_BYTES + use_buffer_load = main_use_buffer_load and extra_use_buffer_load + # AITER_PA_DECODE_FORCE_GLOBAL_LOAD emulates a >2GB production cache without + # allocating one: 1 = both caches, main = main only, extra = extra only. + _fg = _os.environ.get("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "") + if _fg == "1": + main_use_buffer_load = extra_use_buffer_load = False + elif _fg == "main": + main_use_buffer_load = False + elif _fg == "extra": + extra_use_buffer_load = False + use_buffer_load = main_use_buffer_load and extra_use_buffer_load HEAD_ALIGNED = num_heads % BLOCK_M == 0 heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M out = torch.empty_like(q, dtype=torch.bfloat16) @@ -620,6 +678,7 @@ def _pa_decode_sparse_gfx950_gluon( EXTRA_IS_FP8=extra_is_fp8, MAIN_BLOCK_SIZE=main_block, EXTRA_BLOCK_SIZE=extra_block, + CS0_ALIGN=cs0_align, NOPE_DIM=nope_dim, ROPE_DIM=ROPE_DIM, HEAD_SIZE=head_dim, @@ -629,7 +688,8 @@ def _pa_decode_sparse_gfx950_gluon( HEAD_ALIGNED=HEAD_ALIGNED, MFMA_K=MFMA_K, UNIFORM=UNIFORM, - USE_BUFFER_LOAD=use_buffer_load, + MAIN_USE_BUFFER_LOAD=main_use_buffer_load, + EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, num_warps=num_warps, From 23e8319da1d2f6055b8d35831c860de8cc58b472 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 06:33:04 +0000 Subject: [PATCH 02/51] [GLUON][GFX950][DSV4] Chunked fp8 dequant + parallel split-K reduce 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) --- .../gfx950/attention/pa_decode_sparse.py | 183 +++++++++++++++--- .../ops/triton/attention/pa_decode_sparse.py | 49 ++++- 2 files changed, 199 insertions(+), 33 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index ee56afbed1c..a794c8cf627 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -44,6 +44,100 @@ def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) +@gluon.jit +def _split2(x): + """Contiguous register split along dim 1: [A, B] -> two [A, B//2]. + + x0 takes columns [0, B//2), x1 takes [B//2, B). Both halves come back in the + input's own layout, so with warps tiling dim 0 (see ``gather_l``) the column + direction is a pure per-lane register repeat and the split is a rename -- no + cross-lane traffic. ``assert_trivial`` makes a non-free split a compile error + rather than a silent LDS round-trip. + """ + layout: gl.constexpr = x.type.layout + x_r = x.reshape(x.shape[0], 2, x.shape[1] // 2).permute(0, 2, 1) + x0, x1 = gl.split(x_r) + x0 = gl.convert_layout(x0, layout, assert_trivial=True) + x1 = gl.convert_layout(x1, layout, assert_trivial=True) + return x0, x1 + + +@gluon.jit +def _deq_store(x_u8, sc, kv_smem, off, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr): + """Dequant one fp8 slab and write it straight to kv_smem[:, off:off+W]. + + Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path lowers to + an emulated mul and doubles the loop. ``sc`` is the raw per-64-group exponent + byte (packed fp8_ds_mla, UE8M0) or an f32 scale (uniform pool). + """ + W: gl.constexpr = x_u8.shape[1] + if UNIFORM: + scale = sc + else: + scale = gl.exp2(sc.to(gl.float32) - 127.0) + kv_smem.slice(off, W, dim=1).store( + (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) + ) + + +@gluon.jit +def _deq_store_tile( + x_u8, + sc, + kv_smem, + NOPE_CHUNK: gl.constexpr, + UNIFORM: gl.constexpr, + FP8_FNUZ: gl.constexpr, +): + """Dequant a whole gathered fp8 tile into kv_smem in NOPE_CHUNK-wide pieces. + + The f32 expansion is 4x the fp8 tile (a [64, 512] tile is 128 f32 VGPRs/lane), + so materializing it in one shot is what pins 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 -- so only HEAD_SIZE/NOPE_CHUNK-th of it is ever + live. The splits themselves are register renames (see ``_split2``). + """ + W: gl.constexpr = x_u8.shape[1] + if NOPE_CHUNK >= W: + _deq_store(x_u8, sc, kv_smem, 0, UNIFORM, FP8_FNUZ) + else: + x0, x1 = _split2(x_u8) + s0, s1 = _split2(sc) + W2: gl.constexpr = W // 2 + if NOPE_CHUNK >= W2: + _deq_store(x0, s0, kv_smem, 0, UNIFORM, FP8_FNUZ) + _deq_store(x1, s1, kv_smem, W2, UNIFORM, FP8_FNUZ) + else: + x00, x01 = _split2(x0) + s00, s01 = _split2(s0) + x10, x11 = _split2(x1) + s10, s11 = _split2(s1) + W4: gl.constexpr = W // 4 + if NOPE_CHUNK >= W4: + _deq_store(x00, s00, kv_smem, 0, UNIFORM, FP8_FNUZ) + _deq_store(x01, s01, kv_smem, W4, UNIFORM, FP8_FNUZ) + _deq_store(x10, s10, kv_smem, 2 * W4, UNIFORM, FP8_FNUZ) + _deq_store(x11, s11, kv_smem, 3 * W4, UNIFORM, FP8_FNUZ) + else: + W8: gl.constexpr = W // 8 + y0, y1 = _split2(x00) + t0, t1 = _split2(s00) + _deq_store(y0, t0, kv_smem, 0, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, W8, UNIFORM, FP8_FNUZ) + y0, y1 = _split2(x01) + t0, t1 = _split2(s01) + _deq_store(y0, t0, kv_smem, 2 * W8, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 3 * W8, UNIFORM, FP8_FNUZ) + y0, y1 = _split2(x10) + t0, t1 = _split2(s10) + _deq_store(y0, t0, kv_smem, 4 * W8, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 5 * W8, UNIFORM, FP8_FNUZ) + y0, y1 = _split2(x11) + t0, t1 = _split2(s11) + _deq_store(y0, t0, kv_smem, 6 * W8, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 7 * W8, UNIFORM, FP8_FNUZ) + + @gluon.jit def _slots( indices_ptr, @@ -106,6 +200,7 @@ def _decode_tile( HEAD_SIZE: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, + NOPE_CHUNK: gl.constexpr, HEAD_ALIGNED: gl.constexpr, MASKED: gl.constexpr, UNIFORM: gl.constexpr, @@ -154,7 +249,7 @@ def _decode_tile( else: x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) - kv_smem.store((_fp8_to_f32(x_u8, FP8_FNUZ) * sc).to(gl.bfloat16)) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, True, FP8_FNUZ) elif IS_FP8: # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. nope_off = (block_idx_g * cs0 + pos_g * 576)[:, None] + offs_full[None, :] @@ -171,12 +266,7 @@ def _decode_tile( else: x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) - # Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path - # lowers to an emulated mul and doubles the loop. - x_fp8 = x_u8.to(gl.float8e4nv, bitcast=True) - scales = gl.exp2(exps.to(gl.float32) - 127.0) - k_nope = (x_fp8.to(gl.float32) * scales).to(gl.bfloat16) - kv_smem.store(k_nope) + _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, False, FP8_FNUZ) block_idx_gr, pos_gr, valid_gr = _slots( indices_ptr, seg_start, @@ -270,8 +360,14 @@ def _gd_fp8( HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): - """Gather + dequant one full fp8 tile to bf16 regs. Split from the LDS-write/MFMA - so the gather issues an iteration early.""" + """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues + an iteration early. + + The prefetch is carried across the MFMA in **raw fp8**, not dequantized bf16: + the loop keeps two tiles in flight, and a [BLOCK_K, HEAD_SIZE] tile is 32 + VGPRs/lane as u8 but 64 as bf16, so dequantizing here would burn an extra 64 + VGPRs on loop-carried state alone. The consumer dequants (chunked) instead. + """ if not USE_BUFFER_LOAD: cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) bg, pg, valid = _slots( @@ -290,19 +386,14 @@ def _gd_fp8( scl_off = (bg * NGRP)[:, None] + (offs_full[None, :] // 64) x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) - k_nope = (_fp8_to_f32(x_u8, FP8_FNUZ) * sc).to(gl.bfloat16) - k_rope = k_nope # unused for UNIFORM (rope slice-store skipped) -> DCE'd + k_rope = x_u8 # unused for UNIFORM (rope slice-store skipped) -> DCE'd else: nope_off = (bg * cs0 + pg * 576)[:, None] + offs_full[None, :] scl_off = (bg * cs0 + BLOCK_SIZE * 576 + pg * 8)[:, None] + ( offs_full[None, :] // 64 ) x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) - exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) - k_nope = ( - x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) - * gl.exp2(exps.to(gl.float32) - 127.0) - ).to(gl.bfloat16) + sc = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) bgr, pgr, _ = _slots( indices_ptr, seg_start, @@ -315,12 +406,13 @@ def _gd_fp8( ) rope_off = (bgr * (cs0 // 2) + pgr * 288 + 224)[:, None] + offs_rope[None, :] k_rope = _cache_load(cache_bf16_ptr, rope_off, USE_BUFFER_LOAD) - return k_nope, k_rope, valid + return x_u8, sc, k_rope, valid @gluon.jit def _qkpv_fp8( - k_nope, + x_u8, + sc, k_rope, valid, q_dot, @@ -340,15 +432,17 @@ def _qkpv_fp8( HEAD_SIZE: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, + NOPE_CHUNK: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, + FP8_FNUZ: gl.constexpr, HAS_INVALID: gl.constexpr, ): - """Write a prefetched bf16 tile to LDS, then QK -> softmax -> PV. UNIFORM skips - the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, mask the - columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" + """Dequant a prefetched fp8 tile into LDS, then QK -> softmax -> PV. UNIFORM + skips the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, + mask the columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" neg_inf = float("-inf") - kv_smem.store(k_nope) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, UNIFORM, FP8_FNUZ) if not UNIFORM: kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) k = kv_smem.permute([1, 0]).load(k_layout) @@ -416,6 +510,7 @@ def _process_segment( HEAD_SIZE: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, + NOPE_CHUNK: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, @@ -434,7 +529,7 @@ def _process_segment( if IS_FP8: n_full = (hi_full - lo) // BLOCK_K if n_full > 0: - kn, kr, vld = _gd_fp8( + kn, ks, kr, vld = _gd_fp8( cache_ptr, cache_bf16_ptr, indices_ptr, @@ -455,7 +550,7 @@ def _process_segment( FP8_FNUZ, ) for i in range(1, n_full): - kn2, kr2, vld2 = _gd_fp8( + kn2, ks2, kr2, vld2 = _gd_fp8( cache_ptr, cache_bf16_ptr, indices_ptr, @@ -477,6 +572,7 @@ def _process_segment( ) m_i, l_i, acc = _qkpv_fp8( kn, + ks, kr, vld, q_dot, @@ -496,13 +592,16 @@ def _process_segment( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, UNIFORM, + FP8_FNUZ, HAS_INVALID, ) - kn, kr, vld = kn2, kr2, vld2 + kn, ks, kr, vld = kn2, ks2, kr2, vld2 m_i, l_i, acc = _qkpv_fp8( kn, + ks, kr, vld, q_dot, @@ -522,8 +621,10 @@ def _process_segment( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, UNIFORM, + FP8_FNUZ, HAS_INVALID, ) else: @@ -562,6 +663,7 @@ def _process_segment( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, False, UNIFORM, @@ -605,6 +707,7 @@ def _process_segment( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, True, UNIFORM, @@ -668,6 +771,12 @@ def _pa_decode_sparse( HEAD_ALIGNED: gl.constexpr, MFMA_K: gl.constexpr, UNIFORM: gl.constexpr, + # Experiment knobs (driver reads AITER_PA_DECODE_* env vars; defaults = prior + # behaviour). GATHER_W0: put the NoPE gather's warps on dim 0 (rows) instead of + # dim 1 (columns). + GATHER_W0: gl.constexpr, + # NOPE_CHUNK: width of one dequant piece (HEAD_SIZE = one shot, prior behaviour). + NOPE_CHUNK: gl.constexpr, # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache # whose span exceeds that must gather via 64-bit gl.load -- but the two caches # are sized independently (SWA window vs full compressed history), so gating @@ -717,10 +826,17 @@ def _pa_decode_sparse( # 16 uint8 = 128-bit fp8 gather loads GSPT: gl.constexpr = 16 + # Warps on dim 1 (columns) vs dim 0 (rows). Coalescing is identical either way + # -- a row's 128 contiguous bytes always come from 8 threads -- but the row + # index vector lives in SliceLayout(1, gather_l), so warps-on-dim-1 replicates + # it across all NUM_WARPS and every lane carries BLOCK_K/8 slots, while + # warps-on-dim-0 carries BLOCK_K/(8*NUM_WARPS). Dim 0 also makes the column + # direction a pure per-lane register repeat, so a chunked dequant can split it + # with a trivial (register-renaming) layout convert. gather_l: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, GSPT], threads_per_warp=[8, 8], - warps_per_cta=[1, NUM_WARPS], + warps_per_cta=[NUM_WARPS, 1] if GATHER_W0 else [1, NUM_WARPS], order=[1, 0], ) # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting @@ -815,6 +931,7 @@ def _pa_decode_sparse( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, UNIFORM, MAIN_USE_BUFFER_LOAD, @@ -860,6 +977,7 @@ def _pa_decode_sparse( HEAD_SIZE, BLOCK_M, BLOCK_K, + NOPE_CHUNK, HEAD_ALIGNED, UNIFORM, EXTRA_USE_BUFFER_LOAD, @@ -949,15 +1067,26 @@ def _pa_decode_sparse_reduce( """Split-KV combine: merge the per-split partials, fold the attn sink, and write the final output. Partials store m in the base-2 exponent domain (row-max * softmax_scale * log2e), matching the triton reduce. Grid: (num_queries, heads_blocks). + + BLOCK_M is the reduce's own head tile and is deliberately decoupled from the + attention kernel's: the combine is pure bandwidth over + [num_queries, num_heads, HEAD_SIZE] f32, so the only thing that matters is + having enough workgroups to cover the machine. At BLOCK_M = num_heads there is + one workgroup per query -- 64 of them on a 256-CU part, i.e. 3/4 of the GPU + idle and one wave to hide every partial-load latency behind. """ NUM_WARPS: gl.constexpr = gl.num_warps() RCP_LN2: gl.constexpr = 1.4426950408889634 query_idx = gl.program_id(0) pid_h = gl.program_id(1) + # One warp covers [BLOCK_M, 64//BLOCK_M * 8] -- lay the 64 lanes out so a + # small BLOCK_M spends them on the head dim instead of idling them on rows. + TPW0: gl.constexpr = BLOCK_M if BLOCK_M < 8 else 8 + TPW1: gl.constexpr = 64 // TPW0 BLK: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, 8], - threads_per_warp=[8, 8], + threads_per_warp=[TPW0, TPW1], warps_per_cta=[1, NUM_WARPS], order=[1, 0], ) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 0f95b0fd79f..7636a8cbb8c 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -267,6 +267,9 @@ def pa_decode_sparse( # mismatches between 2D (m/l) and 3D (acc) loads. reduce_num_warps = 1 if use_gluon else 4 reduce_waves_per_eu = 4 if use_gluon else 1 + _rw = _os.environ.get("AITER_PA_DECODE_REDUCE_WARPS") + if _rw and use_gluon: + reduce_num_warps = int(_rw) USE_EXP2 = True # Infer KV_SPLITS from inputs when caller doesn't override. @@ -508,7 +511,9 @@ def _pa_decode_sparse_gfx950_gluon( # Tuned launch config (gfx950 / MI355), inlined. BLOCK_M = heads per MFMA M-tile; # BLOCK_K = KV tile; num_warps = BLOCK_K // 16 (warps tile the dot-N, MFMA N=16). - BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 + # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD + # threshold on gfx950's 512-VGPR file. Reachable only with the chunked dequant. + BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 2 # AITER_PA_DECODE_BLOCK_K: experiment override for the KV tile width. # num_warps stays BLOCK_K//16 (warps tile the dot-N, MFMA N=16). import os as _os @@ -521,6 +526,9 @@ def _pa_decode_sparse_gfx950_gluon( if _mk: MFMA_K = int(_mk) num_warps = max(1, BLOCK_K // 16) + # AITER_PA_DECODE_GATHER_W0: NoPE gather warps on dim 0 (rows) instead of dim 1. + GATHER_W0 = _os.environ.get("AITER_PA_DECODE_GATHER_W0", "1") == "1" + # AITER_PA_DECODE_NOPE_CHUNK: width of one fp8 dequant piece (0/unset = one shot). NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -614,7 +622,11 @@ def _pa_decode_sparse_gfx950_gluon( heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M out = torch.empty_like(q, dtype=torch.bfloat16) - if kv_splits is not None: + # AITER_PA_DECODE_SPLITS: experiment override for the split-K count. + _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") + if _sp: + num_splits = max(1, int(_sp)) + elif kv_splits is not None: num_splits = max(1, int(kv_splits)) else: num_splits = _decode_num_splits( @@ -641,6 +653,16 @@ def _pa_decode_sparse_gfx950_gluon( part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 + # 128 is the narrowest piece the gather layout can split to for free (below + # that a split falls inside a thread's 16-byte run, which Triton rejects). + _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") + nope_chunk = int(_nc) if _nc else min(128, head_dim) + # AITER_PA_DECODE_WPEU: amdgpu-waves-per-eu floor. 2 caps the allocator at 256 + # unified VGPRs (the 2 waves/SIMD threshold) instead of letting it run to 512. + _wp = _os.environ.get("AITER_PA_DECODE_WPEU") + if _wp: + waves_per_eu = int(_wp) + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, @@ -688,6 +710,8 @@ def _pa_decode_sparse_gfx950_gluon( HEAD_ALIGNED=HEAD_ALIGNED, MFMA_K=MFMA_K, UNIFORM=UNIFORM, + GATHER_W0=GATHER_W0, + NOPE_CHUNK=nope_chunk, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, HAS_INVALID=has_invalid, @@ -701,7 +725,20 @@ def _pa_decode_sparse_gfx950_gluon( if skip_reduce: return part_acc, part_m, part_l - rgrid = (num_queries, heads_blocks) + # The combine is pure bandwidth over [num_queries, num_heads, head_dim] f32, + # so size its head tile for workgroup count, not for the attention kernel's + # BLOCK_M: at BLOCK_M=num_heads there is one workgroup per query (64 on a + # 256-CU part). One head per workgroup gives num_queries*num_heads of them and + # enough concurrent waves to hide the partial-load latency. + red_block_m = 1 + _rbm = _os.environ.get("AITER_PA_DECODE_REDUCE_BLOCK_M") + if _rbm: + red_block_m = int(_rbm) + red_warps = 1 + _rw = _os.environ.get("AITER_PA_DECODE_REDUCE_WARPS") + if _rw: + red_warps = int(_rw) + rgrid = (num_queries, (num_heads + red_block_m - 1) // red_block_m) _pa_decode_sparse_reduce_gfx950[rgrid]( part_m, part_l, @@ -718,9 +755,9 @@ def _pa_decode_sparse_gfx950_gluon( num_heads, HAS_SINK=has_sink, HEAD_SIZE=head_dim, - BLOCK_M=BLOCK_M, + BLOCK_M=red_block_m, NUM_SPLITS=num_splits, - HEAD_ALIGNED=HEAD_ALIGNED, - num_warps=4, + HEAD_ALIGNED=num_heads % red_block_m == 0, + num_warps=red_warps, ) return out From f944ac3446175c8be222a968c575c649e8cd2e69 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 06:38:29 +0000 Subject: [PATCH 03/51] [GLUON][GFX950][DSV4] pa_decode_sparse: accept a caller-supplied out= 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) --- .../ops/triton/attention/pa_decode_sparse.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 7636a8cbb8c..cd9e9813be5 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -53,6 +53,17 @@ _FP8_DTYPE = torch.float8_e4m3fnuz +def _check_out(out, q, dtype): + """Caller-supplied output buffer, or a fresh one. Writing the caller's buffer + directly saves a full [T, H, D] device copy per call.""" + if out is None: + return torch.empty_like(q, dtype=dtype) + 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 + + def pa_decode_sparse( q: torch.Tensor, unified_kv: torch.Tensor, @@ -70,6 +81,7 @@ def pa_decode_sparse( extra_cache: torch.Tensor | None = None, extra_indices: torch.Tensor | None = None, extra_indptr: torch.Tensor | None = None, + out: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse paged-decode attention with split-K + widened BLOCK_H. @@ -88,6 +100,8 @@ def pa_decode_sparse( auto-infers to fill ~512 total CTAs while capping below the number of K-blocks, then rounds up to a power of 2. num_stages: software-pipeline depth of the K loop (default 2). + out: optional ``[N, H, D]`` destination. Supplied -> written in place and + returned, which saves the caller a full-size device copy. skip_reduce: when the split-K path is active (``kv_splits > 1``), return the pre-reduce ``(acc_partial, m_partial, l_partial)`` partials instead of launching the reduce kernel. Has no effect when @@ -163,6 +177,7 @@ def pa_decode_sparse( skip_reduce=skip_reduce, has_invalid=bool(has_invalid), fp8_fnuz=fp8_fnuz, + out=out, ) assert ( @@ -199,7 +214,7 @@ def pa_decode_sparse( f"PA_DECODE_SPARSE T={T} H={H} D={D} " f"total_indices={kv_indices.shape[0]}" ) - out = torch.empty_like(q) + out = _check_out(out, q, q.dtype) assert kv_indices.dtype == torch.int32 and kv_indices.is_contiguous() assert kv_indptr.dtype == torch.int32 and kv_indptr.is_contiguous() @@ -495,6 +510,7 @@ def _pa_decode_sparse_gfx950_gluon( extra_indptr=None, kv_splits=None, skip_reduce=False, + out=None, has_invalid=False, fp8_fnuz=False, ): @@ -620,7 +636,7 @@ def _pa_decode_sparse_gfx950_gluon( use_buffer_load = main_use_buffer_load and extra_use_buffer_load HEAD_ALIGNED = num_heads % BLOCK_M == 0 heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M - out = torch.empty_like(q, dtype=torch.bfloat16) + out = _check_out(out, q, torch.bfloat16) # AITER_PA_DECODE_SPLITS: experiment override for the split-K count. _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") From 9459334178916a5e046a66eaaec6c1b4e6bb8a6b Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 06:42:09 +0000 Subject: [PATCH 04/51] [GLUON][GFX950][DSV4] Only ask for 2 waves/EU on the buffer_load path 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) --- .../ops/triton/attention/pa_decode_sparse.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index cd9e9813be5..d0be707f614 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -527,9 +527,7 @@ def _pa_decode_sparse_gfx950_gluon( # Tuned launch config (gfx950 / MI355), inlined. BLOCK_M = heads per MFMA M-tile; # BLOCK_K = KV tile; num_warps = BLOCK_K // 16 (warps tile the dot-N, MFMA N=16). - # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD - # threshold on gfx950's 512-VGPR file. Reachable only with the chunked dequant. - BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 2 + BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 # AITER_PA_DECODE_BLOCK_K: experiment override for the KV tile width. # num_warps stays BLOCK_K//16 (warps tile the dot-N, MFMA N=16). import os as _os @@ -670,11 +668,21 @@ def _pa_decode_sparse_gfx950_gluon( pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 # 128 is the narrowest piece the gather layout can split to for free (below - # that a split falls inside a thread's 16-byte run, which Triton rejects). + # that a split falls inside a thread's 16-byte run, which Triton rejects), and + # the split is only trivial when the gather's warps tile dim 0. _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") nope_chunk = int(_nc) if _nc else min(128, head_dim) - # AITER_PA_DECODE_WPEU: amdgpu-waves-per-eu floor. 2 caps the allocator at 256 - # unified VGPRs (the 2 waves/SIMD threshold) instead of letting it run to 512. + if not GATHER_W0: + nope_chunk = head_dim + + # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD + # threshold on gfx950's 512-VGPR file. The chunked dequant gets the buffer_load + # path to 256 with a single spilled VGPR, 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 (+40% at C4A), so only ask for it + # when both caches are on the fast path. + if use_buffer_load and nope_chunk < head_dim: + waves_per_eu = 2 _wp = _os.environ.get("AITER_PA_DECODE_WPEU") if _wp: waves_per_eu = int(_wp) From e079000ef61f622c8be15dcc5f505f3304bee292 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 06:54:48 +0000 Subject: [PATCH 05/51] [GLUON][GFX950][DSV4] Occupancy-aware split-K policy + separate main 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) --- .../gfx950/attention/pa_decode_sparse.py | 13 ++++- .../ops/triton/attention/pa_decode_sparse.py | 54 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index a794c8cf627..74147c0b8a6 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -777,6 +777,15 @@ def _pa_decode_sparse( GATHER_W0: gl.constexpr, # NOPE_CHUNK: width of one dequant piece (HEAD_SIZE = one shot, prior behaviour). NOPE_CHUNK: gl.constexpr, + # How many of the NUM_SPLITS programs share the main (SWA) segment. The two + # segments have very different shapes -- main is a contiguous window whose + # length is fixed by the sliding window, extra is the top-k list -- so one + # split count has to be wrong for one of them. Splitting main past + # main_len/BLOCK_K turns every full tile into a half-empty masked one, while + # extra still wants the CTAs. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole + # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an + # empty main range and contribute an extra-only partial. + MAIN_SPLITS: gl.constexpr, # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache # whose span exceeds that must gather via 64-bit gl.load -- but the two caches # are sized independently (SWA window vs full compressed history), so gating @@ -897,8 +906,8 @@ def _pa_decode_sparse( main_start = gl.load(main_indptr_ptr + query_idx) main_end = gl.load(main_indptr_ptr + query_idx + 1) main_len = main_end - main_start - main_chunk = (main_len + NUM_SPLITS - 1) // NUM_SPLITS - main_lo = split_id * main_chunk + main_chunk = (main_len + MAIN_SPLITS - 1) // MAIN_SPLITS + main_lo = gl.minimum(split_id * main_chunk, main_len) main_hi = gl.minimum(main_lo + main_chunk, main_len) m_i, l_i, acc = _process_segment( q_dot, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index d0be707f614..c9d02d55df3 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -497,6 +497,45 @@ def _waves(s): return best_splits +def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block_k): + """Split-K count for the gfx950 gluon kernel: fill the machine, but never + split a segment finer than one BLOCK_K tile. + + Two facts set this. (a) The kernel is 256 unified VGPRs and 68,608 B of LDS, + so a CU holds two workgroups -- the launch wants ~2*CU programs before extra + splits stop paying. (b) Splitting a segment past its tile count does not divide + the work, it *multiplies* it: every split then owns a partial tile, and a + partial tile costs a full masked gather for a fraction of the tokens. So the + split count is capped by the larger segment's tile count, and the main segment + separately stops at its own (see MAIN_SPLITS in the kernel). + + Measured at C=64 H=16 BLOCK_K=64 (attn = kernel + reduce, us), this rule picks + the sweep optimum at extra in {8, 64, 272, 512, 1024, 2048} and lands within + 1.4% at extra=128: + + extra 8 64 128 272 512 1024 2048 + picked S 2 2 2 5 8 8 8 + best S 2 2 4 5 8 8 8 + old model 2 2 2 3 4 4 4 + old us 14.0 14.6 17.9 21.3 22.0 29.8 45.6 + new us 14.0 14.4 17.4 19.5 20.9 26.9 38.1 + """ + num_sms = get_num_sms() + base_wg = max(1, num_queries * heads_blocks) + cta_cap = max(1, (2 * num_sms) // base_wg) + main_tiles = max(1, math.ceil(avg_main / block_k)) if avg_main > 0 else 0 + extra_tiles = max(1, math.ceil(avg_extra / block_k)) if avg_extra > 0 else 0 + tiles = max(1, main_tiles, extra_tiles) + if base_wg >= num_sms: + # Already at least one workgroup per CU without splitting, so a split buys + # only the second occupancy slot -- worth an extra round trip of the + # [queries, heads, D] f32 partials plus a reduce launch only when each + # split still owns a real run of tiles. At C=256 extra=8 (2 tiles) forcing + # a split costs 31%; at extra=2048 (32 tiles) it saves 19%. + return max(1, min(cta_cap, tiles // 4)) + return max(1, min(cta_cap, tiles)) + + def _pa_decode_sparse_gfx950_gluon( q, cache, @@ -642,6 +681,10 @@ def _pa_decode_sparse_gfx950_gluon( num_splits = max(1, int(_sp)) elif kv_splits is not None: num_splits = max(1, int(kv_splits)) + elif _os.environ.get("AITER_PA_DECODE_SPLIT_POLICY", "occ") == "occ": + num_splits = _decode_num_splits_occ( + num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K + ) else: num_splits = _decode_num_splits( num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K @@ -687,6 +730,16 @@ def _pa_decode_sparse_gfx950_gluon( if _wp: waves_per_eu = int(_wp) + # Cap the main segment's split count at whole BLOCK_K tiles: past that every + # split's main range is a partial tile, which costs a full masked gather for + # half the tokens. extra keeps all num_splits programs. + main_splits = num_splits + if has_extra and avg_main > 0: + main_splits = max(1, min(num_splits, int(math.ceil(avg_main / BLOCK_K)))) + _ms = _os.environ.get("AITER_PA_DECODE_SPLITS_MAIN") + if _ms: + main_splits = max(1, min(num_splits, int(_ms))) + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, @@ -736,6 +789,7 @@ def _pa_decode_sparse_gfx950_gluon( UNIFORM=UNIFORM, GATHER_W0=GATHER_W0, NOPE_CHUNK=nope_chunk, + MAIN_SPLITS=main_splits, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, HAS_INVALID=has_invalid, From 7034e3353c66c26baa0a87cf6f84994edcd83099 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 07:16:43 +0000 Subject: [PATCH 06/51] [GLUON][GFX950][DSV4] Whole-row gather, NaN-propagating max, base-2 softmax 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) --- .../gfx950/attention/pa_decode_sparse.py | 211 ++++++++++++------ .../ops/triton/attention/pa_decode_sparse.py | 26 ++- 2 files changed, 160 insertions(+), 77 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 74147c0b8a6..96317d53cf0 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -19,9 +19,27 @@ from triton.experimental import gluon from triton.experimental.gluon import language as gl +from triton.language.core import PropagateNan from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +# Triton's default max ignores NaN, which on AMD costs a v_max_f32 x, x, x +# canonicalize per operand before the real compare -- 60 of the 96 v_max in this +# kernel were those no-ops. Nothing here produces NaN (masked lanes are -inf and +# the all-masked row is guarded explicitly), so propagate instead. +_MAX_PROP_NAN: gl.constexpr = gl.constexpr(PropagateNan.ALL) + + +@gluon.jit +def _max2(a, b): + return gl.maximum(a, b, propagate_nan=_MAX_PROP_NAN) + + +@gluon.jit +def _rmax(x, axis): + return gl.reduce(x, axis, _max2) + + @gluon.jit def _cache_load(ptr, off, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): @@ -63,21 +81,50 @@ def _split2(x): @gluon.jit -def _deq_store(x_u8, sc, kv_smem, off, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr): +def _split2_dim0(x): + """Contiguous register split along dim 0: [A, B] -> two [A//2, B]. + + The dim-0 counterpart of _split2, for the row-wide gather layout where dim 1 + is spent entirely on threads (one instruction = whole token rows) and the + per-lane register repeats live on dim 0 instead. + """ + layout: gl.constexpr = x.type.layout + x_r = x.reshape(2, x.shape[0] // 2, x.shape[1]).permute(1, 2, 0) + x0, x1 = gl.split(x_r) + x0 = gl.convert_layout(x0, layout, assert_trivial=True) + x1 = gl.convert_layout(x1, layout, assert_trivial=True) + return x0, x1 + + +@gluon.jit +def _split_ax(x, AXIS: gl.constexpr): + if AXIS == 1: + a, b = _split2(x) + else: + a, b = _split2_dim0(x) + return a, b + + +@gluon.jit +def _deq_store( + x_u8, sc, kv_smem, off, + AXIS: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, +): """Dequant one fp8 slab and write it straight to kv_smem[:, off:off+W]. Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path lowers to an emulated mul and doubles the loop. ``sc`` is the raw per-64-group exponent byte (packed fp8_ds_mla, UE8M0) or an f32 scale (uniform pool). """ - W: gl.constexpr = x_u8.shape[1] if UNIFORM: scale = sc else: scale = gl.exp2(sc.to(gl.float32) - 127.0) - kv_smem.slice(off, W, dim=1).store( - (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) - ) + val = (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) + if AXIS == 1: + kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) + else: + kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) @gluon.jit @@ -86,56 +133,58 @@ def _deq_store_tile( sc, kv_smem, NOPE_CHUNK: gl.constexpr, + AXIS: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, ): - """Dequant a whole gathered fp8 tile into kv_smem in NOPE_CHUNK-wide pieces. + """Dequant a whole gathered fp8 tile into kv_smem in NOPE_CHUNK-sized pieces + along CHUNK_AXIS (0 = rows, 1 = columns). The f32 expansion is 4x the fp8 tile (a [64, 512] tile is 128 f32 VGPRs/lane), so materializing it in one shot is what pins 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 -- so only HEAD_SIZE/NOPE_CHUNK-th of it is ever + piece c's ds_writes and die -- so only a 1/pieces fraction of it is ever live. The splits themselves are register renames (see ``_split2``). """ - W: gl.constexpr = x_u8.shape[1] + W: gl.constexpr = x_u8.shape[1] if AXIS == 1 else x_u8.shape[0] if NOPE_CHUNK >= W: - _deq_store(x_u8, sc, kv_smem, 0, UNIFORM, FP8_FNUZ) + _deq_store(x_u8, sc, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) else: - x0, x1 = _split2(x_u8) - s0, s1 = _split2(sc) + x0, x1 = _split_ax(x_u8, AXIS) + s0, s1 = _split_ax(sc, AXIS) W2: gl.constexpr = W // 2 if NOPE_CHUNK >= W2: - _deq_store(x0, s0, kv_smem, 0, UNIFORM, FP8_FNUZ) - _deq_store(x1, s1, kv_smem, W2, UNIFORM, FP8_FNUZ) + _deq_store(x0, s0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x1, s1, kv_smem, W2, AXIS, UNIFORM, FP8_FNUZ) else: - x00, x01 = _split2(x0) - s00, s01 = _split2(s0) - x10, x11 = _split2(x1) - s10, s11 = _split2(s1) + x00, x01 = _split_ax(x0, AXIS) + s00, s01 = _split_ax(s0, AXIS) + x10, x11 = _split_ax(x1, AXIS) + s10, s11 = _split_ax(s1, AXIS) W4: gl.constexpr = W // 4 if NOPE_CHUNK >= W4: - _deq_store(x00, s00, kv_smem, 0, UNIFORM, FP8_FNUZ) - _deq_store(x01, s01, kv_smem, W4, UNIFORM, FP8_FNUZ) - _deq_store(x10, s10, kv_smem, 2 * W4, UNIFORM, FP8_FNUZ) - _deq_store(x11, s11, kv_smem, 3 * W4, UNIFORM, FP8_FNUZ) + _deq_store(x00, s00, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x01, s01, kv_smem, W4, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x10, s10, kv_smem, 2 * W4, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x11, s11, kv_smem, 3 * W4, AXIS, UNIFORM, FP8_FNUZ) else: W8: gl.constexpr = W // 8 - y0, y1 = _split2(x00) - t0, t1 = _split2(s00) - _deq_store(y0, t0, kv_smem, 0, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, W8, UNIFORM, FP8_FNUZ) - y0, y1 = _split2(x01) - t0, t1 = _split2(s01) - _deq_store(y0, t0, kv_smem, 2 * W8, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 3 * W8, UNIFORM, FP8_FNUZ) - y0, y1 = _split2(x10) - t0, t1 = _split2(s10) - _deq_store(y0, t0, kv_smem, 4 * W8, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 5 * W8, UNIFORM, FP8_FNUZ) - y0, y1 = _split2(x11) - t0, t1 = _split2(s11) - _deq_store(y0, t0, kv_smem, 6 * W8, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 7 * W8, UNIFORM, FP8_FNUZ) + y0, y1 = _split_ax(x00, AXIS) + t0, t1 = _split_ax(s00, AXIS) + _deq_store(y0, t0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, W8, AXIS, UNIFORM, FP8_FNUZ) + y0, y1 = _split_ax(x01, AXIS) + t0, t1 = _split_ax(s01, AXIS) + _deq_store(y0, t0, kv_smem, 2 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 3 * W8, AXIS, UNIFORM, FP8_FNUZ) + y0, y1 = _split_ax(x10, AXIS) + t0, t1 = _split_ax(s10, AXIS) + _deq_store(y0, t0, kv_smem, 4 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 5 * W8, AXIS, UNIFORM, FP8_FNUZ) + y0, y1 = _split_ax(x11, AXIS) + t0, t1 = _split_ax(s11, AXIS) + _deq_store(y0, t0, kv_smem, 6 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y1, t1, kv_smem, 7 * W8, AXIS, UNIFORM, FP8_FNUZ) @gluon.jit @@ -201,6 +250,7 @@ def _decode_tile( BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, + CHUNK_AXIS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, MASKED: gl.constexpr, UNIFORM: gl.constexpr, @@ -249,7 +299,7 @@ def _decode_tile( else: x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, True, FP8_FNUZ) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, True, FP8_FNUZ) elif IS_FP8: # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. nope_off = (block_idx_g * cs0 + pos_g * 576)[:, None] + offs_full[None, :] @@ -266,7 +316,7 @@ def _decode_tile( else: x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, False, FP8_FNUZ) + _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, False, FP8_FNUZ) block_idx_gr, pos_gr, valid_gr = _slots( indices_ptr, seg_start, @@ -323,12 +373,18 @@ def _decode_tile( ] S = gl.where(col_mask, S, neg_inf) - m_block = gl.max(S, axis=1) - m_new = gl.maximum(m_i, m_block) + # Online softmax in the base-2 exponent domain: fold qk_scale*log2(e) into S + # once, right out of the MFMA, and carry m_i already scaled. The alternative + # (raw running max, scale at use) needs qk_scale on both the S tile and the two + # m vectors every iteration; here the per-element work in the loop is one + # subtract feeding exp2, and the whole epilogue -- sink combine and partial + # store, both of which want base-2 m -- stops converting. + S = S * qk_scale + m_block = _rmax(S, 1) + m_new = _max2(m_i, m_block) m_new = gl.where(m_new > neg_inf, m_new, 0.0) # guard all-masked rows - m_new_s = m_new * qk_scale - p = gl.exp2(S * qk_scale - m_new_s[:, None]) - alpha = gl.exp2(m_i * qk_scale - m_new_s) + p = gl.exp2(S - m_new[:, None]) + alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(v_layout) # [BLOCK_K, HEAD_SIZE] @@ -433,6 +489,7 @@ def _qkpv_fp8( BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, + CHUNK_AXIS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, @@ -442,7 +499,7 @@ def _qkpv_fp8( skips the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, mask the columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" neg_inf = float("-inf") - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, UNIFORM, FP8_FNUZ) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, UNIFORM, FP8_FNUZ) if not UNIFORM: kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) k = kv_smem.permute([1, 0]).load(k_layout) @@ -463,12 +520,12 @@ def _qkpv_fp8( :, None ] S = gl.where(col_mask, S, neg_inf) - m_block = gl.max(S, axis=1) - m_new = gl.maximum(m_i, m_block) + S = S * qk_scale # see _decode_tile: m_i is carried in base-2 exponent space + m_block = _rmax(S, 1) + m_new = _max2(m_i, m_block) m_new = gl.where(m_new > neg_inf, m_new, 0.0) - m_new_s = m_new * qk_scale - p = gl.exp2(S * qk_scale - m_new_s[:, None]) - alpha = gl.exp2(m_i * qk_scale - m_new_s) + p = gl.exp2(S - m_new[:, None]) + alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(v_layout) p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) @@ -511,6 +568,7 @@ def _process_segment( BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, + CHUNK_AXIS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, @@ -593,6 +651,7 @@ def _process_segment( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, UNIFORM, FP8_FNUZ, @@ -622,6 +681,7 @@ def _process_segment( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, UNIFORM, FP8_FNUZ, @@ -664,6 +724,7 @@ def _process_segment( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, False, UNIFORM, @@ -708,6 +769,7 @@ def _process_segment( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, True, UNIFORM, @@ -771,12 +833,11 @@ def _pa_decode_sparse( HEAD_ALIGNED: gl.constexpr, MFMA_K: gl.constexpr, UNIFORM: gl.constexpr, - # Experiment knobs (driver reads AITER_PA_DECODE_* env vars; defaults = prior - # behaviour). GATHER_W0: put the NoPE gather's warps on dim 0 (rows) instead of - # dim 1 (columns). - GATHER_W0: gl.constexpr, - # NOPE_CHUNK: width of one dequant piece (HEAD_SIZE = one shot, prior behaviour). + GATHER_TW1: gl.constexpr, + # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, 1 = + # columns); >= the tile's extent on that axis means one shot. NOPE_CHUNK: gl.constexpr, + CHUNK_AXIS: gl.constexpr, # How many of the NUM_SPLITS programs share the main (SWA) segment. The two # segments have very different shapes -- main is a contiguous window whose # length is fixed by the sliding window, extra is the top-k list -- so one @@ -842,10 +903,18 @@ def _pa_decode_sparse( # warps-on-dim-0 carries BLOCK_K/(8*NUM_WARPS). Dim 0 also makes the column # direction a pure per-lane register repeat, so a chunked dequant can split it # with a trivial (register-renaming) layout convert. + # GATHER_TW1 = threads spent on the head dim. Each thread loads 16 B, so + # TW1=8 requests 128 B of a token row per instruction and TW1=32 requests the + # whole 512 B row -- for a scattered gather that is one request instead of + # four quarters. The cost is that the row-index vector lives in + # SliceLayout(1, gather_l), so a wider TW1 leaves fewer dim-0 thread slots and + # every lane carries more slots (BLOCK_K*TW1/64 of them). Whichever dim keeps + # per-lane register repeats is the one the chunked dequant can split + # (CHUNK_AXIS). gather_l: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, GSPT], - threads_per_warp=[8, 8], - warps_per_cta=[NUM_WARPS, 1] if GATHER_W0 else [1, NUM_WARPS], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], order=[1, 0], ) # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting @@ -941,6 +1010,7 @@ def _pa_decode_sparse( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, UNIFORM, MAIN_USE_BUFFER_LOAD, @@ -987,6 +1057,7 @@ def _pa_decode_sparse( BLOCK_M, BLOCK_K, NOPE_CHUNK, + CHUNK_AXIS, HEAD_ALIGNED, UNIFORM, EXTRA_USE_BUFFER_LOAD, @@ -1001,15 +1072,15 @@ def _pa_decode_sparse( if NUM_SPLITS == 1: if HAS_SINK: - # m_pv is the RAW row-max; the sink is a scaled-score logit. Combine - # in scaled space (Ms = scale*m) using exp2. - Ms = m_pv * scale + # m_pv is already the row-max in the base-2 exponent domain + # (row_max * softmax_scale * log2e); lift the sink -- a scaled-score + # logit -- into the same domain and combine there. sink = gl.amd.cdna4.buffer_load( ptr=attn_sink_ptr, offsets=h_pv, mask=head_mask_pv, other=float("-inf") - ).to(gl.float32) - m_final = gl.maximum(Ms, sink) - alpha = gl.exp2((Ms - m_final) * RCP_LN2) - l_final = l_pv * alpha + gl.exp2((sink - m_final) * RCP_LN2) + ).to(gl.float32) * RCP_LN2 + m_final = _max2(m_pv, sink) + alpha = gl.exp2(m_pv - m_final) + l_final = l_pv * alpha + gl.exp2(sink - m_final) acc = acc * alpha[:, None] else: l_final = l_pv @@ -1025,12 +1096,12 @@ def _pa_decode_sparse( mask=head_mask_pv[:, None], ) else: - # store un-normalized partials for the reduce kernel. m is stored in the - # base-2 exponent domain (row-max * softmax_scale * log2e) so the reduce/ - # skip_reduce partials match the triton convention. + # store un-normalized partials for the reduce kernel. m is already in the + # base-2 exponent domain (row-max * softmax_scale * log2e), which is the + # reduce/skip_reduce convention the triton kernel also uses. pm_base = query_idx * pm_stride0 + split_id * pm_stride_s gl.amd.cdna4.buffer_store( - m_pv * (scale * RCP_LN2), + m_pv, ptr=part_m_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv, @@ -1115,12 +1186,12 @@ def _pa_decode_sparse_reduce( m_s = gl.amd.cdna4.buffer_load( ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf ) - m_final = gl.maximum(m_final, m_s) # m_s already in base-2 exponent domain + m_final = _max2(m_final, m_s) # m_s already in base-2 exponent domain if HAS_SINK: sink = gl.amd.cdna4.buffer_load( ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf ).to(gl.float32) - m_final = gl.maximum(m_final, sink * RCP_LN2) # lift sink to base-2 + m_final = _max2(m_final, sink * RCP_LN2) # lift sink to base-2 # pass 2: weighted sums l_final = gl.zeros([BLOCK_M], gl.float32, layout=row_l) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index c9d02d55df3..705c3fbe441 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -579,8 +579,12 @@ def _pa_decode_sparse_gfx950_gluon( if _mk: MFMA_K = int(_mk) num_warps = max(1, BLOCK_K // 16) - # AITER_PA_DECODE_GATHER_W0: NoPE gather warps on dim 0 (rows) instead of dim 1. - GATHER_W0 = _os.environ.get("AITER_PA_DECODE_GATHER_W0", "1") == "1" + # AITER_PA_DECODE_TW1: threads the NoPE gather spends on the head dim (16 B + # each). 32 issues a whole 512 B token row per instruction instead of 8 x 128 B + # quarters, which is what a scattered top-k gather wants; the cost is a longer + # per-lane slot vector. Measured best at 32 for every extra >= 128 (-10% at + # extra=128) and within noise at extra=8. + gather_tw1 = int(_os.environ.get("AITER_PA_DECODE_TW1", "32")) # AITER_PA_DECODE_NOPE_CHUNK: width of one fp8 dequant piece (0/unset = one shot). NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -713,10 +717,17 @@ def _pa_decode_sparse_gfx950_gluon( # 128 is the narrowest piece the gather layout can split to for free (below # that a split falls inside a thread's 16-byte run, which Triton rejects), and # the split is only trivial when the gather's warps tile dim 0. + # Row-wide gather chunks rows (4 pieces of BLOCK_K/4) instead of 128-wide + # column pieces; both leave a quarter of the tile's f32 expansion live. + # Split the axis that still has per-lane register repeats: columns when the + # gather leaves >=2 column repeats (TW1=8), rows otherwise. + col_reps = head_dim // (gather_tw1 * 16) + chunk_axis = 1 if col_reps >= 4 else 0 _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") - nope_chunk = int(_nc) if _nc else min(128, head_dim) - if not GATHER_W0: - nope_chunk = head_dim + if chunk_axis == 0: + nope_chunk = int(_nc) if _nc else max(1, BLOCK_K // 4) + else: + nope_chunk = int(_nc) if _nc else min(128, head_dim) # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD # threshold on gfx950's 512-VGPR file. The chunked dequant gets the buffer_load @@ -724,7 +735,7 @@ def _pa_decode_sparse_gfx950_gluon( # offset gathers through 64-bit addresses and lands at ~321 -- there the cap # buys nothing and costs ~150 scratch stores (+40% at C4A), so only ask for it # when both caches are on the fast path. - if use_buffer_load and nope_chunk < head_dim: + if use_buffer_load and (chunk_axis == 0 or nope_chunk < head_dim): waves_per_eu = 2 _wp = _os.environ.get("AITER_PA_DECODE_WPEU") if _wp: @@ -787,8 +798,9 @@ def _pa_decode_sparse_gfx950_gluon( HEAD_ALIGNED=HEAD_ALIGNED, MFMA_K=MFMA_K, UNIFORM=UNIFORM, - GATHER_W0=GATHER_W0, + GATHER_TW1=gather_tw1, NOPE_CHUNK=nope_chunk, + CHUNK_AXIS=chunk_axis, MAIN_SPLITS=main_splits, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, From e3b8129cf6764a90e9af44c3aa959e13641f3e18 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 07:28:53 +0000 Subject: [PATCH 07/51] [GLUON][GFX950][DSV4] Split policy: do not inflate partial tiles when 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 c6f04892e: 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) --- aiter/ops/triton/attention/pa_decode_sparse.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 705c3fbe441..07aaa6fb136 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -533,7 +533,19 @@ def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block # split still owns a real run of tiles. At C=256 extra=8 (2 tiles) forcing # a split costs 31%; at extra=2048 (32 tiles) it saves 19%. return max(1, min(cta_cap, tiles // 4)) - return max(1, min(cta_cap, tiles)) + if tiles <= cta_cap: + return max(1, tiles) + # The cap binds, so every split owns several tiles and the only question is + # whether the last one is partial. S * ceil(tiles/S) is the number of tiles + # actually executed: at 18 tiles, S=8 runs 24 of them and S=6 runs 18. Take the + # largest S in range that executes the fewest -- largest because, ties aside, + # more programs is more memory parallelism. + best_s, best_cost = 1, tiles + for s in range(2, cta_cap + 1): + cost = s * math.ceil(tiles / s) + if cost <= best_cost: + best_s, best_cost = s, cost + return best_s def _pa_decode_sparse_gfx950_gluon( From 42ebba54ffc43c1efaeafa790dd3e033c83b530d Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 07:50:26 +0000 Subject: [PATCH 08/51] [GLUON][GFX950][DSV4] Per-query adaptive split count with early CTA exit 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) --- .../gfx950/attention/pa_decode_sparse.py | 74 +++++++++++++++++-- .../ops/triton/attention/pa_decode_sparse.py | 28 +++---- 2 files changed, 82 insertions(+), 20 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 96317d53cf0..403e9292b5b 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -847,6 +847,8 @@ def _pa_decode_sparse( # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an # empty main range and contribute an extra-only partial. MAIN_SPLITS: gl.constexpr, + # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. + ADAPTIVE_SPLITS: gl.constexpr, # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache # whose span exceeds that must gather via 64-bit gl.load -- but the two caches # are sized independently (SWA window vs full compressed history), so gating @@ -971,11 +973,59 @@ def _pa_decode_sparse( RCP_LN2: gl.constexpr = 1.4426950408889634 qk_scale = scale * RCP_LN2 - # ---- main (SWA) segment ---- + # ---- segment lengths ---- main_start = gl.load(main_indptr_ptr + query_idx) main_end = gl.load(main_indptr_ptr + query_idx + 1) main_len = main_end - main_start - main_chunk = (main_len + MAIN_SPLITS - 1) // MAIN_SPLITS + if HAS_EXTRA: + extra_start = gl.load(extra_indptr_ptr + query_idx) + extra_end = gl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + else: + extra_start = 0 + extra_len = 0 + + # ---- how many of the launched splits this query actually wants ---- + # NUM_SPLITS / MAIN_SPLITS come from the host, which only knows the batch's + # AVERAGE segment lengths. In a ragged batch (the SWA window saturates at + # sliding_window while top-k keeps growing with context) the per-query lengths + # differ a lot, and a split count sized for the average over-splits the short + # queries -- each surplus program still gathers a mostly-masked tile and writes + # a full [BLOCK_M, HEAD_SIZE] f32 partial. Recompute the useful count from this + # query's own lengths and let the surplus programs write a neutral partial and + # leave. The reduce skips a split whose m is -inf, so their part_acc is never + # read and does not have to be written. + if ADAPTIVE_SPLITS: + m_tiles = (main_len + BLOCK_K - 1) // BLOCK_K + e_tiles = (extra_len + BLOCK_K - 1) // BLOCK_K + work_splits = gl.minimum(gl.maximum(gl.maximum(m_tiles, e_tiles), 1), NUM_SPLITS) + main_splits = gl.minimum(gl.maximum(m_tiles, 1), work_splits) + if split_id >= work_splits: + pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + gl.amd.cdna4.buffer_store( + gl.full( + [BLOCK_M], + float("-inf"), + gl.float32, + layout=gl.SliceLayout(1, pv_layout), + ), + ptr=part_m_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + gl.amd.cdna4.buffer_store( + gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, pv_layout)), + ptr=part_l_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) + return + else: + work_splits = NUM_SPLITS + main_splits = MAIN_SPLITS + + # ---- main (SWA) segment ---- + main_chunk = (main_len + main_splits - 1) // main_splits main_lo = gl.minimum(split_id * main_chunk, main_len) main_hi = gl.minimum(main_lo + main_chunk, main_len) m_i, l_i, acc = _process_segment( @@ -1019,10 +1069,7 @@ def _pa_decode_sparse( ) if HAS_EXTRA: - extra_start = gl.load(extra_indptr_ptr + query_idx) - extra_end = gl.load(extra_indptr_ptr + query_idx + 1) - extra_len = extra_end - extra_start - extra_chunk = (extra_len + NUM_SPLITS - 1) // NUM_SPLITS + extra_chunk = (extra_len + work_splits - 1) // work_splits extra_lo = split_id * extra_chunk extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) m_i, l_i, acc = _process_segment( @@ -1143,6 +1190,7 @@ def _pa_decode_sparse_reduce( BLOCK_M: gl.constexpr, NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, + ADAPTIVE_SPLITS: gl.constexpr, ): """Split-KV combine: merge the per-split partials, fold the attn sink, and write the final output. Partials store m in the base-2 exponent domain (row-max * @@ -1179,6 +1227,11 @@ def _pa_decode_sparse_reduce( offs_d = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, BLK)) neg_inf = float("-inf") + # NOT made adaptive. Bounding these loops by the per-query split count instead + # of NUM_SPLITS 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 the + # launch over-splits. The acc load is masked on m > -inf instead, which is + # what actually has to be right (a split that bowed out wrote no accumulator). m_final = gl.full([BLOCK_M], neg_inf, gl.float32, layout=row_l) # pass 1: global max over splits for s in range(NUM_SPLITS): @@ -1208,8 +1261,15 @@ def _pa_decode_sparse_reduce( l_final = l_final + w * l_s a_base = query_idx * pa_stride0 + s * pa_stride_s a_off = (a_base + h[:, None] * pa_stride_h + offs_d[None, :]).to(gl.int32) + # An adaptive-split program that bowed out wrote m = -inf and no + # accumulator, so its part_acc slot is uninitialized -- mask the load + # rather than relying on w == 0 (0 * NaN is NaN). + if ADAPTIVE_SPLITS: + acc_mask = head_mask[:, None] & (m_s > neg_inf)[:, None] + else: + acc_mask = head_mask[:, None] acc_s = gl.amd.cdna4.buffer_load( - ptr=part_acc_ptr, offsets=a_off, mask=head_mask[:, None], other=0.0 + ptr=part_acc_ptr, offsets=a_off, mask=acc_mask, other=0.0 ) acc = acc + w[:, None] * acc_s diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 07aaa6fb136..64d2cf201d3 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -533,19 +533,12 @@ def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block # split still owns a real run of tiles. At C=256 extra=8 (2 tiles) forcing # a split costs 31%; at extra=2048 (32 tiles) it saves 19%. return max(1, min(cta_cap, tiles // 4)) - if tiles <= cta_cap: - return max(1, tiles) - # The cap binds, so every split owns several tiles and the only question is - # whether the last one is partial. S * ceil(tiles/S) is the number of tiles - # actually executed: at 18 tiles, S=8 runs 24 of them and S=6 runs 18. Take the - # largest S in range that executes the fewest -- largest because, ties aside, - # more programs is more memory parallelism. - best_s, best_cost = 1, tiles - for s in range(2, cta_cap + 1): - cost = s * math.ceil(tiles / s) - if cost <= best_cost: - best_s, best_cost = s, cost - return best_s + # Deliberately NOT refined toward split counts that divide `tiles` evenly. + # Minimizing the executed tile count S*ceil(tiles/S) is only right for a + # uniform batch -- it is worth ~5% on the 1-loop seq=1152 shape -- but `tiles` + # here is a batch AVERAGE, and wall clock is set by the longest query. On a + # 16x-ragged batch that refinement picks 5 splits where 8 is 22% faster. + return max(1, min(cta_cap, tiles)) def _pa_decode_sparse_gfx950_gluon( @@ -763,6 +756,13 @@ def _pa_decode_sparse_gfx950_gluon( if _ms: main_splits = max(1, min(num_splits, int(_ms))) + # AITER_PA_DECODE_ADAPTIVE: per-query split count decided in-kernel. Only + # meaningful when there is more than one split to give back. + adaptive_splits = ( + num_splits > 1 + and _os.environ.get("AITER_PA_DECODE_ADAPTIVE", "1") == "1" + ) + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, @@ -814,6 +814,7 @@ def _pa_decode_sparse_gfx950_gluon( NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, MAIN_SPLITS=main_splits, + ADAPTIVE_SPLITS=adaptive_splits, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, HAS_INVALID=has_invalid, @@ -860,6 +861,7 @@ def _pa_decode_sparse_gfx950_gluon( BLOCK_M=red_block_m, NUM_SPLITS=num_splits, HEAD_ALIGNED=num_heads % red_block_m == 0, + ADAPTIVE_SPLITS=adaptive_splits, num_warps=red_warps, ) return out From 8c8e1c33db19dec08fe9352cd853ef0397ee8fd3 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 08:07:06 +0000 Subject: [PATCH 09/51] [GLUON][GFX950][DSV4] Fix the prologue's serial chain and the gather 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) --- .../gfx950/attention/pa_decode_sparse.py | 47 ++++++++++++------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 403e9292b5b..e187b040b53 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -306,16 +306,16 @@ def _decode_tile( scl_off = (block_idx_g * cs0 + BLOCK_SIZE * 576 + pos_g * 8)[:, None] + ( offs_full[None, :] // 64 ) - if MASKED: - x_u8 = _cache_load( - cache_ptr, nope_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 - ) + if MASKED: # scales first: see _gd_fp8 exps = _cache_load( cache_ptr, scl_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=127 ) + x_u8 = _cache_load( + cache_ptr, nope_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 + ) else: - x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) + x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, False, FP8_FNUZ) block_idx_gr, pos_gr, valid_gr = _slots( indices_ptr, @@ -448,8 +448,13 @@ def _gd_fp8( scl_off = (bg * cs0 + BLOCK_SIZE * 576 + pg * 8)[:, None] + ( offs_full[None, :] // 64 ) - x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) + # UE8M0 scales first, bulk fp8 after. vmcnt is one in-order FIFO, so a + # wait can only name "at most N outstanding", never a specific load: + # issuing the scales last would make the first dequant piece wait behind + # every data load as well. Issued first, the scales are covered by the + # wait that the first piece's own data needs anyway. sc = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) + x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) bgr, pgr, _ = _slots( indices_ptr, seg_start, @@ -941,6 +946,24 @@ def _pa_decode_sparse( h_off = pid_h * BLOCK_M + # ---- segment lengths, issued first ---- + # These gate the whole KV chain: indptr -> segment range -> indices gather -> + # cache addresses -> cache gather, three dependent memory round trips deep. + # Q is independent of all of it, so issue the indptr loads before Q rather + # than after: the scalar readfirstlane they feed needs s_waitcnt vmcnt(0), and + # behind Q's load plus its through-LDS layout conversion (two barriers) that + # wait lands at the end of a long serial prologue instead of overlapping it. + main_start = gl.load(main_indptr_ptr + query_idx) + main_end = gl.load(main_indptr_ptr + query_idx + 1) + if HAS_EXTRA: + extra_start = gl.load(extra_indptr_ptr + query_idx) + extra_end = gl.load(extra_indptr_ptr + query_idx + 1) + extra_len = extra_end - extra_start + else: + extra_start = 0 + extra_len = 0 + main_len = main_end - main_start + # ---- load Q [BLOCK_M, HEAD_SIZE] ---- offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, blocked_q)) offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, blocked_q)) @@ -973,18 +996,6 @@ def _pa_decode_sparse( RCP_LN2: gl.constexpr = 1.4426950408889634 qk_scale = scale * RCP_LN2 - # ---- segment lengths ---- - main_start = gl.load(main_indptr_ptr + query_idx) - main_end = gl.load(main_indptr_ptr + query_idx + 1) - main_len = main_end - main_start - if HAS_EXTRA: - extra_start = gl.load(extra_indptr_ptr + query_idx) - extra_end = gl.load(extra_indptr_ptr + query_idx + 1) - extra_len = extra_end - extra_start - else: - extra_start = 0 - extra_len = 0 - # ---- how many of the launched splits this query actually wants ---- # NUM_SPLITS / MAIN_SPLITS come from the host, which only knows the batch's # AVERAGE segment lengths. In a ragged batch (the SWA window saturates at From e9cb7b32651667de7783b229dd82af715ddd3a0f Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 09:25:41 +0000 Subject: [PATCH 10/51] [GLUON][GFX950][DSV4] Narrow the UE8M0 scale gather on the 64-bit cache 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> data: tensor<32x128x...> gluon scales: tensor<64x512x!tt.ptr> 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) --- .../gfx950/attention/pa_decode_sparse.py | 228 ++++++++++++++---- 1 file changed, 186 insertions(+), 42 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index e187b040b53..03b944e74e7 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -42,14 +42,33 @@ def _rmax(x, axis): @gluon.jit -def _cache_load(ptr, off, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): - # gfx950 buffer_load carries a 32-bit offset (2 GB cap), a cache past that gathers - # via 64-bit gl.load instead. +def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): + """Gather rows[i] + col[j] out of a cache. + + ``row`` is the per-token byte offset ([BLOCK_K]); ``col`` is the in-row offset + ([W]), which is always a small compile-time arange. Keeping them apart is what + makes the >2 GB path affordable: buffer_load carries a 32-bit offset (2 GB cap), + so a bigger cache has to gather through 64-bit addresses -- and adding the + pointer to a fully materialized [BLOCK_K, W] offset tensor makes that a 64-bit + add *per element*. Resolving one pointer per token instead costs BLOCK_K of + them, and the leftover column offset is a constant the load can fold into its + immediate field. (Same shape as the in-tree triton kernel's token_data_ptr.) + """ if USE_BUFFER_LOAD: return gl.amd.cdna4.buffer_load( - ptr=ptr, offsets=off.to(gl.int32), mask=mask, other=other, cache=".cg" + ptr=ptr, + offsets=row.to(gl.int32)[:, None] + col.to(gl.int32)[None, :], + mask=mask, + other=other, + cache=".cg", ) - return gl.load(ptr + off.to(gl.int64), mask=mask, other=other, cache_modifier=".cg") + row_ptr = ptr + row.to(gl.int64) + return gl.load( + row_ptr[:, None] + col[None, :], + mask=mask, + other=other, + cache_modifier=".cg", + ) @gluon.jit @@ -62,6 +81,68 @@ def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) +@gluon.jit +def _scale_load( + ptr, + row, + valid, + USE_BUFFER_LOAD: gl.constexpr, + gather_l: gl.constexpr, + scl_l: gl.constexpr, + NG: gl.constexpr, + HEAD_SIZE: gl.constexpr, + MASKED: gl.constexpr, +): + """Gather the NG per-64-group scale bytes of each token and broadcast them out + to HEAD_SIZE in registers. + + The obvious spelling -- index the full row with ``offs_full // 64`` -- builds a + [BLOCK_K, HEAD_SIZE] pointer tensor for a value that has only NG distinct + entries per row. On the buffer_load path the identical 32-bit offsets CSE away + and it costs 48 byte loads; on the 64-bit path they do not, and it costs 288. + (The in-tree triton kernel loads tensor<32x2x!tt.ptr> here, we were loading + tensor<64x512x!tt.ptr>.) Gather NG wide instead and broadcast: ``scl_l`` is + the dim-2 slice of a 3-D layout picked so that reshaping [BLOCK_K, NG, 64] back + to [BLOCK_K, HEAD_SIZE] lands exactly on ``gather_l`` -- column c = g*64 + j + maps to thread 4g + j//16 -- so the broadcast is a register rename. + """ + cols = gl.arange(0, NG, layout=gl.SliceLayout(0, scl_l)) + rows = gl.convert_layout(row, gl.SliceLayout(1, scl_l)) + if MASKED: + m = gl.convert_layout(valid, gl.SliceLayout(1, scl_l))[:, None] + if USE_BUFFER_LOAD: + sc = gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=rows.to(gl.int32)[:, None] + cols[None, :], + mask=m, + other=127, + cache=".cg", + ) + else: + sc = gl.load( + (ptr + rows.to(gl.int64))[:, None] + cols[None, :], + mask=m, + other=127, + cache_modifier=".cg", + ) + else: + if USE_BUFFER_LOAD: + sc = gl.amd.cdna4.buffer_load( + ptr=ptr, + offsets=rows.to(gl.int32)[:, None] + cols[None, :], + cache=".cg", + ) + else: + sc = gl.load( + (ptr + rows.to(gl.int64))[:, None] + cols[None, :], + cache_modifier=".cg", + ) + wide = gl.expand_dims(sc, 2).broadcast_to([sc.shape[0], NG, HEAD_SIZE // NG]) + return gl.convert_layout( + wide.reshape([sc.shape[0], HEAD_SIZE]), gather_l, assert_trivial=True + ) + + @gluon.jit def _split2(x): """Contiguous register split along dim 1: [A, B] -> two [A, B//2]. @@ -242,6 +323,8 @@ def _decode_tile( p_layout: gl.constexpr, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + scl_l: gl.constexpr, + NARROW_SCALE: gl.constexpr, IS_FP8: gl.constexpr, BLOCK_SIZE: gl.constexpr, NOPE_DIM: gl.constexpr, @@ -283,39 +366,55 @@ def _decode_tile( if IS_FP8 and UNIFORM: # uniform pool: one fp8 gather over the whole head + separate fp32 scales. NGRP: gl.constexpr = HEAD_SIZE // 64 - kv_off = (block_idx_g * cs0 + pos_g * HEAD_SIZE)[:, None] + offs_full[None, :] - scl_off = (block_idx_g * NGRP)[:, None] + (offs_full[None, :] // 64) + kv_row = block_idx_g * cs0 + pos_g * HEAD_SIZE + scl_row = block_idx_g * NGRP + scl_col = offs_full // 64 if MASKED: x_u8 = _cache_load( - cache_ptr, kv_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 + cache_ptr, kv_row, offs_full, USE_BUFFER_LOAD, + mask=valid_g[:, None], other=0, ) sc = _cache_load( cache_bf16_ptr, - scl_off, + scl_row, + scl_col, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, - ) + ) # uniform pool: fp32 scales, left wide (see _scale_load) else: - x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) - sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) + x_u8 = _cache_load(cache_ptr, kv_row, offs_full, USE_BUFFER_LOAD) + sc = _cache_load(cache_bf16_ptr, scl_row, scl_col, USE_BUFFER_LOAD) _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, True, FP8_FNUZ) elif IS_FP8: # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. - nope_off = (block_idx_g * cs0 + pos_g * 576)[:, None] + offs_full[None, :] - scl_off = (block_idx_g * cs0 + BLOCK_SIZE * 576 + pos_g * 8)[:, None] + ( - offs_full[None, :] // 64 - ) + nope_row = block_idx_g * cs0 + pos_g * 576 + scl_row = block_idx_g * cs0 + BLOCK_SIZE * 576 + pos_g * 8 + scl_col = offs_full // 64 if MASKED: # scales first: see _gd_fp8 - exps = _cache_load( - cache_ptr, scl_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=127 - ) + if NARROW_SCALE and not USE_BUFFER_LOAD: + exps = _scale_load( + cache_ptr, scl_row, valid1d, USE_BUFFER_LOAD, gather_l, scl_l, + HEAD_SIZE // 64, HEAD_SIZE, True, + ) + else: + exps = _cache_load( + cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD, + mask=valid_g[:, None], other=127, + ) x_u8 = _cache_load( - cache_ptr, nope_off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0 + cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD, + mask=valid_g[:, None], other=0, ) else: - exps = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) - x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) + if NARROW_SCALE and not USE_BUFFER_LOAD: + exps = _scale_load( + cache_ptr, scl_row, scl_row, USE_BUFFER_LOAD, gather_l, scl_l, + HEAD_SIZE // 64, HEAD_SIZE, False, + ) + else: + exps = _cache_load(cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD) + x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, False, FP8_FNUZ) block_idx_gr, pos_gr, valid_gr = _slots( indices_ptr, @@ -327,28 +426,28 @@ def _decode_tile( MASKED, HAS_INVALID, ) - rope_off = (block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224)[ - :, None - ] + offs_rope[None, :] + rope_row = block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224 if MASKED: k_rope = _cache_load( cache_bf16_ptr, - rope_off, + rope_row, + offs_rope, USE_BUFFER_LOAD, mask=valid_gr[:, None], other=0.0, ) else: - k_rope = _cache_load(cache_bf16_ptr, rope_off, USE_BUFFER_LOAD) + k_rope = _cache_load(cache_bf16_ptr, rope_row, offs_rope, USE_BUFFER_LOAD) kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) else: - off = (block_idx_g * cs0 + pos_g * HEAD_SIZE)[:, None] + offs_full[None, :] + kv_row2 = block_idx_g * cs0 + pos_g * HEAD_SIZE if MASKED: kv = _cache_load( - cache_bf16_ptr, off, USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0 + cache_bf16_ptr, kv_row2, offs_full, USE_BUFFER_LOAD, + mask=valid_g[:, None], other=0.0, ) else: - kv = _cache_load(cache_bf16_ptr, off, USE_BUFFER_LOAD) + kv = _cache_load(cache_bf16_ptr, kv_row2, offs_full, USE_BUFFER_LOAD) kv_smem.store(kv) k = kv_smem.permute([1, 0]).load(k_layout) # [HEAD_SIZE, BLOCK_K] @@ -409,6 +508,8 @@ def _gd_fp8( k_rng_rope, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + scl_l: gl.constexpr, + NARROW_SCALE: gl.constexpr, BLOCK_SIZE: gl.constexpr, HEAD_SIZE: gl.constexpr, UNIFORM: gl.constexpr, @@ -438,23 +539,27 @@ def _gd_fp8( ) if UNIFORM: NGRP: gl.constexpr = HEAD_SIZE // 64 - kv_off = (bg * cs0 + pg * HEAD_SIZE)[:, None] + offs_full[None, :] - scl_off = (bg * NGRP)[:, None] + (offs_full[None, :] // 64) - x_u8 = _cache_load(cache_ptr, kv_off, USE_BUFFER_LOAD) - sc = _cache_load(cache_bf16_ptr, scl_off, USE_BUFFER_LOAD) + x_u8 = _cache_load( + cache_ptr, bg * cs0 + pg * HEAD_SIZE, offs_full, USE_BUFFER_LOAD + ) + sc = _cache_load(cache_bf16_ptr, bg * NGRP, offs_full // 64, USE_BUFFER_LOAD) k_rope = x_u8 # unused for UNIFORM (rope slice-store skipped) -> DCE'd else: - nope_off = (bg * cs0 + pg * 576)[:, None] + offs_full[None, :] - scl_off = (bg * cs0 + BLOCK_SIZE * 576 + pg * 8)[:, None] + ( - offs_full[None, :] // 64 - ) + nope_row = bg * cs0 + pg * 576 + scl_row = bg * cs0 + BLOCK_SIZE * 576 + pg * 8 # UE8M0 scales first, bulk fp8 after. vmcnt is one in-order FIFO, so a # wait can only name "at most N outstanding", never a specific load: # issuing the scales last would make the first dequant piece wait behind # every data load as well. Issued first, the scales are covered by the # wait that the first piece's own data needs anyway. - sc = _cache_load(cache_ptr, scl_off, USE_BUFFER_LOAD) - x_u8 = _cache_load(cache_ptr, nope_off, USE_BUFFER_LOAD) + if NARROW_SCALE and not USE_BUFFER_LOAD: + sc = _scale_load( + cache_ptr, scl_row, scl_row, USE_BUFFER_LOAD, gather_l, scl_l, + HEAD_SIZE // 64, HEAD_SIZE, False, + ) + else: + sc = _cache_load(cache_ptr, scl_row, offs_full // 64, USE_BUFFER_LOAD) + x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) bgr, pgr, _ = _slots( indices_ptr, seg_start, @@ -465,8 +570,9 @@ def _gd_fp8( False, HAS_INVALID, ) - rope_off = (bgr * (cs0 // 2) + pgr * 288 + 224)[:, None] + offs_rope[None, :] - k_rope = _cache_load(cache_bf16_ptr, rope_off, USE_BUFFER_LOAD) + k_rope = _cache_load( + cache_bf16_ptr, bgr * (cs0 // 2) + pgr * 288 + 224, offs_rope, USE_BUFFER_LOAD + ) return x_u8, sc, k_rope, valid @@ -564,6 +670,8 @@ def _process_segment( p_layout: gl.constexpr, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + scl_l: gl.constexpr, + NARROW_SCALE: gl.constexpr, slot_l: gl.constexpr, IS_FP8: gl.constexpr, BLOCK_SIZE: gl.constexpr, @@ -605,6 +713,8 @@ def _process_segment( k_rng_rope, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, BLOCK_SIZE, HEAD_SIZE, UNIFORM, @@ -626,6 +736,8 @@ def _process_segment( k_rng_rope, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, BLOCK_SIZE, HEAD_SIZE, UNIFORM, @@ -721,6 +833,8 @@ def _process_segment( p_layout, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, IS_FP8, BLOCK_SIZE, NOPE_DIM, @@ -766,6 +880,8 @@ def _process_segment( p_layout, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, IS_FP8, BLOCK_SIZE, NOPE_DIM, @@ -933,6 +1049,30 @@ def _pa_decode_sparse( warps_per_cta=[NUM_WARPS, 1], order=[1, 0], ) + # 3-D companion of gather_l used by _scale_load: dim 1 carries the NG scale + # groups (one thread each) and dim 2 the 64 columns inside a group, so that + # reshaping [BLOCK_K, NG, 64] back to 2-D reproduces gather_l exactly. Legal + # only when the group's columns fill whole threads, i.e. GSPT * (TW1 // NG) is + # the group width; otherwise fall back to the wide (redundant) gather. + NGRP_S: gl.constexpr = HEAD_SIZE // 64 + NARROW_SCALE: gl.constexpr = ( + GATHER_TW1 % NGRP_S == 0 and GSPT * (GATHER_TW1 // NGRP_S) == 64 + ) + # ...and only worth it on the 64-bit path. With buffer_load the wide gather's + # identical 32-bit offsets CSE to the same 48 byte loads, so the narrow form + # only adds a layout convert per tile (+3.4% at extra=1024). Without it they do + # not CSE and the wide form costs 288 loads. + scl_l3: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 1, GSPT], + threads_per_warp=[ + 64 // GATHER_TW1, + NGRP_S if NARROW_SCALE else 1, + (GATHER_TW1 // NGRP_S) if NARROW_SCALE else GATHER_TW1, + ], + warps_per_cta=[NUM_WARPS, 1, 1], + order=[2, 1, 0], + ) + scl_l: gl.constexpr = gl.SliceLayout(2, scl_l3) slot_l: gl.constexpr = gl.SliceLayout(1, gather_l) blocked_q: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, 8], @@ -1062,6 +1202,8 @@ def _pa_decode_sparse( p_layout, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, slot_l, MAIN_IS_FP8, MAIN_BLOCK_SIZE, @@ -1106,6 +1248,8 @@ def _pa_decode_sparse( p_layout, gather_l, gather_rope_l, + scl_l, + NARROW_SCALE, slot_l, EXTRA_IS_FP8, EXTRA_BLOCK_SIZE, From a62d73dc22943d2bb0a67f98aab5269d463ece2b Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 09:39:35 +0000 Subject: [PATCH 11/51] [GLUON][GFX950][DSV4] Add AITER_PA_DECODE_LDS_PAD knob (default unchanged) 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) --- .../_gluon_kernels/gfx950/attention/pa_decode_sparse.py | 8 +++++++- aiter/ops/triton/attention/pa_decode_sparse.py | 3 +++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 03b944e74e7..b20b7f23050 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -955,6 +955,7 @@ def _pa_decode_sparse( MFMA_K: gl.constexpr, UNIFORM: gl.constexpr, GATHER_TW1: gl.constexpr, + LDS_PAD: gl.constexpr, # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, 1 = # columns); >= the tile's extent on that axis means one shot. NOPE_CHUNK: gl.constexpr, @@ -1080,8 +1081,13 @@ def _pa_decode_sparse( warps_per_cta=[1, NUM_WARPS], order=[1, 0], ) + # LDS pad after every row. Row pitch (512 + LDS_PAD) bf16 sets which banks a + # transposed K read (ds_read_b64_tr_b16 walks down a column) lands on: + # bank = (row * pitch_dwords) mod 32, so pitch_dwords mod 32 == 4 at PAD=8 means + # 32 lanes share 8 banks. Kept as a knob because the conflict share is high + # (46%) but a high sub-metric ratio is not proof of a bottleneck. kv_shared: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( - [[HEAD_SIZE, 8]], [BLOCK_K, HEAD_SIZE], [1, 0] + [[HEAD_SIZE, LDS_PAD]], [BLOCK_K, HEAD_SIZE], [1, 0] ) h_off = pid_h * BLOCK_M diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 64d2cf201d3..4ce2c35116c 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -590,6 +590,8 @@ def _pa_decode_sparse_gfx950_gluon( # per-lane slot vector. Measured best at 32 for every extra >= 128 (-10% at # extra=128) and within noise at extra=8. gather_tw1 = int(_os.environ.get("AITER_PA_DECODE_TW1", "32")) + # AITER_PA_DECODE_LDS_PAD: bf16 elements of padding after each kv_smem row. + lds_pad = int(_os.environ.get("AITER_PA_DECODE_LDS_PAD", "8")) # AITER_PA_DECODE_NOPE_CHUNK: width of one fp8 dequant piece (0/unset = one shot). NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -811,6 +813,7 @@ def _pa_decode_sparse_gfx950_gluon( MFMA_K=MFMA_K, UNIFORM=UNIFORM, GATHER_TW1=gather_tw1, + LDS_PAD=lds_pad, NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, MAIN_SPLITS=main_splits, From 5a82565392a39c53a9bb6471e7f3293c331a5797 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 09:54:29 +0000 Subject: [PATCH 12/51] [GLUON][GFX950][DSV4] Add AITER_PA_DECODE_BLOCK_M knob (default unchanged) 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) --- aiter/ops/triton/attention/pa_decode_sparse.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 4ce2c35116c..14d42995b9f 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -578,6 +578,11 @@ def _pa_decode_sparse_gfx950_gluon( _bk = _os.environ.get("AITER_PA_DECODE_BLOCK_K") if _bk: BLOCK_K = int(_bk) + # AITER_PA_DECODE_BLOCK_M: heads per MFMA M-tile. 16 is both the MFMA M and the + # DSv4 head count; larger pads and quadruples the accumulator (see RESULTS.md). + _bm = _os.environ.get("AITER_PA_DECODE_BLOCK_M") + if _bm: + BLOCK_M = int(_bm) # AITER_PA_DECODE_MFMA_K: 32 selects CDNA4's double-rate v_mfma_f32_16x16x32_bf16 # (16 does 16x16x16, i.e. half the K per instruction). _mk = _os.environ.get("AITER_PA_DECODE_MFMA_K") From 241098970962679a1ab0f6a5ad00d9f2223c805d Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 10:06:52 +0000 Subject: [PATCH 13/51] [GLUON][GFX950][DSV4] Store the split-K accumulator partials in bf16 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) --- .../gfx950/attention/pa_decode_sparse.py | 35 +++++++++++++++---- .../ops/triton/attention/pa_decode_sparse.py | 16 ++++++++- 2 files changed, 43 insertions(+), 8 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index b20b7f23050..43f2f6269a4 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -960,6 +960,7 @@ def _pa_decode_sparse( # columns); >= the tile's extent on that axis means one shot. NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, + PART_STORE_CACHE: gl.constexpr, # How many of the NUM_SPLITS programs share the main (SWA) segment. The two # segments have very different shapes -- main is a contiguous window whose # length is fixed by the sliding window, extra is the top-k list -- so one @@ -1308,20 +1309,35 @@ def _pa_decode_sparse( # base-2 exponent domain (row-max * softmax_scale * log2e), which is the # reduce/skip_reduce convention the triton kernel also uses. pm_base = query_idx * pm_stride0 + split_id * pm_stride_s + # The reduce kernel reads these back immediately, so they want to stay + # resident rather than stream to memory. ATT puts 8.5% of all stall cycles + # on the 68 buffer_store_dwordx4 of the accumulator alone (~154 cycles + # each), and the partials are ~31% of the kernel's HBM traffic. gl.amd.cdna4.buffer_store( m_pv, ptr=part_m_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv, + cache=PART_STORE_CACHE, ) gl.amd.cdna4.buffer_store( - l_pv, ptr=part_l_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv + l_pv, + ptr=part_l_ptr + pm_base, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + cache=PART_STORE_CACHE, ) offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) a_base = query_idx * pa_stride0 + split_id * pa_stride_s a_off = (a_base + h_pv[:, None] * pa_stride_h + offs_d_a[None, :]).to(gl.int32) + # Follow part_acc's own dtype: at bf16 this halves both the partial traffic + # (~31% of the kernel's HBM bytes) and the store instruction count. gl.amd.cdna4.buffer_store( - acc, ptr=part_acc_ptr, offsets=a_off, mask=head_mask_pv[:, None] + acc.to(part_acc_ptr.dtype.element_ty), + ptr=part_acc_ptr, + offsets=a_off, + mask=head_mask_pv[:, None], + cache=PART_STORE_CACHE, ) @@ -1352,6 +1368,7 @@ def _pa_decode_sparse_reduce( NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, ADAPTIVE_SPLITS: gl.constexpr, + PART_LOAD_CACHE: gl.constexpr, ): """Split-KV combine: merge the per-split partials, fold the attn sink, and write the final output. Partials store m in the base-2 exponent domain (row-max * @@ -1398,7 +1415,8 @@ def _pa_decode_sparse_reduce( for s in range(NUM_SPLITS): base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( - ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf + ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, + cache=PART_LOAD_CACHE, ) m_final = _max2(m_final, m_s) # m_s already in base-2 exponent domain if HAS_SINK: @@ -1413,10 +1431,12 @@ def _pa_decode_sparse_reduce( for s in range(NUM_SPLITS): base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( - ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf + ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, + cache=PART_LOAD_CACHE, ) l_s = gl.amd.cdna4.buffer_load( - ptr=part_l_ptr + base, offsets=h, mask=head_mask, other=0.0 + ptr=part_l_ptr + base, offsets=h, mask=head_mask, other=0.0, + cache=PART_LOAD_CACHE, ) w = gl.exp2(m_s - m_final) l_final = l_final + w * l_s @@ -1430,9 +1450,10 @@ def _pa_decode_sparse_reduce( else: acc_mask = head_mask[:, None] acc_s = gl.amd.cdna4.buffer_load( - ptr=part_acc_ptr, offsets=a_off, mask=acc_mask, other=0.0 + ptr=part_acc_ptr, offsets=a_off, mask=acc_mask, other=0.0, + cache=PART_LOAD_CACHE, ) - acc = acc + w[:, None] * acc_s + acc = acc + w[:, None] * acc_s.to(gl.float32) if HAS_SINK: sink = gl.amd.cdna4.buffer_load( diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 14d42995b9f..220c77be7f6 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -711,9 +711,21 @@ def _pa_decode_sparse_gfx950_gluon( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device ) part_l = torch.empty_like(part_m) + # Split-K accumulator partials in bf16. They are ~31% of the kernel's HBM + # traffic and ATT puts 8.5% of all stall cycles on the 68 buffer_store of + # them, so halving both is worth 6-23% depending on split count. The + # mantissa bits it costs are far inside the error the fp8 KV format already + # carries: against an fp64 reference, max|delta|/max|ref| goes 3.517e-2 -> + # 3.567e-2 at C128A, and the bf16-vs-f32 delta (4-6e-3) is ~6x smaller than + # the error already present. Not used on the skip_reduce path, which hands + # part_acc back to the caller and must keep its dtype. + _pab = ( + _os.environ.get("AITER_PA_DECODE_PART_BF16", "1") == "1" + and not skip_reduce + ) part_acc = torch.empty( (num_queries, num_splits, num_heads, head_dim), - dtype=torch.float32, + dtype=torch.bfloat16 if _pab else torch.float32, device=q.device, ) pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) @@ -821,6 +833,7 @@ def _pa_decode_sparse_gfx950_gluon( LDS_PAD=lds_pad, NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, + PART_STORE_CACHE=_os.environ.get('AITER_PA_DECODE_PART_ST', ''), MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, @@ -870,6 +883,7 @@ def _pa_decode_sparse_gfx950_gluon( NUM_SPLITS=num_splits, HEAD_ALIGNED=num_heads % red_block_m == 0, ADAPTIVE_SPLITS=adaptive_splits, + PART_LOAD_CACHE=_os.environ.get("AITER_PA_DECODE_PART_LD", ""), num_warps=red_warps, ) return out From 1f62fea0d977805e05b30c6c72df2c7af8efba46 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 14:40:59 +0000 Subject: [PATCH 14/51] [GLUON][GFX950][DSV4] Ask for 2 waves/EU on the 64-bit gather path too 945933417 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 (e9cb7b326) 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) --- aiter/ops/triton/attention/pa_decode_sparse.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 220c77be7f6..97f992019de 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -754,12 +754,16 @@ def _pa_decode_sparse_gfx950_gluon( nope_chunk = int(_nc) if _nc else min(128, head_dim) # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD - # threshold on gfx950's 512-VGPR file. The chunked dequant gets the buffer_load - # path to 256 with a single spilled VGPR, 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 (+40% at C4A), so only ask for it - # when both caches are on the fast path. - if use_buffer_load and (chunk_axis == 0 or nope_chunk < head_dim): + # threshold on gfx950's 512-VGPR file. Asked for on both load paths. + # + # This used to be gated on buffer_load, because a cache past buffer_load's 2 GB + # offset gathers through 64-bit addresses and used to land at ~321 VGPRs, where + # the cap bought nothing and cost ~150 scratch stores (+40% at C4A). Narrowing + # the scale gather brought that path to 283, close enough that the cap now + # reaches 256 with 17 spill slots instead of 148 -- worth 0.745x at extra=272 + # and 0.833x at extra=1024 on the 64-bit path. Re-measure this gate whenever the + # register picture moves. + if chunk_axis == 0 or nope_chunk < head_dim: waves_per_eu = 2 _wp = _os.environ.get("AITER_PA_DECODE_WPEU") if _wp: From d7302ea4654d9133d38244519cc8259311689b1c Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 14:49:15 +0000 Subject: [PATCH 15/51] [GLUON][GFX950][DSV4] buffer_load the slot index gather 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) --- .../gfx950/attention/pa_decode_sparse.py | 33 +++++++++++++++++-- .../ops/triton/attention/pa_decode_sparse.py | 9 +++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 43f2f6269a4..c0ad3d1762c 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -278,17 +278,32 @@ def _slots( BLOCK_SIZE: gl.constexpr, MASKED: gl.constexpr, HAS_INVALID: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, ): # Returns in whatever layout k_pos carries. Called once per gather layout: NoPE # and RoPE tile their warps differently, and re-loading this tiny broadcast # vector is cheaper than a cross-lane convert between the two. + # 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 is also the largest single source + # of exec-mask branching in the kernel (20 of 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. if MASKED: in_range = k_pos < hi - slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1) + if IDX_BUFFER_LOAD: + slot = gl.amd.cdna4.buffer_load( + ptr=indices_ptr + seg_start, offsets=k_pos, mask=in_range, other=-1 + ) + else: + slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1) valid = in_range & (slot >= 0) & (slot < num_rows) slot = gl.where(valid, slot, 0) else: - slot = gl.load(indices_ptr + seg_start + k_pos) + if IDX_BUFFER_LOAD: + slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=k_pos) + else: + slot = gl.load(indices_ptr + seg_start + k_pos) valid = slot >= 0 # -1 sentinels: clamp in-bounds, mask score below if HAS_INVALID: slot = gl.where(valid, slot, 0) @@ -338,6 +353,7 @@ def _decode_tile( MASKED: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -357,6 +373,7 @@ def _decode_tile( BLOCK_SIZE, MASKED, HAS_INVALID, + IDX_BUFFER_LOAD, ) block_idx_g = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_l)) pos_g = gl.convert_layout(pos, gl.SliceLayout(1, gather_l)) @@ -425,6 +442,7 @@ def _decode_tile( BLOCK_SIZE, MASKED, HAS_INVALID, + IDX_BUFFER_LOAD, ) rope_row = block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224 if MASKED: @@ -514,6 +532,7 @@ def _gd_fp8( HEAD_SIZE: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -536,6 +555,7 @@ def _gd_fp8( BLOCK_SIZE, False, HAS_INVALID, + IDX_BUFFER_LOAD, ) if UNIFORM: NGRP: gl.constexpr = HEAD_SIZE // 64 @@ -569,6 +589,7 @@ def _gd_fp8( BLOCK_SIZE, False, HAS_INVALID, + IDX_BUFFER_LOAD, ) k_rope = _cache_load( cache_bf16_ptr, bgr * (cs0 // 2) + pgr * 288 + 224, offs_rope, USE_BUFFER_LOAD @@ -685,6 +706,7 @@ def _process_segment( HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -719,6 +741,7 @@ def _process_segment( HEAD_SIZE, UNIFORM, USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -742,6 +765,7 @@ def _process_segment( HEAD_SIZE, UNIFORM, USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -848,6 +872,7 @@ def _process_segment( False, UNIFORM, USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -895,6 +920,7 @@ def _process_segment( True, UNIFORM, USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -979,6 +1005,7 @@ def _pa_decode_sparse( # partner is large. At tiny top-k the main/SWA gather is ~94% of the tokens. MAIN_USE_BUFFER_LOAD: gl.constexpr, EXTRA_USE_BUFFER_LOAD: gl.constexpr, + IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -1224,6 +1251,7 @@ def _pa_decode_sparse( HEAD_ALIGNED, UNIFORM, MAIN_USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) @@ -1270,6 +1298,7 @@ def _pa_decode_sparse( HEAD_ALIGNED, UNIFORM, EXTRA_USE_BUFFER_LOAD, + IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, ) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 97f992019de..91dca01f79b 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -676,6 +676,14 @@ def _pa_decode_sparse_gfx950_gluon( # arch/format asserts below.) main_use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES extra_use_buffer_load = max_addressable_bytes(extra_cache) < MAX_BYTES + # The index lists are one int32 per gathered token, orders of magnitude under + # the 2 GB buffer offset limit even for a full batch, so they keep the fast path + # regardless of how big the KV caches are. + idx_use_buffer_load = ( + _os.environ.get("AITER_PA_DECODE_IDX_BUF", "1") == "1" + and max_addressable_bytes(indices) < MAX_BYTES + and max_addressable_bytes(extra_indices) < MAX_BYTES + ) use_buffer_load = main_use_buffer_load and extra_use_buffer_load # AITER_PA_DECODE_FORCE_GLOBAL_LOAD emulates a >2GB production cache without # allocating one: 1 = both caches, main = main only, extra = extra only. @@ -842,6 +850,7 @@ def _pa_decode_sparse_gfx950_gluon( ADAPTIVE_SPLITS=adaptive_splits, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, + IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, num_warps=num_warps, From 59617a41a49ca59069e40bc3fdf04b9adcfd098d Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 15:36:27 +0000 Subject: [PATCH 16/51] [GLUON][GFX950][DSV4] Reciprocal-multiply instead of divide in both epilogues 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) --- .../_gluon_kernels/gfx950/attention/pa_decode_sparse.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index c0ad3d1762c..f4d2fbcd780 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1322,7 +1322,8 @@ def _pa_decode_sparse( acc = acc * alpha[:, None] else: l_final = l_pv - out = acc / l_final[:, None] + one_over_l = 1.0 / l_final + out = acc * one_over_l[:, None] offs_d_o = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) o_off = ( query_idx * out_stride0 + h_pv[:, None] * out_stride1 + offs_d_o[None, :] @@ -1490,7 +1491,11 @@ def _pa_decode_sparse_reduce( ).to(gl.float32) l_final = l_final + gl.exp2(sink * RCP_LN2 - m_final) - out = acc / l_final[:, None] + # One reciprocal per row, then a broadcast multiply: an f32 divide lowers to a + # ~10-instruction sequence, so dividing the whole [BLOCK_M, HEAD_SIZE] tile costs + # that per element. + one_over_l = 1.0 / l_final + out = acc * one_over_l[:, None] o_off = (query_idx * out_stride0 + h[:, None] * out_stride1 + offs_d[None, :]).to( gl.int32 ) From 9cf6a0c1b13537c5516cb6ef988ee9e1b55321b4 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 19:57:02 +0000 Subject: [PATCH 17/51] [GLUON][GFX950][DSV4] Add AITER_PA_DECODE_CS0_ALIGN override (default 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) --- aiter/ops/triton/attention/pa_decode_sparse.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 91dca01f79b..25cf430d3e1 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -671,6 +671,13 @@ def _pa_decode_sparse_gfx950_gluon( if _s0 % _a == 0 and _s1 % _a == 0: cs0_align = _a break + # AITER_PA_DECODE_CS0_ALIGN: override for re-measuring the hint (1 = off). As of + # the row/column gather split the hint is codegen-neutral -- the column offsets are + # compile-time constants, so vectorization no longer depends on knowing cs0's + # divisibility, and the emitted load mix is identical with it on and off. + _ca = _os.environ.get("AITER_PA_DECODE_CS0_ALIGN") + if _ca: + cs0_align = int(_ca) # Gate each cache on its OWN span: one oversized cache must not disable the # fast path for the other. (cache_bytes, the max of the two, is kept for the # arch/format asserts below.) From 19f5977b04d01ff123bcff46a5439f056c876626 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 21:25:27 +0000 Subject: [PATCH 18/51] [GLUON][GFX950][DSV4] Fuse the fp8xE8M0 dequant, and spend the registers 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 --- .../gfx950/attention/pa_decode_sparse.py | 212 +++++++++++++++--- .../ops/triton/attention/pa_decode_sparse.py | 80 ++++++- 2 files changed, 248 insertions(+), 44 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index f4d2fbcd780..dc52e5b3742 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -186,10 +186,59 @@ def _split_ax(x, AXIS: gl.constexpr): return a, b +# fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. $0 takes the low two +# fp8 bytes of the 32-bit source, $1 the high two via op_sel. "=&v" (early clobber) is +# load-bearing: without it the allocator may give an output the same VGPR as the scale +# input, and the first convert then clobbers the scale before the second reads it -- +# silently wrong in exactly the elements $1 produces. +_DEQ_ASM: gl.constexpr = gl.constexpr( + "v_cvt_scalef32_pk_bf16_fp8 $0, $2, $3\n" + "v_cvt_scalef32_pk_bf16_fp8 $1, $2, $3 op_sel:[1,0,0]" +) +# $0 must be early-clobber: the second convert re-reads $2/$3 after $0 is written, +# so an output sharing a register with an input would read a clobbered scale. $1 is +# written last, after every read, so it may alias an input -- one more reuse for the +# allocator. 4 constraints, not 5: an i16 scale packs into one register where an f32 +# scale took two ($4 was never named). +_DEQ_CONS: gl.constexpr = gl.constexpr("=&v,=v,v,v") + + +@gluon.jit +def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): + """[BLOCK_K, W8/2] int16 (4 packed fp8 per 32 bits) + raw E8M0 byte -> [BLOCK_K, W8] bf16. + + The scale operand is read as E8M0 -- the hardware takes bits [30:23] and nothing + else. Those are bits [14:7] of the operand's high half, so a 16-bit `e << 7` in + both halves of the register puts the exponent exactly where an f32 `e << 23` + would, and an i16 operand costs ONE register at pack=2 where f32 cost two. + Converting to f32 with exp2 also works but is a longer chain for nothing. Because the scale is a power of two this is bit-identical to the f32 + chain: fp8 -> bf16 is exact (3 mantissa bits into 8) and scaling by 2^k stays exact. + """ + lo, hi = gl.inline_asm_elementwise( + _DEQ_ASM, + _DEQ_CONS, + [x16, (e_u8.to(gl.uint16) << 7)], + dtype=(gl.bfloat16, gl.bfloat16), + is_pure=True, + pack=2, + ) + # lo carries fp8 elements 4i+{0,1} and hi 4i+{2,3}, so the flat order is + # 4i + 2*lohi + s. A lane's 8 int16 are the same 16-byte run the u8 gather used, + # so this whole interleave stays inside one lane -- assert_trivial proves it. + W16: gl.constexpr = W8 // 2 + lo3 = lo.reshape(lo.shape[0], W16 // 2, 2) + hi3 = hi.reshape(hi.shape[0], W16 // 2, 2) + both = gl.join(lo3, hi3).permute(0, 1, 3, 2).reshape(lo.shape[0], W8) + # Back to the byte-gather layout, or the kv_smem store sees the join/permute layout + # and lowers to narrow ds_writes (112 ds_write_b128 -> 28, and ~1.5x slower). + return gl.convert_layout(both, out_l, assert_trivial=True) + + @gluon.jit def _deq_store( x_u8, sc, kv_smem, off, AXIS: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, + ASM_DEQ: gl.constexpr, gather_l: gl.constexpr, ): """Dequant one fp8 slab and write it straight to kv_smem[:, off:off+W]. @@ -197,15 +246,30 @@ def _deq_store( an emulated mul and doubles the loop. ``sc`` is the raw per-64-group exponent byte (packed fp8_ds_mla, UE8M0) or an f32 scale (uniform pool). """ - if UNIFORM: - scale = sc - else: - scale = gl.exp2(sc.to(gl.float32) - 127.0) - val = (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) - if AXIS == 1: - kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) + if ASM_DEQ: + # x_u8 is the int16 view here (see _gd_fp8), so its column count is W8/2. + W8: gl.constexpr = x_u8.shape[1] * 2 + # Adjacent fp8 columns share a scale (groups are 64 wide, so even), which makes + # dropping every other broadcast column exact. + s_even, _ = gl.split(sc.reshape(sc.shape[0], sc.shape[1] // 2, 2)) + # split leaves a SliceLayout of a 3-D parent; the asm needs both operands in + # the int16 gather's own layout. + s_even = gl.convert_layout(s_even, x_u8.type.layout) + val = _deq_asm(x_u8.to(gl.int16, bitcast=True), s_even, W8, gather_l) + if AXIS == 1: + kv_smem.slice(off, W8, dim=1).store(val) + else: + kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) else: - kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) + if UNIFORM: + scale = sc + else: + scale = gl.exp2(sc.to(gl.float32) - 127.0) + val = (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) + if AXIS == 1: + kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) + else: + kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) @gluon.jit @@ -215,6 +279,8 @@ def _deq_store_tile( kv_smem, NOPE_CHUNK: gl.constexpr, AXIS: gl.constexpr, + ASM_DEQ: gl.constexpr, + gather_l: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, ): @@ -229,14 +295,14 @@ def _deq_store_tile( """ W: gl.constexpr = x_u8.shape[1] if AXIS == 1 else x_u8.shape[0] if NOPE_CHUNK >= W: - _deq_store(x_u8, sc, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x_u8, sc, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) else: x0, x1 = _split_ax(x_u8, AXIS) s0, s1 = _split_ax(sc, AXIS) W2: gl.constexpr = W // 2 if NOPE_CHUNK >= W2: - _deq_store(x0, s0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(x1, s1, kv_smem, W2, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x0, s0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x1, s1, kv_smem, W2, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) else: x00, x01 = _split_ax(x0, AXIS) s00, s01 = _split_ax(s0, AXIS) @@ -244,28 +310,28 @@ def _deq_store_tile( s10, s11 = _split_ax(s1, AXIS) W4: gl.constexpr = W // 4 if NOPE_CHUNK >= W4: - _deq_store(x00, s00, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(x01, s01, kv_smem, W4, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(x10, s10, kv_smem, 2 * W4, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(x11, s11, kv_smem, 3 * W4, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(x00, s00, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x01, s01, kv_smem, W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x10, s10, kv_smem, 2 * W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x11, s11, kv_smem, 3 * W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) else: W8: gl.constexpr = W // 8 y0, y1 = _split_ax(x00, AXIS) t0, t1 = _split_ax(s00, AXIS) - _deq_store(y0, t0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y0, t0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y1, t1, kv_smem, W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) y0, y1 = _split_ax(x01, AXIS) t0, t1 = _split_ax(s01, AXIS) - _deq_store(y0, t0, kv_smem, 2 * W8, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 3 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y0, t0, kv_smem, 2 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y1, t1, kv_smem, 3 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) y0, y1 = _split_ax(x10, AXIS) t0, t1 = _split_ax(s10, AXIS) - _deq_store(y0, t0, kv_smem, 4 * W8, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 5 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y0, t0, kv_smem, 4 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y1, t1, kv_smem, 5 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) y0, y1 = _split_ax(x11, AXIS) t0, t1 = _split_ax(s11, AXIS) - _deq_store(y0, t0, kv_smem, 6 * W8, AXIS, UNIFORM, FP8_FNUZ) - _deq_store(y1, t1, kv_smem, 7 * W8, AXIS, UNIFORM, FP8_FNUZ) + _deq_store(y0, t0, kv_smem, 6 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y1, t1, kv_smem, 7 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) @gluon.jit @@ -328,6 +394,7 @@ def _decode_tile( qk_scale, kv_smem, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, @@ -338,6 +405,7 @@ def _decode_tile( p_layout: gl.constexpr, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + gather16_l: gl.constexpr, scl_l: gl.constexpr, NARROW_SCALE: gl.constexpr, IS_FP8: gl.constexpr, @@ -349,6 +417,7 @@ def _decode_tile( BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, + ASM_DEQ: gl.constexpr, HEAD_ALIGNED: gl.constexpr, MASKED: gl.constexpr, UNIFORM: gl.constexpr, @@ -402,7 +471,7 @@ def _decode_tile( else: x_u8 = _cache_load(cache_ptr, kv_row, offs_full, USE_BUFFER_LOAD) sc = _cache_load(cache_bf16_ptr, scl_row, scl_col, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, True, FP8_FNUZ) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, True, FP8_FNUZ) elif IS_FP8: # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. nope_row = block_idx_g * cs0 + pos_g * 576 @@ -419,10 +488,25 @@ def _decode_tile( cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD, mask=valid_g[:, None], other=127, ) - x_u8 = _cache_load( - cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD, - mask=valid_g[:, None], other=0, - ) + if ASM_DEQ: + # Same byte as nope_row, addressed in 2-byte units: one shift on a + # live vector, and block_idx_g * cs0 stays common with the scale + # gather. Exact -- cs0 and 576 are both even. + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, gather16_l) + ) + x_u8 = _cache_load( + cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD, + mask=gl.convert_layout( + valid1d, gl.SliceLayout(1, gather16_l) + )[:, None], + other=0.0, + ) + else: + x_u8 = _cache_load( + cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD, + mask=valid_g[:, None], other=0, + ) else: if NARROW_SCALE and not USE_BUFFER_LOAD: exps = _scale_load( @@ -431,8 +515,19 @@ def _decode_tile( ) else: exps = _cache_load(cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD) - x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, False, FP8_FNUZ) + if ASM_DEQ: + # Same byte as nope_row, addressed in 2-byte units: one shift on a + # live vector, and block_idx_g * cs0 stays common with the scale + # gather. Exact -- cs0 and 576 are both even. + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, gather16_l) + ) + x_u8 = _cache_load( + cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD + ) + else: + x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) + _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, False, FP8_FNUZ) block_idx_gr, pos_gr, valid_gr = _slots( indices_ptr, seg_start, @@ -521,13 +616,16 @@ def _gd_fp8( k_start, cs0, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + gather16_l: gl.constexpr, scl_l: gl.constexpr, NARROW_SCALE: gl.constexpr, + ASM_DEQ: gl.constexpr, BLOCK_SIZE: gl.constexpr, HEAD_SIZE: gl.constexpr, UNIFORM: gl.constexpr, @@ -579,7 +677,22 @@ def _gd_fp8( ) else: sc = _cache_load(cache_ptr, scl_row, offs_full // 64, USE_BUFFER_LOAD) - x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) + if ASM_DEQ: + # 2-byte elements so the asm gets <2 x i16> = 4 packed fp8 in one VGPR out + # of a single dword load. Same bytes, same 16-B per-lane run, half the + # columns; cache_bf16_ptr is already the 2-byte view of this cache and the + # RoPE gather below uses the same cs0//2 / *288 addressing. + # The row vector comes from _slots in SliceLayout(1, gather_l); the i16 + # columns live in gather16_l, and a broadcast needs one parent layout. The + # two layouts share their dim-0 tiling, so this convert is a rename. + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, gather16_l) + ) + x_u8 = _cache_load( + cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD + ) + else: + x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) bgr, pgr, _ = _slots( indices_ptr, seg_start, @@ -615,6 +728,7 @@ def _qkpv_fp8( k_layout: gl.constexpr, v_layout: gl.constexpr, p_layout: gl.constexpr, + gather_l: gl.constexpr, NOPE_DIM: gl.constexpr, ROPE_DIM: gl.constexpr, HEAD_SIZE: gl.constexpr, @@ -622,6 +736,7 @@ def _qkpv_fp8( BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, + ASM_DEQ: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, @@ -631,7 +746,7 @@ def _qkpv_fp8( skips the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, mask the columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" neg_inf = float("-inf") - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, UNIFORM, FP8_FNUZ) + _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, UNIFORM, FP8_FNUZ) if not UNIFORM: kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) k = kv_smem.permute([1, 0]).load(k_layout) @@ -691,6 +806,7 @@ def _process_segment( p_layout: gl.constexpr, gather_l: gl.constexpr, gather_rope_l: gl.constexpr, + gather16_l: gl.constexpr, scl_l: gl.constexpr, NARROW_SCALE: gl.constexpr, slot_l: gl.constexpr, @@ -703,6 +819,7 @@ def _process_segment( BLOCK_K: gl.constexpr, NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, + ASM_DEQ: gl.constexpr, HEAD_ALIGNED: gl.constexpr, UNIFORM: gl.constexpr, USE_BUFFER_LOAD: gl.constexpr, @@ -711,6 +828,8 @@ def _process_segment( FP8_FNUZ: gl.constexpr, ): offs_full = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, gather_l)) + # int16 column offsets for the asm dequant's 2-byte gather (half the columns). + offs_full16 = gl.arange(0, HEAD_SIZE // 2, layout=gl.SliceLayout(0, gather16_l)) offs_rope = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, gather_rope_l)) k_rng_slot = gl.arange(0, BLOCK_K, layout=slot_l) k_rng_rope = gl.arange(0, BLOCK_K, layout=gl.SliceLayout(1, gather_rope_l)) @@ -730,13 +849,16 @@ def _process_segment( lo, cs0, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, + ASM_DEQ, BLOCK_SIZE, HEAD_SIZE, UNIFORM, @@ -754,13 +876,16 @@ def _process_segment( lo + i * BLOCK_K, cs0, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, + ASM_DEQ, BLOCK_SIZE, HEAD_SIZE, UNIFORM, @@ -786,6 +911,7 @@ def _process_segment( k_layout, v_layout, p_layout, + gather_l, NOPE_DIM, ROPE_DIM, HEAD_SIZE, @@ -793,6 +919,7 @@ def _process_segment( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, UNIFORM, FP8_FNUZ, @@ -816,6 +943,7 @@ def _process_segment( k_layout, v_layout, p_layout, + gather_l, NOPE_DIM, ROPE_DIM, HEAD_SIZE, @@ -823,6 +951,7 @@ def _process_segment( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, UNIFORM, FP8_FNUZ, @@ -847,6 +976,7 @@ def _process_segment( qk_scale, kv_smem, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, @@ -857,6 +987,7 @@ def _process_segment( p_layout, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, IS_FP8, @@ -868,6 +999,7 @@ def _process_segment( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, False, UNIFORM, @@ -895,6 +1027,7 @@ def _process_segment( qk_scale, kv_smem, offs_full, + offs_full16, offs_rope, k_rng_slot, k_rng_rope, @@ -905,6 +1038,7 @@ def _process_segment( p_layout, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, IS_FP8, @@ -916,6 +1050,7 @@ def _process_segment( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, True, UNIFORM, @@ -998,6 +1133,7 @@ def _pa_decode_sparse( MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. ADAPTIVE_SPLITS: gl.constexpr, + ASM_DEQ: gl.constexpr, # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache # whose span exceeds that must gather via 64-bit gl.load -- but the two caches # are sized independently (SWA window vs full compressed history), so gating @@ -1102,6 +1238,14 @@ def _pa_decode_sparse( order=[2, 1, 0], ) scl_l: gl.constexpr = gl.SliceLayout(2, scl_l3) + # Same tiling as gather_l but 2-byte elements: half the per-thread run, so a lane + # still covers the same 16 bytes and the lo/hi interleave stays lane-local. + gather16_l: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, GSPT // 2], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) slot_l: gl.constexpr = gl.SliceLayout(1, gather_l) blocked_q: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, 8], @@ -1236,6 +1380,7 @@ def _pa_decode_sparse( p_layout, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, slot_l, @@ -1248,6 +1393,7 @@ def _pa_decode_sparse( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, UNIFORM, MAIN_USE_BUFFER_LOAD, @@ -1283,6 +1429,7 @@ def _pa_decode_sparse( p_layout, gather_l, gather_rope_l, + gather16_l, scl_l, NARROW_SCALE, slot_l, @@ -1295,6 +1442,7 @@ def _pa_decode_sparse( BLOCK_K, NOPE_CHUNK, CHUNK_AXIS, + ASM_DEQ, HEAD_ALIGNED, UNIFORM, EXTRA_USE_BUFFER_LOAD, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 25cf430d3e1..b95abeb1628 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -708,18 +708,49 @@ def _pa_decode_sparse_gfx950_gluon( # AITER_PA_DECODE_SPLITS: experiment override for the split-K count. _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") - if _sp: - num_splits = max(1, int(_sp)) - elif kv_splits is not None: - num_splits = max(1, int(kv_splits)) - elif _os.environ.get("AITER_PA_DECODE_SPLIT_POLICY", "occ") == "occ": - num_splits = _decode_num_splits_occ( - num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K - ) - else: - num_splits = _decode_num_splits( - num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K - ) + + def _pick_splits(hb): + if _sp: + return max(1, int(_sp)) + if kv_splits is not None: + return max(1, int(kv_splits)) + if _os.environ.get("AITER_PA_DECODE_SPLIT_POLICY", "occ") == "occ": + return _decode_num_splits_occ( + num_queries, hb, avg_main, avg_extra, BLOCK_K + ) + return _decode_num_splits(num_queries, hb, avg_main, avg_extra, BLOCK_K) + + num_splits = _pick_splits(heads_blocks) + + # Short top-k is latency- and coverage-bound rather than throughput-bound: at + # <= 2 KV tiles per query the whole launch is a few dozen workgroups on a + # 256-CU part, so the win is MORE programs, not fatter ones. Halving BLOCK_M + # doubles heads_blocks, which is real parallelism -- unlike a finer split-K, + # which hands every split a partial tile (1.14-1.36x worse when forced). Both + # gates matter: measured 2-5% better at every concurrency on top-k 8/64, and + # 1.06-1.22x WORSE on top-k 272/1024, where the halved MFMA M dominates. + _tiles = max( + math.ceil(avg_main / BLOCK_K) if avg_main > 0 else 1, + math.ceil(avg_extra / BLOCK_K) if avg_extra > 0 else 1, + ) + if ( + _bm is None + and BLOCK_M == 16 + and num_heads % 8 == 0 + and _tiles <= 2 + and num_queries * heads_blocks * num_splits <= get_num_sms() + ): + _hb8 = (num_heads + 7) // 8 + _ns8 = _pick_splits(_hb8) + # Doubling heads_blocks can push the launch past one workgroup per CU, and + # the split policy then answers with fewer splits. At C=128 short top-k that + # is 2 -> 1: split-K disappears and the shape goes 0.94x -> 1.11x. Take the + # halving only when it costs neither a split nor the occupancy property. + if _ns8 >= num_splits and num_queries * _hb8 * _ns8 <= get_num_sms(): + BLOCK_M = 8 + HEAD_ALIGNED = num_heads % BLOCK_M == 0 + heads_blocks = _hb8 + num_splits = _ns8 if num_splits > 1: part_m = torch.empty( @@ -780,6 +811,16 @@ def _pa_decode_sparse_gfx950_gluon( # register picture moves. if chunk_axis == 0 or nope_chunk < head_dim: waves_per_eu = 2 + # num_warps is 4 -- one wave per SIMD per workgroup -- so the second occupancy + # slot exists only if a CU gets a second workgroup. When the entire launch is at + # most one workgroup per CU it cannot, and the 2-waves/EU register cap is then + # pure loss: it spills ~29 dwords across the per-tile softmax row-max reduction + # and buys no overlap. Measured on the main kernel at C in {4,16,32,64,128} x + # top-k in {64,272,1024}: 1 wave/EU + ASM_DEQ is 0.927-0.953x at every shape with + # launch <= CU, and 1.19-1.43x *worse* at every shape above it. + one_wg_per_cu = num_queries * heads_blocks * num_splits <= get_num_sms() + if one_wg_per_cu: + waves_per_eu = 1 _wp = _os.environ.get("AITER_PA_DECODE_WPEU") if _wp: waves_per_eu = int(_wp) @@ -801,6 +842,20 @@ def _pa_decode_sparse_gfx950_gluon( and _os.environ.get("AITER_PA_DECODE_ADAPTIVE", "1") == "1" ) + # AITER_PA_DECODE_ASM_DEQ: fuse the fp8 x E8M0 dequant into + # v_cvt_scalef32_pk_bf16_fp8 via inline asm. gfx950 only, packed OCP fp8 only + # (UNIFORM scales are f32 and FNUZ is a different encoding), and bit-identical + # there because the scale is a power of two. + # Default follows the occupancy decision: the fused convert is a shorter + # dependent chain but a wider live range, so it pays only where the extra + # registers are free (see one_wg_per_cu). 0/1 forces it either way. + asm_deq = ( + _os.environ.get("AITER_PA_DECODE_ASM_DEQ", "1" if one_wg_per_cu else "0") == "1" + and not UNIFORM + and not fp8_fnuz + and main_is_fp8 + ) + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, @@ -855,6 +910,7 @@ def _pa_decode_sparse_gfx950_gluon( PART_STORE_CACHE=_os.environ.get('AITER_PA_DECODE_PART_ST', ''), MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, + ASM_DEQ=asm_deq, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, IDX_BUFFER_LOAD=idx_use_buffer_load, From 814d53d7c93d6713a2cb4aa70d4d70ed0a235494 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 23:00:03 +0000 Subject: [PATCH 19/51] [GLUON][GFX950][DSV4] Keep the occupancy rule off the 64-bit gather path; add GRID_ORDER Follow-up to 19f5977b0, 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 --- .../gfx950/attention/pa_decode_sparse.py | 9 ++++++--- aiter/ops/triton/attention/pa_decode_sparse.py | 14 ++++++++++++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index dc52e5b3742..2622204cac2 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1130,6 +1130,7 @@ def _pa_decode_sparse( # extra still wants the CTAs. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an # empty main range and contribute an extra-only partial. + GRID_ORDER: gl.constexpr, MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. ADAPTIVE_SPLITS: gl.constexpr, @@ -1160,9 +1161,11 @@ def _pa_decode_sparse( if CS0_ALIGN > 1: main_cs0 = gl.multiple_of(main_cs0, CS0_ALIGN) extra_cs0 = gl.multiple_of(extra_cs0, CS0_ALIGN) - query_idx = gl.program_id(0) - split_id = gl.program_id(1) - pid_h = gl.program_id(2) + # GRID_ORDER names the launch axes in grid-dim order (see the driver): dim 0 + # varies fastest, which is what decides XCD/L2 sharing. + query_idx = gl.program_id(GRID_ORDER.index("q")) + split_id = gl.program_id(GRID_ORDER.index("s")) + pid_h = gl.program_id(GRID_ORDER.index("h")) qk_layout: gl.constexpr = gl.amd.AMDMFMALayout( version=4, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index b95abeb1628..112c230137d 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -818,7 +818,10 @@ def _pick_splits(hb): # and buys no overlap. Measured on the main kernel at C in {4,16,32,64,128} x # top-k in {64,272,1024}: 1 wave/EU + ASM_DEQ is 0.927-0.953x at every shape with # launch <= CU, and 1.19-1.43x *worse* at every shape above it. - one_wg_per_cu = num_queries * heads_blocks * num_splits <= get_num_sms() + one_wg_per_cu = ( + use_buffer_load + and num_queries * heads_blocks * num_splits <= get_num_sms() + ) if one_wg_per_cu: waves_per_eu = 1 _wp = _os.environ.get("AITER_PA_DECODE_WPEU") @@ -856,7 +859,13 @@ def _pick_splits(hb): and main_is_fp8 ) - grid = (num_queries, num_splits, heads_blocks) + # AITER_PA_DECODE_GRID_ORDER: which launch axis varies fastest. Dim 0 is the + # fastest-varying and XCD assignment is round-robin over the linear workgroup id, + # so this decides what shares an XCD's L2. "qsh" is the historical order. + _go = _os.environ.get("AITER_PA_DECODE_GRID_ORDER", "qsh") + _ax = {"q": num_queries, "s": num_splits, "h": heads_blocks} + assert sorted(_go) == ["h", "q", "s"], f"bad AITER_PA_DECODE_GRID_ORDER {_go!r}" + grid = tuple(_ax[c] for c in _go) _pa_decode_sparse_gfx950[grid]( q, cache, @@ -908,6 +917,7 @@ def _pick_splits(hb): NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE=_os.environ.get('AITER_PA_DECODE_PART_ST', ''), + GRID_ORDER=_go, MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, ASM_DEQ=asm_deq, From 7cd9465848bc2716f1406028e84234c1f8e8c90a Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 18 Aug 2026 23:34:42 +0000 Subject: [PATCH 20/51] [GLUON][GFX950][DSV4] Emit the fused dequant as two single-output asm calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../gfx950/attention/pa_decode_sparse.py | 42 +++++++++---------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 2622204cac2..44e5f3cb9ab 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -186,21 +186,22 @@ def _split_ax(x, AXIS: gl.constexpr): return a, b -# fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. $0 takes the low two -# fp8 bytes of the 32-bit source, $1 the high two via op_sel. "=&v" (early clobber) is -# load-bearing: without it the allocator may give an output the same VGPR as the scale -# input, and the first convert then clobbers the scale before the second reads it -- -# silently wrong in exactly the elements $1 produces. -_DEQ_ASM: gl.constexpr = gl.constexpr( - "v_cvt_scalef32_pk_bf16_fp8 $0, $2, $3\n" - "v_cvt_scalef32_pk_bf16_fp8 $1, $2, $3 op_sel:[1,0,0]" +# fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. Emitted as two +# separate single-output asm calls, NOT as one two-output blob: a single-instruction +# asm reads all its sources before writing dst, so no output can clobber an input and +# neither needs early clobber. The blob form does -- without `=&v` on its first output +# the allocator may give it the scale's register and the second convert then reads a +# clobbered scale, wrong in exactly the elements it produces -- and that forced +# liveness costs spills: 21 st/28 ld against 11/12 here, on a kernel with zero +# headroom at the 256-VGPR cap. The split form costs +184 instructions (each call +# re-reads the sources) and is worth it: 0.970-0.977x against the blob's 0.983-1.017x +# across C=16/64/128 at short top-k, and it removes the reduce knock-on that the +# blob's scratch traffic caused by evicting the split-K partials from L2. +_DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") +_DEQ_HI: gl.constexpr = gl.constexpr( + "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" ) -# $0 must be early-clobber: the second convert re-reads $2/$3 after $0 is written, -# so an output sharing a register with an input would read a clobbered scale. $1 is -# written last, after every read, so it may alias an input -- one more reuse for the -# allocator. 4 constraints, not 5: an i16 scale packs into one register where an f32 -# scale took two ($4 was never named). -_DEQ_CONS: gl.constexpr = gl.constexpr("=&v,=v,v,v") +_DEQ_CONS: gl.constexpr = gl.constexpr("=v,v,v") @gluon.jit @@ -214,13 +215,12 @@ def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): Converting to f32 with exp2 also works but is a longer chain for nothing. Because the scale is a power of two this is bit-identical to the f32 chain: fp8 -> bf16 is exact (3 mantissa bits into 8) and scaling by 2^k stays exact. """ - lo, hi = gl.inline_asm_elementwise( - _DEQ_ASM, - _DEQ_CONS, - [x16, (e_u8.to(gl.uint16) << 7)], - dtype=(gl.bfloat16, gl.bfloat16), - is_pure=True, - pack=2, + sc16 = e_u8.to(gl.uint16) << 7 + lo = gl.inline_asm_elementwise( + _DEQ_LO, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 + ) + hi = gl.inline_asm_elementwise( + _DEQ_HI, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 ) # lo carries fp8 elements 4i+{0,1} and hi 4i+{2,3}, so the flat order is # 4i + 2*lohi + s. A lane's 8 int16 are the same 16-byte run the u8 gather used, From b585afce539fe700e7b3a9babb315310ca713fce Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 00:37:06 +0000 Subject: [PATCH 21/51] [GLUON][GFX950][DSV4] UNI_TILE: one tile body for full and partial tiles 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 --- .../gfx950/attention/pa_decode_sparse.py | 83 ++++++++++++++++--- .../ops/triton/attention/pa_decode_sparse.py | 9 ++ 2 files changed, 82 insertions(+), 10 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 44e5f3cb9ab..0b9068b2b21 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -345,6 +345,7 @@ def _slots( MASKED: gl.constexpr, HAS_INVALID: gl.constexpr, IDX_BUFFER_LOAD: gl.constexpr, + UNI_TILE: gl.constexpr = False, ): # Returns in whatever layout k_pos carries. Called once per gather layout: NoPE # and RoPE tile their warps differently, and re-loading this tiny broadcast @@ -355,7 +356,23 @@ def _slots( # of exec-mask branching in the kernel (20 of 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. - if MASKED: + if UNI_TILE: + # One body for full and partial tiles. Read a REAL index -- k_pos clamped into + # the segment -- instead of masking the load, and fold the range test into + # `valid`, which _qkpv_fp8 already turns into a -inf score mask. An + # out-of-range lane then gathers a genuine row whose score is -inf, which is + # exactly how -1 sentinels are already handled on the full-tile path. + in_range = k_pos < hi + kp = gl.minimum(k_pos, gl.maximum(hi - 1, 0)) + if IDX_BUFFER_LOAD: + slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=kp) + else: + slot = gl.load(indices_ptr + seg_start + kp) + valid = in_range + if HAS_INVALID: + valid = valid & (slot >= 0) & (slot < num_rows) + slot = gl.where(valid, slot, 0) + elif MASKED: in_range = k_pos < hi if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load( @@ -633,6 +650,9 @@ def _gd_fp8( IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, + UNI_TILE: gl.constexpr = False, + seg_hi=0, + num_rows=0, ): """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues an iteration early. @@ -648,12 +668,13 @@ def _gd_fp8( indices_ptr, seg_start, k_start + k_rng_slot, - 0, # hi/num_rows: unused at MASKED=False - 0, + seg_hi, # hi/num_rows: unused unless UNI_TILE + num_rows, BLOCK_SIZE, False, HAS_INVALID, IDX_BUFFER_LOAD, + UNI_TILE, ) if UNIFORM: NGRP: gl.constexpr = HEAD_SIZE // 64 @@ -697,12 +718,13 @@ def _gd_fp8( indices_ptr, seg_start, k_start + k_rng_rope, - 0, - 0, + seg_hi, + num_rows, BLOCK_SIZE, False, HAS_INVALID, IDX_BUFFER_LOAD, + UNI_TILE, ) k_rope = _cache_load( cache_bf16_ptr, bgr * (cs0 // 2) + pgr * 288 + 224, offs_rope, USE_BUFFER_LOAD @@ -741,6 +763,9 @@ def _qkpv_fp8( UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, HAS_INVALID: gl.constexpr, + UNI_TILE: gl.constexpr = False, + k_start=0, + seg_hi=0, ): """Dequant a prefetched fp8 tile into LDS, then QK -> softmax -> PV. UNIFORM skips the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, @@ -753,10 +778,23 @@ def _qkpv_fp8( S = gl.amd.cdna4.mfma( q_dot, k, gl.zeros([BLOCK_M, BLOCK_K], gl.float32, layout=qk_layout) ) - NEED_MASK: gl.constexpr = HAS_INVALID or (not HEAD_ALIGNED) + # UNI_TILE folds the tile's range test into `valid`, so the score mask is no + # longer only about -1 sentinels: it is what makes the partial last tile correct. + COL_MASK: gl.constexpr = HAS_INVALID or UNI_TILE + NEED_MASK: gl.constexpr = COL_MASK or (not HEAD_ALIGNED) if NEED_MASK: - if HAS_INVALID: - col_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] + if COL_MASK: + if UNI_TILE and not HAS_INVALID: + # Build the range mask directly in the MFMA layout. Converting the + # slot-layout `valid` vector instead costs a cross-lane layout change + # per tile, which at 16 tiles ate the whole saving. + col_mask = ( + k_start + + gl.arange(0, BLOCK_K, layout=gl.SliceLayout(0, qk_layout)) + < seg_hi + )[None, :] + else: + col_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] if not HEAD_ALIGNED: col_mask = ( gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] @@ -784,6 +822,7 @@ def _qkpv_fp8( @gluon.jit def _process_segment( + UNI_TILE: gl.constexpr, q_dot, cache_ptr, cache_bf16_ptr, @@ -839,7 +878,12 @@ def _process_segment( hi_full = lo + ((hi - lo) // BLOCK_K) * BLOCK_K if IS_FP8: - n_full = (hi_full - lo) // BLOCK_K + # UNI_TILE: the partial tile is just the last iteration -- no peeled masked + # copy of the body (RESULTS.md 19). + if UNI_TILE: + n_full = (hi - lo + BLOCK_K - 1) // BLOCK_K + else: + n_full = (hi_full - lo) // BLOCK_K if n_full > 0: kn, ks, kr, vld = _gd_fp8( cache_ptr, @@ -866,6 +910,9 @@ def _process_segment( IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, + UNI_TILE, + hi, + num_rows, ) for i in range(1, n_full): kn2, ks2, kr2, vld2 = _gd_fp8( @@ -893,6 +940,9 @@ def _process_segment( IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, + UNI_TILE, + hi, + num_rows, ) m_i, l_i, acc = _qkpv_fp8( kn, @@ -924,6 +974,9 @@ def _process_segment( UNIFORM, FP8_FNUZ, HAS_INVALID, + UNI_TILE, + lo + (i - 1) * BLOCK_K, + hi, ) kn, ks, kr, vld = kn2, ks2, kr2, vld2 m_i, l_i, acc = _qkpv_fp8( @@ -956,6 +1009,9 @@ def _process_segment( UNIFORM, FP8_FNUZ, HAS_INVALID, + UNI_TILE, + lo + (n_full - 1) * BLOCK_K, + hi, ) else: for k_start in range(lo, hi_full, BLOCK_K): @@ -1007,9 +1063,13 @@ def _process_segment( IDX_BUFFER_LOAD, HAS_INVALID, FP8_FNUZ, + UNI_TILE, + hi, + num_rows, ) - if hi_full < hi: + if (not UNI_TILE) or (not IS_FP8): + if hi_full < hi: m_i, l_i, acc = _decode_tile( q_dot, cache_ptr, @@ -1130,6 +1190,7 @@ def _pa_decode_sparse( # extra still wants the CTAs. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an # empty main range and contribute an extra-only partial. + UNI_TILE: gl.constexpr, GRID_ORDER: gl.constexpr, MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. @@ -1361,6 +1422,7 @@ def _pa_decode_sparse( main_lo = gl.minimum(split_id * main_chunk, main_len) main_hi = gl.minimum(main_lo + main_chunk, main_len) m_i, l_i, acc = _process_segment( + UNI_TILE, q_dot, main_cache_ptr, main_cache_bf16_ptr, @@ -1410,6 +1472,7 @@ def _pa_decode_sparse( extra_lo = split_id * extra_chunk extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) m_i, l_i, acc = _process_segment( + UNI_TILE, q_dot, extra_cache_ptr, extra_cache_bf16_ptr, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 112c230137d..6c764ea7927 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -865,6 +865,14 @@ def _pick_splits(hb): _go = _os.environ.get("AITER_PA_DECODE_GRID_ORDER", "qsh") _ax = {"q": num_queries, "s": num_splits, "h": heads_blocks} assert sorted(_go) == ["h", "q", "s"], f"bad AITER_PA_DECODE_GRID_ORDER {_go!r}" + # AITER_PA_DECODE_UNI_TILE: run the partial last tile through the same body as the + # full ones instead of peeling a masked copy of it. Gluon inlines, so the peeled + # copy is a second ~1000-instruction gather+dequant+MFMA body whose register demand + # spills the tile loop: worth 0.900-0.907x on the buffer path and 0.943-0.972x on + # the 64-bit path at >=2 tiles/split, scaling with tiles per split. 0 restores the + # peeled tail. + uni_tile = _os.environ.get("AITER_PA_DECODE_UNI_TILE", "1") == "1" + grid = tuple(_ax[c] for c in _go) _pa_decode_sparse_gfx950[grid]( q, @@ -918,6 +926,7 @@ def _pick_splits(hb): CHUNK_AXIS=chunk_axis, PART_STORE_CACHE=_os.environ.get('AITER_PA_DECODE_PART_ST', ''), GRID_ORDER=_go, + UNI_TILE=uni_tile, MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, ASM_DEQ=asm_deq, From a5e8595d1e124d4d6930590b8ed3eb5f1b77a99e Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 00:43:27 +0000 Subject: [PATCH 22/51] [GLUON][GFX950][DSV4] Simplify the UNI_TILE slot path; three of its steps were dead Once the score mask is built natively in the MFMA layout (b585afce5), 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 --- .../gfx950/attention/pa_decode_sparse.py | 46 ++++++++----------- 1 file changed, 18 insertions(+), 28 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 0b9068b2b21..e92af5e72e8 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -356,23 +356,9 @@ def _slots( # of exec-mask branching in the kernel (20 of 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. - if UNI_TILE: - # One body for full and partial tiles. Read a REAL index -- k_pos clamped into - # the segment -- instead of masking the load, and fold the range test into - # `valid`, which _qkpv_fp8 already turns into a -inf score mask. An - # out-of-range lane then gathers a genuine row whose score is -inf, which is - # exactly how -1 sentinels are already handled on the full-tile path. - in_range = k_pos < hi - kp = gl.minimum(k_pos, gl.maximum(hi - 1, 0)) - if IDX_BUFFER_LOAD: - slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=kp) - else: - slot = gl.load(indices_ptr + seg_start + kp) - valid = in_range - if HAS_INVALID: - valid = valid & (slot >= 0) & (slot < num_rows) - slot = gl.where(valid, slot, 0) - elif MASKED: + if MASKED: + # Legacy peeled tail (UNI_TILE=0, and the bf16 path). Predicating the read is + # what makes this expensive: a masked gl.load branches on exec. in_range = k_pos < hi if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load( @@ -383,13 +369,21 @@ def _slots( valid = in_range & (slot >= 0) & (slot < num_rows) slot = gl.where(valid, slot, 0) else: + # UNI_TILE lets the partial tile ride this same body: clamp the index-list read + # into the segment rather than predicating it, so every lane reads a genuine + # in-segment index, and add the range test to `valid` -- which becomes a -inf + # score mask, exactly how -1 sentinels are already handled here. hi >= 1 + # whenever this runs (_gd_fp8 is guarded by n_full > 0, so hi > lo >= 0). + off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos if IDX_BUFFER_LOAD: - slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=k_pos) + slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off) else: - slot = gl.load(indices_ptr + seg_start + k_pos) - valid = slot >= 0 # -1 sentinels: clamp in-bounds, mask score below + slot = gl.load(indices_ptr + seg_start + off) + valid = (k_pos < hi) if UNI_TILE else (slot >= 0) if HAS_INVALID: - slot = gl.where(valid, slot, 0) + if UNI_TILE: + valid = valid & (slot >= 0) + slot = gl.where(valid, slot, 0) # -1 sentinels: clamp, mask score below return (slot // BLOCK_SIZE).to(gl.int32), (slot % BLOCK_SIZE).to(gl.int32), valid @@ -652,7 +646,6 @@ def _gd_fp8( FP8_FNUZ: gl.constexpr, UNI_TILE: gl.constexpr = False, seg_hi=0, - num_rows=0, ): """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues an iteration early. @@ -668,8 +661,8 @@ def _gd_fp8( indices_ptr, seg_start, k_start + k_rng_slot, - seg_hi, # hi/num_rows: unused unless UNI_TILE - num_rows, + seg_hi, # hi: unused unless UNI_TILE; num_rows: MASKED only + 0, BLOCK_SIZE, False, HAS_INVALID, @@ -719,7 +712,7 @@ def _gd_fp8( seg_start, k_start + k_rng_rope, seg_hi, - num_rows, + 0, BLOCK_SIZE, False, HAS_INVALID, @@ -912,7 +905,6 @@ def _process_segment( FP8_FNUZ, UNI_TILE, hi, - num_rows, ) for i in range(1, n_full): kn2, ks2, kr2, vld2 = _gd_fp8( @@ -942,7 +934,6 @@ def _process_segment( FP8_FNUZ, UNI_TILE, hi, - num_rows, ) m_i, l_i, acc = _qkpv_fp8( kn, @@ -1065,7 +1056,6 @@ def _process_segment( FP8_FNUZ, UNI_TILE, hi, - num_rows, ) if (not UNI_TILE) or (not IS_FP8): From 645a30b0edbb3c58a6c5a02f9b8211a5cc6afd75 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 00:53:38 +0000 Subject: [PATCH 23/51] [GLUON][GFX950][DSV4] Fold qk_scale into the row max so the exp2 argument 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 --- .../gfx950/attention/pa_decode_sparse.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index e92af5e72e8..b7b53caaa76 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -798,11 +798,25 @@ def _qkpv_fp8( :, None ] S = gl.where(col_mask, S, neg_inf) - S = S * qk_scale # see _decode_tile: m_i is carried in base-2 exponent space - m_block = _rmax(S, 1) + # m_i is carried in base-2 exponent space (see _decode_tile), hence qk_scale. + # max commutes with a positive scale (qk_scale = scale * log2(e); any real attention + # scale is > 0), so scale the row max -- one value per lane -- rather than every + # element of S. That leaves `S * qk_scale - m_new` as the only use of the scaled S, + # which lowers to v_fma_f32 (16 of them) where the split form needed a separate + # multiply over the tile. -inf * positive = -inf, so masked columns stay masked, and + # the row max is the same product bit for bit; only the fused subtract differs, one + # rounding instead of two -- measured <= 0.23 of a bf16 output ulp. + # + # PERFORMANCE-NEUTRAL: 0.989-1.006x across C=32/64/128 x top-k 64/272/1024 and the + # 64-bit path, inside the control arm's own 0.996-1.003x. S is [BLOCK_M, BLOCK_K], + # only 4 f32 per lane per tile, so nothing here can move a 10 us kernel. Kept + # because it is strictly fewer instructions and one fewer rounding. The multiply + # that *would* matter is `acc * alpha_pv` (128 f32 per lane) and MFMA has no + # accumulator-scale operand, so it cannot fuse. + m_block = _rmax(S, 1) * qk_scale m_new = _max2(m_i, m_block) m_new = gl.where(m_new > neg_inf, m_new, 0.0) - p = gl.exp2(S - m_new[:, None]) + p = gl.exp2(S * qk_scale - m_new[:, None]) alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(v_layout) From a1a1a2208448c1e25884580430b6b694248b5658 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 03:13:35 +0000 Subject: [PATCH 24/51] [GLUON][GFX950][DSV4] Restructure sparse decode around Cfg/Fmt/Seg aggregates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../gfx950/attention/pa_decode_sparse.py | 1866 ++++++++++------- .../ops/triton/attention/pa_decode_sparse.py | 20 +- 2 files changed, 1088 insertions(+), 798 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index b7b53caaa76..dad942a4b9a 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1,27 +1,69 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Gluon (gfx950) DeepSeek-V4 sparse-MLA decode. Adapted from the vLLM DSv4 sparse -attention kernels: -https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py - -K == V, so each tile is gathered once into one [BLOCK_K, HEAD] bf16 LDS buffer and -read permuted for QK, direct for PV. fp8 dequants to bf16 on the way in. One kernel -serves three KV formats via the UNIFORM constexpr on the fp8 gather: - - UNIFORM=False packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + RoPE bf16 - UNIFORM=True uniform pool: whole head fp8 + a separate fp32 kv_scales - IS_FP8=False bf16 pool - -Two-loop (SWA + top-k) or a single merged segment; 2D and 3D (split-K + reduce) -share one kernel. Launchers: aiter/ops/triton/attention/pa_decode_sparse.py. +"""Gluon (gfx950) sparse-MLA decode for DeepSeek-V4 *and* GLM-5 (MLA) geometries. + +Restructuring of the previous DSv4-only kernel (last DSv4-only head: +``645a30b0e``); for the existing DSv4/uniform configurations it compiles to +byte-identical .amdgcn (gate: ``compile_check.py`` in the port workspace, +``/app/scripts/glm_gluon_port_8_19``, which keeps the frozen reference). +Two changes: + +1. **Aggregate config objects** (the ``unified_attention_2d.py`` pattern): the + ~40-entry constexpr parameter lists threaded through every helper are replaced by + three aggregates — + + Cfg compile-time geometry, layouts, and behavior knobs (one per launch) + Fmt compile-time cache-format description (one per segment) + Seg runtime segment context: pointers + bounds (one per segment) + + Helpers take ``(cfg, seg, ...)``; the flat constexpr list exists only at the + kernel entry, where the aggregates are built. + +2. **Two geometries under one constexpr, ``ROPE_SEPARATE``**: + + ROPE_SEPARATE=False (DSv4): rope lives *inside* the 512-wide row. One LDS + plane ``kv_smem [BLOCK_K, KV_DIM]``, one QK MFMA chain over KV_DIM, + V = the same plane. Bit-for-bit today's kernel (compile_check.py + diffs the .amdgcn against the aiter reference to prove it). + ROPE_SEPARATE=True (GLM-5): rope is *appended* (512 latent + 64 rope = 576 + QK; V = the 512 latent only). Two LDS planes, both pow-2: + ``kv_smem [BLOCK_K, KV_DIM]`` + ``rope_smem [BLOCK_K, ROPE_DIM]``. + QK = mfma(q_rope, k_rope, mfma(q_nope, k_nope, 0)) — MFMA accumulates + natively so the 576-wide contraction is two chained dots. The PV dot + and the accumulator are unchanged: V is always exactly plane 0. + +Cache formats (``Fmt.KIND``), all sharing one gather/stage/dot pipeline: + + KIND row layout (per token) scales rope + "bf16" [KV_DIM (+ROPE_DIM if separate)] bf16 — in-row + "dsv4" 448 fp8 | 64 bf16 rope (576 B), 8 B E8M0/64 E8M0, block in-plane0 + trailer per 64-token block (aiter fp8_ds_mla-packed DSv4) + "uniform" [KV_DIM] fp8 + separate [tokens, NG] f32 f32 per-64 in-row + "tensor" [KV_DIM (+ROPE_DIM if separate)] fp8, one scalar fp8 tail + per-tensor f32 scale (GLM-5's production asm format; also the + DSv4 per-tensor experiment) + "dsmla" 512 fp8 | 4 f32 | 64 bf16 rope (656 B) f32 per-128 bf16 tail + (vLLM CacheDType "fp8_ds_mla"; requires ROPE_SEPARATE) + +Per-tensor scale handling ("tensor") never touches the tile loop: the K-side scale +folds into the segment's qk_scale (max/exp2 commute with a positive scale) and the +V-side scale multiplies ``p`` — 4 f32 per lane — right before the PV dot, so the +staged tile is a raw fp8→bf16 convert with no multiplies at all. + +Two-loop (SWA + top-k) or a single segment; 2D and 3D (split-K + reduce) share one +kernel. Launchers: ``aiter/ops/triton/attention/pa_decode_sparse.py`` (DSv4 / +uniform-pool, geometry unchanged) and +``aiter/ops/triton/attention/sparse_mla_decode.py`` (GLM-5 / separated-rope MLA). """ from triton.experimental import gluon from triton.experimental.gluon import language as gl from triton.language.core import PropagateNan +from triton.language.core import _aggregate as aggregate from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr +from aiter.ops.triton.utils.common_utils import strip_annotate # Triton's default max ignores NaN, which on AMD costs a v_max_f32 x, x, x # canonicalize per operand before the real compare -- 60 of the 96 v_max in this @@ -40,19 +82,18 @@ def _rmax(x, axis): return gl.reduce(x, axis, _max2) - @gluon.jit def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): """Gather rows[i] + col[j] out of a cache. - ``row`` is the per-token byte offset ([BLOCK_K]); ``col`` is the in-row offset - ([W]), which is always a small compile-time arange. Keeping them apart is what - makes the >2 GB path affordable: buffer_load carries a 32-bit offset (2 GB cap), - so a bigger cache has to gather through 64-bit addresses -- and adding the - pointer to a fully materialized [BLOCK_K, W] offset tensor makes that a 64-bit - add *per element*. Resolving one pointer per token instead costs BLOCK_K of - them, and the leftover column offset is a constant the load can fold into its - immediate field. (Same shape as the in-tree triton kernel's token_data_ptr.) + ``row`` is the per-token offset in ptr's element units ([BLOCK_K]); ``col`` is + the in-row offset ([W]), always a small compile-time arange. Keeping them apart + is what makes the >2 GB path affordable: buffer_load carries a 32-bit offset + (2 GB cap), so a bigger cache has to gather through 64-bit addresses -- and + adding the pointer to a fully materialized [BLOCK_K, W] offset tensor makes + that a 64-bit add *per element*. Resolving one pointer per token instead costs + BLOCK_K of them, and the leftover column offset is a constant the load can fold + into its immediate field. """ if USE_BUFFER_LOAD: return gl.amd.cdna4.buffer_load( @@ -81,6 +122,15 @@ def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) +@gluon.jit +def _fp8_to_bf16(x_u8, FP8_FNUZ: gl.constexpr): + # fp8 -> bf16 is exact (3 mantissa bits into 8), so the per-tensor mode can + # stage the raw code points without any scale arithmetic at all. + if FP8_FNUZ: + return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16) + return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.bfloat16) + + @gluon.jit def _scale_load( ptr, @@ -90,21 +140,25 @@ def _scale_load( gather_l: gl.constexpr, scl_l: gl.constexpr, NG: gl.constexpr, - HEAD_SIZE: gl.constexpr, + W_FULL: gl.constexpr, MASKED: gl.constexpr, + OTHER: gl.constexpr, ): - """Gather the NG per-64-group scale bytes of each token and broadcast them out - to HEAD_SIZE in registers. - - The obvious spelling -- index the full row with ``offs_full // 64`` -- builds a - [BLOCK_K, HEAD_SIZE] pointer tensor for a value that has only NG distinct - entries per row. On the buffer_load path the identical 32-bit offsets CSE away - and it costs 48 byte loads; on the 64-bit path they do not, and it costs 288. - (The in-tree triton kernel loads tensor<32x2x!tt.ptr> here, we were loading - tensor<64x512x!tt.ptr>.) Gather NG wide instead and broadcast: ``scl_l`` is - the dim-2 slice of a 3-D layout picked so that reshaping [BLOCK_K, NG, 64] back - to [BLOCK_K, HEAD_SIZE] lands exactly on ``gather_l`` -- column c = g*64 + j - maps to thread 4g + j//16 -- so the broadcast is a register rename. + """Gather the NG per-group scale entries of each token and broadcast them out + to W_FULL columns in registers. + + The obvious spelling -- index the full row with ``offs_full // GROUP`` -- builds + a [BLOCK_K, W_FULL] pointer tensor for a value that has only NG distinct entries + per row. On the buffer_load path the identical 32-bit offsets CSE away; on the + 64-bit path they do not, and the wide form costs 6x the loads. Gather NG wide + instead and broadcast: ``scl_l`` is the dim-2 slice of a 3-D layout picked so + that reshaping [BLOCK_K, NG, GROUP] back to [BLOCK_K, W_FULL] lands exactly on + ``gather_l`` -- so the broadcast is a register rename (assert_trivial proves it + at compile time; GLM5_PORT_PLAN flagged re-verifying this for NG != 8). + + Element type follows ``ptr``: u8 E8M0 bytes for the DSv4 packed cache + (OTHER=127 -> 2^0), f32 group scales for fp8_ds_mla (OTHER=0.0 -> masked + lanes dequant to 0). """ cols = gl.arange(0, NG, layout=gl.SliceLayout(0, scl_l)) rows = gl.convert_layout(row, gl.SliceLayout(1, scl_l)) @@ -115,14 +169,14 @@ def _scale_load( ptr=ptr, offsets=rows.to(gl.int32)[:, None] + cols[None, :], mask=m, - other=127, + other=OTHER, cache=".cg", ) else: sc = gl.load( (ptr + rows.to(gl.int64))[:, None] + cols[None, :], mask=m, - other=127, + other=OTHER, cache_modifier=".cg", ) else: @@ -137,9 +191,9 @@ def _scale_load( (ptr + rows.to(gl.int64))[:, None] + cols[None, :], cache_modifier=".cg", ) - wide = gl.expand_dims(sc, 2).broadcast_to([sc.shape[0], NG, HEAD_SIZE // NG]) + wide = gl.expand_dims(sc, 2).broadcast_to([sc.shape[0], NG, W_FULL // NG]) return gl.convert_layout( - wide.reshape([sc.shape[0], HEAD_SIZE]), gather_l, assert_trivial=True + wide.reshape([sc.shape[0], W_FULL]), gather_l, assert_trivial=True ) @@ -148,7 +202,7 @@ def _split2(x): """Contiguous register split along dim 1: [A, B] -> two [A, B//2]. x0 takes columns [0, B//2), x1 takes [B//2, B). Both halves come back in the - input's own layout, so with warps tiling dim 0 (see ``gather_l``) the column + input's own layout, so with warps tiling dim 0 (see ``Cfg.gather_l``) the column direction is a pure per-lane register repeat and the split is a rename -- no cross-lane traffic. ``assert_trivial`` makes a non-free split a compile error rather than a silent LDS round-trip. @@ -234,20 +288,343 @@ def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): return gl.convert_layout(both, out_l, assert_trivial=True) +# --------------------------------------------------------------------------- +# Aggregates. Compile-time state lives in Cfg (shared) and Fmt (per segment); +# runtime pointers/bounds live in Seg. Built once at the kernel entry, passed +# in place of the former ~40-parameter lists. +# --------------------------------------------------------------------------- + + +@aggregate +@strip_annotate +class Cfg: + """Compile-time geometry, layouts, and behavior knobs shared by both segments.""" + + # geometry + BLOCK_M: gl.constexpr + BLOCK_K: gl.constexpr + KV_DIM: gl.constexpr # plane-0 width = LDS tile width = V width (pow-2) + ROPE_DIM: gl.constexpr # plane-1 width when ROPE_SEPARATE, else the bf16 + # tail width inside plane 0 (DSv4 packed) + ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0 (DSv4). True: rope is + # a K-only second plane (GLM-5); QK contracts over + # KV_DIM + ROPE_DIM via two chained MFMAs. + QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) + MFMA_K: gl.constexpr + NUM_WARPS: gl.constexpr + GATHER_TW1: gl.constexpr + LDS_PAD: gl.constexpr + # behavior knobs + UNI_TILE: gl.constexpr + HAS_INVALID: gl.constexpr + HEAD_ALIGNED: gl.constexpr + IDX_BUFFER_LOAD: gl.constexpr + # operator layouts + qk_layout: gl.constexpr + pv_layout: gl.constexpr + q_layout: gl.constexpr + k_layout: gl.constexpr + p_layout: gl.constexpr + v_layout: gl.constexpr + # memory layouts + gather_l: gl.constexpr + gather_rope_l: gl.constexpr + gather16_l: gl.constexpr + slot_l: gl.constexpr + blocked_q: gl.constexpr + kv_shared: gl.constexpr + rope_shared: gl.constexpr + + @gluon.constexpr_function + def __init__( + self, + BLOCK_M, + BLOCK_K, + KV_DIM, + ROPE_DIM, + ROPE_SEPARATE, + MFMA_K, + NUM_WARPS, + GATHER_TW1, + LDS_PAD, + UNI_TILE, + HAS_INVALID, + HEAD_ALIGNED, + IDX_BUFFER_LOAD, + ): + self.BLOCK_M = gl.constexpr(BLOCK_M) + self.BLOCK_K = gl.constexpr(BLOCK_K) + self.KV_DIM = gl.constexpr(KV_DIM) + self.ROPE_DIM = gl.constexpr(ROPE_DIM) + self.ROPE_SEPARATE = gl.constexpr(ROPE_SEPARATE) + self.QK_DIM = gl.constexpr(KV_DIM + (ROPE_DIM if ROPE_SEPARATE else 0)) + self.MFMA_K = gl.constexpr(MFMA_K) + self.NUM_WARPS = gl.constexpr(NUM_WARPS) + self.GATHER_TW1 = gl.constexpr(GATHER_TW1) + self.LDS_PAD = gl.constexpr(LDS_PAD) + self.UNI_TILE = gl.constexpr(UNI_TILE) + self.HAS_INVALID = gl.constexpr(HAS_INVALID) + self.HEAD_ALIGNED = gl.constexpr(HEAD_ALIGNED) + self.IDX_BUFFER_LOAD = gl.constexpr(IDX_BUFFER_LOAD) + + self.qk_layout = gl.constexpr( + gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + ) + self.pv_layout = gl.constexpr( + gl.amd.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[1, NUM_WARPS], + ) + ) + KW = MFMA_K // 2 + self.q_layout = gl.constexpr(gl.DotOperandLayout(0, self.qk_layout, KW)) + self.k_layout = gl.constexpr(gl.DotOperandLayout(1, self.qk_layout, KW)) + self.p_layout = gl.constexpr(gl.DotOperandLayout(0, self.pv_layout, KW)) + self.v_layout = gl.constexpr(gl.DotOperandLayout(1, self.pv_layout, KW)) + + # 16 uint8 = 128-bit fp8 gather loads + GSPT = 16 + # Warps on dim 1 (columns) vs dim 0 (rows). Coalescing is identical either way + # -- a row's 128 contiguous bytes always come from 8 threads -- but the row + # index vector lives in SliceLayout(1, gather_l), so warps-on-dim-1 replicates + # it across all NUM_WARPS and every lane carries BLOCK_K/8 slots, while + # warps-on-dim-0 carries BLOCK_K/(8*NUM_WARPS). Dim 0 also makes the column + # direction a pure per-lane register repeat, so a chunked dequant can split it + # with a trivial (register-renaming) layout convert. + # GATHER_TW1 = threads spent on the head dim. Each thread loads 16 B, so + # TW1=8 requests 128 B of a token row per instruction and TW1=32 requests the + # whole 512 B row -- for a scattered gather that is one request instead of + # four quarters. The cost is that the row-index vector lives in + # SliceLayout(1, gather_l), so a wider TW1 leaves fewer dim-0 thread slots and + # every lane carries more slots (BLOCK_K*TW1/64 of them). Whichever dim keeps + # per-lane register repeats is the one the chunked dequant can split + # (CHUNK_AXIS). + self.gather_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, GSPT], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting + # them on dim 1 (as the NoPE gather does) overshoots 4x and every warp re-gathers + # the same tile. Worth ~10% on packed fp8. + self.gather_rope_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + # Same tiling as gather_l but 2-byte elements: half the per-thread run, so a lane + # still covers the same 16 bytes and the lo/hi interleave stays lane-local. + self.gather16_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, GSPT // 2], + threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + self.slot_l = gl.constexpr(gl.SliceLayout(1, self.gather_l.value)) + self.blocked_q = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[1, NUM_WARPS], + order=[1, 0], + ) + ) + # LDS pad after every row. Row pitch (512 + LDS_PAD) bf16 sets which banks a + # transposed K read (ds_read_b64_tr_b16 walks down a column) lands on: + # bank = (row * pitch_dwords) mod 32, so pitch_dwords mod 32 == 4 at PAD=8 means + # 32 lanes share 8 banks. Kept as a knob because the conflict share is high + # (46%) but a high sub-metric ratio is not proof of a bottleneck. + self.kv_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[KV_DIM, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] + ) + ) + # Plane 1 (K-only rope) when ROPE_SEPARATE; same padding rule as plane 0 + # (its K read is the same transposed ds_read pattern). Built regardless -- + # a dead constexpr field costs nothing when the plane is never allocated. + self.rope_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[ROPE_DIM, LDS_PAD]], [BLOCK_K, ROPE_DIM], [1, 0] + ) + ) + + +@aggregate +@strip_annotate +class Fmt: + """Compile-time description of one segment's cache format.""" + + KIND: gl.constexpr # "bf16" | "dsv4" | "uniform" | "tensor" | "dsmla" + IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop + BLOCK_SIZE: gl.constexpr + USE_BUFFER_LOAD: gl.constexpr + FP8_FNUZ: gl.constexpr + ASM_DEQ: gl.constexpr + NOPE_CHUNK: gl.constexpr + CHUNK_AXIS: gl.constexpr + NOPE_DIM: gl.constexpr # fp8 payload width (448 dsv4; KV_DIM elsewhere) + GROUP: gl.constexpr # scale group width (64 dsv4/uniform, 128 dsmla) + NG: gl.constexpr # scale groups per row (KV_DIM // GROUP) + NARROW_SCALE: gl.constexpr + scl_l: gl.constexpr + # packed-row addressing constants (element units named in the suffix) + TOK_U8: gl.constexpr # bytes per token row inside a block (576 / 656) + TOK_U16: gl.constexpr # the same row in bf16-view units + ROPE_U16_OFF: gl.constexpr # bf16-view offset of the rope tail in a row + SCL_TRAILER_U8: gl.constexpr # dsv4: scale bytes per token in the block trailer + TOK_F32: gl.constexpr # dsmla: f32-view stride per token row + SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the 4 group scales + TOK_EL: gl.constexpr # flat formats: cache elements per token row + + @gluon.constexpr_function + def __init__( + self, + cfg, + KIND, + BLOCK_SIZE, + USE_BUFFER_LOAD, + FP8_FNUZ, + ASM_DEQ, + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, + ): + self.KIND = gl.constexpr(KIND) + self.IS_FP8 = gl.constexpr(KIND != "bf16") + self.BLOCK_SIZE = gl.constexpr(BLOCK_SIZE) + self.USE_BUFFER_LOAD = gl.constexpr(USE_BUFFER_LOAD) + self.FP8_FNUZ = gl.constexpr(FP8_FNUZ) + self.ASM_DEQ = gl.constexpr(ASM_DEQ) + self.NOPE_CHUNK = gl.constexpr(NOPE_CHUNK) + self.CHUNK_AXIS = gl.constexpr(CHUNK_AXIS) + self.NOPE_DIM = gl.constexpr(NOPE_DIM) + + KV_DIM = cfg.KV_DIM.value + ROPE_DIM = cfg.ROPE_DIM.value + GROUP = 128 if KIND == "dsmla" else 64 + NG = KV_DIM // GROUP + self.GROUP = gl.constexpr(GROUP) + self.NG = gl.constexpr(NG) + # 3-D companion of gather_l used by _scale_load: dim 1 carries the NG scale + # groups (one thread each) and dim 2 the GROUP columns inside a group, so + # that reshaping [BLOCK_K, NG, GROUP] back to 2-D reproduces gather_l + # exactly. Legal only when the group's columns fill whole threads, i.e. + # GSPT * (TW1 // NG) is the group width; otherwise fall back to the wide + # (redundant) gather. + GSPT = 16 + TW1 = cfg.GATHER_TW1.value + NARROW_SCALE = TW1 % NG == 0 and GSPT * (TW1 // NG) == GROUP + self.NARROW_SCALE = gl.constexpr(NARROW_SCALE) + # ...and only worth it on the 64-bit path. With buffer_load the wide gather's + # identical 32-bit offsets CSE to the same few loads, so the narrow form + # only adds a layout convert per tile (+3.4% at extra=1024). Without it they + # do not CSE and the wide form costs 6x the loads. + scl_l3 = gl.BlockedLayout( + size_per_thread=[1, 1, GSPT], + threads_per_warp=[ + 64 // TW1, + NG if NARROW_SCALE else 1, + (TW1 // NG) if NARROW_SCALE else TW1, + ], + warps_per_cta=[cfg.NUM_WARPS.value, 1, 1], + order=[2, 1, 0], + ) + self.scl_l = gl.constexpr(gl.SliceLayout(2, scl_l3)) + + # Packed-row constants. dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] with an + # 8-byte-per-token E8M0 trailer after the block; dsmla row: + # [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] all inline. + if KIND == "dsv4": + TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 + elif KIND == "dsmla": + TOK_U8 = KV_DIM + 4 * NG + 2 * ROPE_DIM # 512 + 16 + 128 = 656 + else: + TOK_U8 = 0 + self.TOK_U8 = gl.constexpr(TOK_U8) + self.TOK_U16 = gl.constexpr(TOK_U8 // 2) + self.ROPE_U16_OFF = gl.constexpr( + (NOPE_DIM // 2) if KIND == "dsv4" else ((KV_DIM + 4 * NG) // 2) + ) + self.SCL_TRAILER_U8 = gl.constexpr(8) + self.TOK_F32 = gl.constexpr(TOK_U8 // 4) + self.SCL_F32_OFF = gl.constexpr(KV_DIM // 4) + # Flat formats gather in cache elements (bf16 or fp8-byte): the row is + # KV_DIM wide, plus the appended rope when the geometry separates it. + self.TOK_EL = gl.constexpr( + KV_DIM + (ROPE_DIM if cfg.ROPE_SEPARATE.value and KIND != "uniform" else 0) + ) + + +@aggregate +@strip_annotate +class Seg: + """Runtime context of one segment: cache pointers, index list, bounds. + + Pointer roles by format ("--" = unused, carries a harmless duplicate): + + KIND cache_ptr alt_ptr scl_ptr + bf16 -- bf16 cache -- + dsv4 u8 cache bf16 view of the cache -- + uniform u8 cache f32 per-64 kv_scales -- + tensor u8 cache -- f32 scalar k_scale + dsmla u8 cache bf16 view (rope tail) f32 view (group scales) + """ + + fmt: Fmt + cache_ptr: gl.tensor + alt_ptr: gl.tensor + scl_ptr: gl.tensor + indices_ptr: gl.tensor + seg_start: gl.tensor + cs0: gl.tensor + num_rows: gl.tensor + + @gluon.constexpr_function + def __init__( + self, fmt, cache_ptr, alt_ptr, scl_ptr, indices_ptr, seg_start, cs0, num_rows + ): + self.fmt = fmt + self.cache_ptr = cache_ptr + self.alt_ptr = alt_ptr + self.scl_ptr = scl_ptr + self.indices_ptr = indices_ptr + self.seg_start = seg_start + self.cs0 = cs0 + self.num_rows = num_rows + + +# --------------------------------------------------------------------------- +# Dequant-and-store into the LDS plane(s) +# --------------------------------------------------------------------------- + + @gluon.jit -def _deq_store( - x_u8, sc, kv_smem, off, - AXIS: gl.constexpr, UNIFORM: gl.constexpr, FP8_FNUZ: gl.constexpr, - ASM_DEQ: gl.constexpr, gather_l: gl.constexpr, -): +def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): """Dequant one fp8 slab and write it straight to kv_smem[:, off:off+W]. Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path lowers to - an emulated mul and doubles the loop. ``sc`` is the raw per-64-group exponent - byte (packed fp8_ds_mla, UE8M0) or an f32 scale (uniform pool). + an emulated mul and doubles the loop. ``sc`` is the raw per-group exponent + byte (dsv4 packed, UE8M0), an f32 scale (uniform pool / fp8_ds_mla), or unused + ("tensor": the per-tensor scale is folded outside the tile loop, so staging is + a bare fp8 -> bf16 convert). """ - if ASM_DEQ: - # x_u8 is the int16 view here (see _gd_fp8), so its column count is W8/2. + if fmt.ASM_DEQ: + # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. W8: gl.constexpr = x_u8.shape[1] * 2 # Adjacent fp8 columns share a scale (groups are 64 wide, so even), which makes # dropping every other broadcast column exact. @@ -255,17 +632,20 @@ def _deq_store( # split leaves a SliceLayout of a 3-D parent; the asm needs both operands in # the int16 gather's own layout. s_even = gl.convert_layout(s_even, x_u8.type.layout) - val = _deq_asm(x_u8.to(gl.int16, bitcast=True), s_even, W8, gather_l) + val = _deq_asm(x_u8.to(gl.int16, bitcast=True), s_even, W8, cfg.gather_l) if AXIS == 1: kv_smem.slice(off, W8, dim=1).store(val) else: kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) else: - if UNIFORM: - scale = sc + if fmt.KIND == "tensor": + val = _fp8_to_bf16(x_u8, fmt.FP8_FNUZ) else: - scale = gl.exp2(sc.to(gl.float32) - 127.0) - val = (_fp8_to_f32(x_u8, FP8_FNUZ) * scale).to(gl.bfloat16) + if fmt.KIND == "uniform" or fmt.KIND == "dsmla": + scale = sc + else: + scale = gl.exp2(sc.to(gl.float32) - 127.0) + val = (_fp8_to_f32(x_u8, fmt.FP8_FNUZ) * scale).to(gl.bfloat16) if AXIS == 1: kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) else: @@ -273,17 +653,7 @@ def _deq_store( @gluon.jit -def _deq_store_tile( - x_u8, - sc, - kv_smem, - NOPE_CHUNK: gl.constexpr, - AXIS: gl.constexpr, - ASM_DEQ: gl.constexpr, - gather_l: gl.constexpr, - UNIFORM: gl.constexpr, - FP8_FNUZ: gl.constexpr, -): +def _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt): """Dequant a whole gathered fp8 tile into kv_smem in NOPE_CHUNK-sized pieces along CHUNK_AXIS (0 = rows, 1 = columns). @@ -293,16 +663,18 @@ def _deq_store_tile( piece c's ds_writes and die -- so only a 1/pieces fraction of it is ever live. The splits themselves are register renames (see ``_split2``). """ + AXIS: gl.constexpr = fmt.CHUNK_AXIS + NOPE_CHUNK: gl.constexpr = fmt.NOPE_CHUNK W: gl.constexpr = x_u8.shape[1] if AXIS == 1 else x_u8.shape[0] if NOPE_CHUNK >= W: - _deq_store(x_u8, sc, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x_u8, sc, kv_smem, 0, cfg, fmt, AXIS) else: x0, x1 = _split_ax(x_u8, AXIS) s0, s1 = _split_ax(sc, AXIS) W2: gl.constexpr = W // 2 if NOPE_CHUNK >= W2: - _deq_store(x0, s0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(x1, s1, kv_smem, W2, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x0, s0, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(x1, s1, kv_smem, W2, cfg, fmt, AXIS) else: x00, x01 = _split_ax(x0, AXIS) s00, s01 = _split_ax(s0, AXIS) @@ -310,41 +682,38 @@ def _deq_store_tile( s10, s11 = _split_ax(s1, AXIS) W4: gl.constexpr = W // 4 if NOPE_CHUNK >= W4: - _deq_store(x00, s00, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(x01, s01, kv_smem, W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(x10, s10, kv_smem, 2 * W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(x11, s11, kv_smem, 3 * W4, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(x00, s00, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(x01, s01, kv_smem, W4, cfg, fmt, AXIS) + _deq_store(x10, s10, kv_smem, 2 * W4, cfg, fmt, AXIS) + _deq_store(x11, s11, kv_smem, 3 * W4, cfg, fmt, AXIS) else: W8: gl.constexpr = W // 8 y0, y1 = _split_ax(x00, AXIS) t0, t1 = _split_ax(s00, AXIS) - _deq_store(y0, t0, kv_smem, 0, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(y1, t1, kv_smem, W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y0, t0, kv_smem, 0, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, W8, cfg, fmt, AXIS) y0, y1 = _split_ax(x01, AXIS) t0, t1 = _split_ax(s01, AXIS) - _deq_store(y0, t0, kv_smem, 2 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(y1, t1, kv_smem, 3 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y0, t0, kv_smem, 2 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 3 * W8, cfg, fmt, AXIS) y0, y1 = _split_ax(x10, AXIS) t0, t1 = _split_ax(s10, AXIS) - _deq_store(y0, t0, kv_smem, 4 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(y1, t1, kv_smem, 5 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y0, t0, kv_smem, 4 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 5 * W8, cfg, fmt, AXIS) y0, y1 = _split_ax(x11, AXIS) t0, t1 = _split_ax(s11, AXIS) - _deq_store(y0, t0, kv_smem, 6 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) - _deq_store(y1, t1, kv_smem, 7 * W8, AXIS, UNIFORM, FP8_FNUZ, ASM_DEQ, gather_l) + _deq_store(y0, t0, kv_smem, 6 * W8, cfg, fmt, AXIS) + _deq_store(y1, t1, kv_smem, 7 * W8, cfg, fmt, AXIS) @gluon.jit def _slots( - indices_ptr, - seg_start, + cfg, + seg, k_pos, hi, num_rows, - BLOCK_SIZE: gl.constexpr, MASKED: gl.constexpr, - HAS_INVALID: gl.constexpr, - IDX_BUFFER_LOAD: gl.constexpr, UNI_TILE: gl.constexpr = False, ): # Returns in whatever layout k_pos carries. Called once per gather layout: NoPE @@ -356,6 +725,11 @@ def _slots( # of exec-mask branching in the kernel (20 of 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. + indices_ptr = seg.indices_ptr + seg_start = seg.seg_start + BLOCK_SIZE: gl.constexpr = seg.fmt.BLOCK_SIZE + HAS_INVALID: gl.constexpr = cfg.HAS_INVALID + IDX_BUFFER_LOAD: gl.constexpr = cfg.IDX_BUFFER_LOAD if MASKED: # Legacy peeled tail (UNI_TILE=0, and the bf16 path). Predicating the read is # what makes this expensive: a masked gl.load branches on exec. @@ -373,7 +747,7 @@ def _slots( # into the segment rather than predicating it, so every lane reads a genuine # in-segment index, and add the range test to `valid` -- which becomes a -inf # score mask, exactly how -1 sentinels are already handled here. hi >= 1 - # whenever this runs (_gd_fp8 is guarded by n_full > 0, so hi > lo >= 0). + # whenever this runs (_gather_full is guarded by n_full > 0, so hi > lo >= 0). off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off) @@ -387,211 +761,492 @@ def _slots( return (slot // BLOCK_SIZE).to(gl.int32), (slot % BLOCK_SIZE).to(gl.int32), valid +@gluon.jit +def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): + """QK scores for one tile. ROPE_SEPARATE=False is a single MFMA chain over + plane 0; ROPE_SEPARATE=True chains a second MFMA over the rope plane (MFMA + accumulates natively, so the KV_DIM + ROPE_DIM contraction is two dots).""" + k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] + S = gl.amd.cdna4.mfma( + q_dot, + k, + gl.zeros([cfg.BLOCK_M, cfg.BLOCK_K], gl.float32, layout=cfg.qk_layout), + ) + if cfg.ROPE_SEPARATE: + k_rope = rope_smem.permute([1, 0]).load(cfg.k_layout) # [ROPE_DIM, BLOCK_K] + S = gl.amd.cdna4.mfma(q_rope_dot, k_rope, S) + return S + + +# --------------------------------------------------------------------------- +# The prefetched fp8 pipeline: _gather_full issues tile N+1's loads while +# _qkpv stages + dots tile N. +# --------------------------------------------------------------------------- + + +@gluon.jit +def _gather_full( + cfg, + seg, + k_start, + seg_hi, + offs_full, + offs_full16, + offs_rope, + k_rng_slot, + k_rng_rope, +): + """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues + an iteration early. + + The prefetch is carried across the MFMA in **raw fp8**, not dequantized bf16: + the loop keeps two tiles in flight, and a [BLOCK_K, KV_DIM] tile is 32 + VGPRs/lane as u8 but 64 as bf16, so dequantizing here would burn an extra 64 + VGPRs on loop-carried state alone. The consumer dequants (chunked) instead. + + Returns (x, sc, k_rope, valid); unused slots carry a harmless duplicate that + DCE removes ("tensor" has no scale vector, "uniform" no rope side-channel). + """ + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) + bg, pg, valid = _slots( + cfg, + seg, + k_start + k_rng_slot, + seg_hi, # hi: unused unless UNI_TILE + 0, + False, + cfg.UNI_TILE, + ) + if fmt.KIND == "uniform": + NGRP: gl.constexpr = cfg.KV_DIM // 64 + x_u8 = _cache_load( + seg.cache_ptr, bg * cs0 + pg * cfg.KV_DIM, offs_full, fmt.USE_BUFFER_LOAD + ) + sc = _cache_load( + seg.alt_ptr, bg * NGRP, offs_full // 64, fmt.USE_BUFFER_LOAD + ) + k_rope = x_u8 # unused for "uniform" (rope slice-store skipped) -> DCE'd + elif fmt.KIND == "tensor": + # Flat fp8 rows, one per-tensor scale for the whole cache. No scale work + # here at all: the K-side scale is folded into this segment's qk_scale and + # the V-side scale into p (see _qkpv), so staging is a bare convert. + x_u8 = _cache_load( + seg.cache_ptr, bg * cs0 + pg * fmt.TOK_EL, offs_full, fmt.USE_BUFFER_LOAD + ) + sc = x_u8 # no scale vector -> DCE'd + if cfg.ROPE_SEPARATE: + # K-only rope tail: fp8 columns [KV_DIM, KV_DIM + ROPE_DIM) of the row. + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.cache_ptr, + bgr * cs0 + pgr * fmt.TOK_EL + cfg.KV_DIM, + offs_rope, + fmt.USE_BUFFER_LOAD, + ) + else: + k_rope = x_u8 # rope lives inside plane 0 -> DCE'd + elif fmt.KIND == "dsmla": + # vLLM fp8_ds_mla row: [KV_DIM fp8 | NG f32 scales | ROPE_DIM bf16]. + nope_row = bg * cs0 + pg * fmt.TOK_U8 + scl_row = bg * (cs0 // 4) + pg * fmt.TOK_F32 + fmt.SCL_F32_OFF + # f32 group scales first, bulk fp8 after (same vmcnt-FIFO argument as dsv4). + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + sc = _scale_load( + seg.scl_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, + fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 0.0, + ) + else: + sc = _cache_load( + seg.scl_ptr, scl_row, offs_full // fmt.GROUP, fmt.USE_BUFFER_LOAD + ) + x_u8 = _cache_load(seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD) + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.alt_ptr, + bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, + offs_rope, + fmt.USE_BUFFER_LOAD, + ) + else: # "dsv4" + nope_row = bg * cs0 + pg * fmt.TOK_U8 + scl_row = bg * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pg * fmt.SCL_TRAILER_U8 + # UE8M0 scales first, bulk fp8 after. vmcnt is one in-order FIFO, so a + # wait can only name "at most N outstanding", never a specific load: + # issuing the scales last would make the first dequant piece wait behind + # every data load as well. Issued first, the scales are covered by the + # wait that the first piece's own data needs anyway. + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: + sc = _scale_load( + seg.cache_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, + fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 127, + ) + else: + sc = _cache_load( + seg.cache_ptr, scl_row, offs_full // 64, fmt.USE_BUFFER_LOAD + ) + if fmt.ASM_DEQ: + # 2-byte elements so the asm gets <2 x i16> = 4 packed fp8 in one VGPR out + # of a single dword load. Same bytes, same 16-B per-lane run, half the + # columns; alt_ptr is already the 2-byte view of this cache and the + # RoPE gather below uses the same cs0//2 / TOK_U16 addressing. + # The row vector comes from _slots in SliceLayout(1, gather_l); the i16 + # columns live in gather16_l, and a broadcast needs one parent layout. The + # two layouts share their dim-0 tiling, so this convert is a rename. + row16 = gl.convert_layout( + nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) + ) + x_u8 = _cache_load( + seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD + ) + else: + x_u8 = _cache_load(seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD) + bgr, pgr, _ = _slots( + cfg, + seg, + k_start + k_rng_rope, + seg_hi, + 0, + False, + cfg.UNI_TILE, + ) + k_rope = _cache_load( + seg.alt_ptr, + bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, + offs_rope, + fmt.USE_BUFFER_LOAD, + ) + return x_u8, sc, k_rope, valid + + +@gluon.jit +def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): + """Write one prefetched tile into the LDS plane(s). Plane 0 is always the full + KV_DIM-wide dequant ("dsv4" gathers KV_DIM bytes too: the last 64 are the bf16 + rope read as garbage fp8, dequanted, and then overwritten by the real rope + slice-store below -- that keeps the gather pow-2 wide).""" + fmt = seg.fmt + _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) + if fmt.KIND == "dsv4": + kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) + elif fmt.KIND == "dsmla": + rope_smem.store(k_rope) + elif fmt.KIND == "tensor": + if cfg.ROPE_SEPARATE: + rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) + # "uniform": the whole head is one fp8 tile; nothing else to store. + + +@gluon.jit +def _qkpv( + cfg, + seg, + x_u8, + sc, + k_rope, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start=0, + seg_hi=0, +): + """Dequant a prefetched fp8 tile into LDS, then QK -> softmax -> PV. When + HAS_INVALID, mask the columns of -1-sentinel slots (``valid`` from + _gather_full) to -inf.""" + neg_inf = float("-inf") + _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem) + S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) + # UNI_TILE folds the tile's range test into `valid`, so the score mask is no + # longer only about -1 sentinels: it is what makes the partial last tile correct. + COL_MASK: gl.constexpr = cfg.HAS_INVALID or cfg.UNI_TILE + NEED_MASK: gl.constexpr = COL_MASK or (not cfg.HEAD_ALIGNED) + if NEED_MASK: + if COL_MASK: + if cfg.UNI_TILE and not cfg.HAS_INVALID: + # Build the range mask directly in the MFMA layout. Converting the + # slot-layout `valid` vector instead costs a cross-lane layout change + # per tile, which at 16 tiles ate the whole saving. + col_mask = ( + k_start + + gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(0, cfg.qk_layout)) + < seg_hi + )[None, :] + else: + col_mask = gl.convert_layout( + valid, gl.SliceLayout(0, cfg.qk_layout) + )[None, :] + if not cfg.HEAD_ALIGNED: + col_mask = ( + gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + & col_mask + ) + else: + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] + S = gl.where(col_mask, S, neg_inf) + # m_i is carried in base-2 exponent space, hence qk_scale. max commutes with a + # positive scale (qk_scale = scale * log2(e) [* k_scale for "tensor"]; any real + # attention scale is > 0), so scale the row max -- one value per lane -- rather + # than every element of S. That leaves `S * qk_scale - m_new` as the only use of + # the scaled S, which lowers to v_fma_f32 where the split form needed a separate + # multiply over the tile. -inf * positive = -inf, so masked columns stay masked, + # and the row max is the same product bit for bit; only the fused subtract + # differs, one rounding instead of two -- measured <= 0.23 of a bf16 output ulp. + m_block = _rmax(S, 1) * qk_scale + m_new = _max2(m_i, m_block) + m_new = gl.where(m_new > neg_inf, m_new, 0.0) + p = gl.exp2(S * qk_scale - m_new[:, None]) + alpha = gl.exp2(m_i - m_new) + l_new = l_i * alpha + gl.sum(p, axis=1) + v = kv_smem.load(cfg.v_layout) + # "tensor": V was staged as raw fp8 code points, so apply the per-tensor scale + # on the small side -- p is [BLOCK_M, BLOCK_K], 4 f32 per lane, vs 32 for the + # accumulator. l_new stays scale-free (it normalizes p, not p*V), so + # out = sum(p * v_scale * V_raw) / l is exactly sum(p * (v_scale * V_raw)) / l. + if seg.fmt.KIND == "tensor": + p = p * v_scale + p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.pv_layout)) + acc = acc * alpha_pv[:, None] + acc = gl.amd.cdna4.mfma(p_dot, v, acc) + return m_new, l_new, acc + + +# --------------------------------------------------------------------------- +# The non-prefetched path: bf16 segments, and the peeled masked tail when +# UNI_TILE is off (dsv4/uniform only; the new fp8 formats require UNI_TILE). +# --------------------------------------------------------------------------- + + @gluon.jit def _decode_tile( + cfg, + seg, q_dot, - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, + q_rope_dot, k_start, hi, - cs0, - num_rows, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, + rope_smem, offs_full, offs_full16, offs_rope, k_rng_slot, k_rng_rope, - qk_layout: gl.constexpr, - pv_layout: gl.constexpr, - k_layout: gl.constexpr, - v_layout: gl.constexpr, - p_layout: gl.constexpr, - gather_l: gl.constexpr, - gather_rope_l: gl.constexpr, - gather16_l: gl.constexpr, - scl_l: gl.constexpr, - NARROW_SCALE: gl.constexpr, - IS_FP8: gl.constexpr, - BLOCK_SIZE: gl.constexpr, - NOPE_DIM: gl.constexpr, - ROPE_DIM: gl.constexpr, - HEAD_SIZE: gl.constexpr, - BLOCK_M: gl.constexpr, - BLOCK_K: gl.constexpr, - NOPE_CHUNK: gl.constexpr, - CHUNK_AXIS: gl.constexpr, - ASM_DEQ: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, MASKED: gl.constexpr, - UNIFORM: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, - IDX_BUFFER_LOAD: gl.constexpr, - HAS_INVALID: gl.constexpr, - FP8_FNUZ: gl.constexpr, ): """One KV tile -> online-softmax update. MASKED=False (peeled full tiles) drops the in-range / gather / score masking; MASKED=True (the tail) keeps it. When HAS_INVALID, full tiles also clamp -1 sentinels in-bounds for the gather and mask their scores to -inf (matching the tail's slot-validity handling).""" neg_inf = float("-inf") - if not USE_BUFFER_LOAD: + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) block_idx, pos, valid1d = _slots( - indices_ptr, - seg_start, + cfg, + seg, k_start + k_rng_slot, hi, - num_rows, - BLOCK_SIZE, + seg.num_rows, MASKED, - HAS_INVALID, - IDX_BUFFER_LOAD, ) - block_idx_g = gl.convert_layout(block_idx, gl.SliceLayout(1, gather_l)) - pos_g = gl.convert_layout(pos, gl.SliceLayout(1, gather_l)) + block_idx_g = gl.convert_layout(block_idx, gl.SliceLayout(1, cfg.gather_l)) + pos_g = gl.convert_layout(pos, gl.SliceLayout(1, cfg.gather_l)) if MASKED: - valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, gather_l)) + valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather_l)) - if IS_FP8 and UNIFORM: + if fmt.KIND == "uniform": # uniform pool: one fp8 gather over the whole head + separate fp32 scales. - NGRP: gl.constexpr = HEAD_SIZE // 64 - kv_row = block_idx_g * cs0 + pos_g * HEAD_SIZE + NGRP: gl.constexpr = cfg.KV_DIM // 64 + kv_row = block_idx_g * cs0 + pos_g * cfg.KV_DIM scl_row = block_idx_g * NGRP scl_col = offs_full // 64 if MASKED: x_u8 = _cache_load( - cache_ptr, kv_row, offs_full, USE_BUFFER_LOAD, + seg.cache_ptr, kv_row, offs_full, fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0, ) sc = _cache_load( - cache_bf16_ptr, + seg.alt_ptr, scl_row, scl_col, - USE_BUFFER_LOAD, + fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, ) # uniform pool: fp32 scales, left wide (see _scale_load) else: - x_u8 = _cache_load(cache_ptr, kv_row, offs_full, USE_BUFFER_LOAD) - sc = _cache_load(cache_bf16_ptr, scl_row, scl_col, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, True, FP8_FNUZ) - elif IS_FP8: + x_u8 = _cache_load(seg.cache_ptr, kv_row, offs_full, fmt.USE_BUFFER_LOAD) + sc = _cache_load(seg.alt_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) + _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) + elif fmt.KIND == "dsv4": # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. - nope_row = block_idx_g * cs0 + pos_g * 576 - scl_row = block_idx_g * cs0 + BLOCK_SIZE * 576 + pos_g * 8 + nope_row = block_idx_g * cs0 + pos_g * fmt.TOK_U8 + scl_row = ( + block_idx_g * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pos_g * fmt.SCL_TRAILER_U8 + ) scl_col = offs_full // 64 - if MASKED: # scales first: see _gd_fp8 - if NARROW_SCALE and not USE_BUFFER_LOAD: + if MASKED: # scales first: see _gather_full + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: exps = _scale_load( - cache_ptr, scl_row, valid1d, USE_BUFFER_LOAD, gather_l, scl_l, - HEAD_SIZE // 64, HEAD_SIZE, True, + seg.cache_ptr, scl_row, valid1d, fmt.USE_BUFFER_LOAD, + cfg.gather_l, fmt.scl_l, fmt.NG, cfg.KV_DIM, True, 127, ) else: exps = _cache_load( - cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD, + seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=127, ) - if ASM_DEQ: + if fmt.ASM_DEQ: # Same byte as nope_row, addressed in 2-byte units: one shift on a # live vector, and block_idx_g * cs0 stays common with the scale - # gather. Exact -- cs0 and 576 are both even. + # gather. Exact -- cs0 and TOK_U8 are both even. row16 = gl.convert_layout( - nope_row >> 1, gl.SliceLayout(1, gather16_l) + nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) x_u8 = _cache_load( - cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD, + seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD, mask=gl.convert_layout( - valid1d, gl.SliceLayout(1, gather16_l) + valid1d, gl.SliceLayout(1, cfg.gather16_l) )[:, None], other=0.0, ) else: x_u8 = _cache_load( - cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD, + seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0, ) else: - if NARROW_SCALE and not USE_BUFFER_LOAD: + if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: exps = _scale_load( - cache_ptr, scl_row, scl_row, USE_BUFFER_LOAD, gather_l, scl_l, - HEAD_SIZE // 64, HEAD_SIZE, False, + seg.cache_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, + cfg.gather_l, fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 127, ) else: - exps = _cache_load(cache_ptr, scl_row, scl_col, USE_BUFFER_LOAD) - if ASM_DEQ: + exps = _cache_load(seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) + if fmt.ASM_DEQ: # Same byte as nope_row, addressed in 2-byte units: one shift on a # live vector, and block_idx_g * cs0 stays common with the scale - # gather. Exact -- cs0 and 576 are both even. + # gather. Exact -- cs0 and TOK_U8 are both even. row16 = gl.convert_layout( - nope_row >> 1, gl.SliceLayout(1, gather16_l) + nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) x_u8 = _cache_load( - cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD + seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD ) else: - x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) - _deq_store_tile(x_u8, exps, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, False, FP8_FNUZ) + x_u8 = _cache_load( + seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD + ) + _deq_store_tile(x_u8, exps, kv_smem, cfg, fmt) block_idx_gr, pos_gr, valid_gr = _slots( - indices_ptr, - seg_start, + cfg, + seg, k_start + k_rng_rope, hi, - num_rows, - BLOCK_SIZE, + seg.num_rows, MASKED, - HAS_INVALID, - IDX_BUFFER_LOAD, ) - rope_row = block_idx_gr * (cs0 // 2) + pos_gr * 288 + 224 + rope_row = block_idx_gr * (cs0 // 2) + pos_gr * fmt.TOK_U16 + fmt.ROPE_U16_OFF if MASKED: k_rope = _cache_load( - cache_bf16_ptr, + seg.alt_ptr, rope_row, offs_rope, - USE_BUFFER_LOAD, + fmt.USE_BUFFER_LOAD, mask=valid_gr[:, None], other=0.0, ) else: - k_rope = _cache_load(cache_bf16_ptr, rope_row, offs_rope, USE_BUFFER_LOAD) - kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) - else: - kv_row2 = block_idx_g * cs0 + pos_g * HEAD_SIZE + k_rope = _cache_load(seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD) + kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) + else: # "bf16" (the new fp8 formats require UNI_TILE, so they never come here) + kv_row2 = block_idx_g * cs0 + pos_g * fmt.TOK_EL if MASKED: kv = _cache_load( - cache_bf16_ptr, kv_row2, offs_full, USE_BUFFER_LOAD, + seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, ) else: - kv = _cache_load(cache_bf16_ptr, kv_row2, offs_full, USE_BUFFER_LOAD) + kv = _cache_load(seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD) kv_smem.store(kv) + if cfg.ROPE_SEPARATE: + # Appended bf16 rope columns [KV_DIM, KV_DIM + ROPE_DIM) -> plane 1. + block_idx_gr, pos_gr, valid_gr = _slots( + cfg, + seg, + k_start + k_rng_rope, + hi, + seg.num_rows, + MASKED, + ) + rope_row = block_idx_gr * cs0 + pos_gr * fmt.TOK_EL + cfg.KV_DIM + if MASKED: + k_rope = _cache_load( + seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD, + mask=valid_gr[:, None], other=0.0, + ) + else: + k_rope = _cache_load( + seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD + ) + rope_smem.store(k_rope) - k = kv_smem.permute([1, 0]).load(k_layout) # [HEAD_SIZE, BLOCK_K] - S = gl.amd.cdna4.mfma( - q_dot, k, gl.zeros([BLOCK_M, BLOCK_K], gl.float32, layout=qk_layout) - ) + S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) # exp2 softmax: qk_scale folds in log2(e) so we hit the HW exp2 directly. # Running max stays in raw-score space; masked cols (-inf) give exp2=0. - COL_VALID: gl.constexpr = MASKED or HAS_INVALID # valid1d defined in both cases - NEED_MASK: gl.constexpr = COL_VALID or (not HEAD_ALIGNED) + COL_VALID: gl.constexpr = MASKED or cfg.HAS_INVALID # valid1d defined in both + NEED_MASK: gl.constexpr = COL_VALID or (not cfg.HEAD_ALIGNED) if NEED_MASK: if COL_VALID: - col_mask = gl.convert_layout(valid1d, gl.SliceLayout(0, qk_layout))[None, :] - if not HEAD_ALIGNED: + col_mask = gl.convert_layout( + valid1d, gl.SliceLayout(0, cfg.qk_layout) + )[None, :] + if not cfg.HEAD_ALIGNED: col_mask = ( - gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] + gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ + :, None + ] & col_mask ) else: # not HEAD_ALIGNED and no col invalidity -> head mask only - col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[ + col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ :, None ] S = gl.where(col_mask, S, neg_inf) @@ -610,218 +1265,11 @@ def _decode_tile( alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) - v = kv_smem.load(v_layout) # [BLOCK_K, HEAD_SIZE] - p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) - alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) - acc = acc * alpha_pv[:, None] - acc = gl.amd.cdna4.mfma(p_dot, v, acc) - return m_new, l_new, acc - - -@gluon.jit -def _gd_fp8( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, - k_start, - cs0, - offs_full, - offs_full16, - offs_rope, - k_rng_slot, - k_rng_rope, - gather_l: gl.constexpr, - gather_rope_l: gl.constexpr, - gather16_l: gl.constexpr, - scl_l: gl.constexpr, - NARROW_SCALE: gl.constexpr, - ASM_DEQ: gl.constexpr, - BLOCK_SIZE: gl.constexpr, - HEAD_SIZE: gl.constexpr, - UNIFORM: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, - IDX_BUFFER_LOAD: gl.constexpr, - HAS_INVALID: gl.constexpr, - FP8_FNUZ: gl.constexpr, - UNI_TILE: gl.constexpr = False, - seg_hi=0, -): - """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues - an iteration early. - - The prefetch is carried across the MFMA in **raw fp8**, not dequantized bf16: - the loop keeps two tiles in flight, and a [BLOCK_K, HEAD_SIZE] tile is 32 - VGPRs/lane as u8 but 64 as bf16, so dequantizing here would burn an extra 64 - VGPRs on loop-carried state alone. The consumer dequants (chunked) instead. - """ - if not USE_BUFFER_LOAD: - cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) - bg, pg, valid = _slots( - indices_ptr, - seg_start, - k_start + k_rng_slot, - seg_hi, # hi: unused unless UNI_TILE; num_rows: MASKED only - 0, - BLOCK_SIZE, - False, - HAS_INVALID, - IDX_BUFFER_LOAD, - UNI_TILE, - ) - if UNIFORM: - NGRP: gl.constexpr = HEAD_SIZE // 64 - x_u8 = _cache_load( - cache_ptr, bg * cs0 + pg * HEAD_SIZE, offs_full, USE_BUFFER_LOAD - ) - sc = _cache_load(cache_bf16_ptr, bg * NGRP, offs_full // 64, USE_BUFFER_LOAD) - k_rope = x_u8 # unused for UNIFORM (rope slice-store skipped) -> DCE'd - else: - nope_row = bg * cs0 + pg * 576 - scl_row = bg * cs0 + BLOCK_SIZE * 576 + pg * 8 - # UE8M0 scales first, bulk fp8 after. vmcnt is one in-order FIFO, so a - # wait can only name "at most N outstanding", never a specific load: - # issuing the scales last would make the first dequant piece wait behind - # every data load as well. Issued first, the scales are covered by the - # wait that the first piece's own data needs anyway. - if NARROW_SCALE and not USE_BUFFER_LOAD: - sc = _scale_load( - cache_ptr, scl_row, scl_row, USE_BUFFER_LOAD, gather_l, scl_l, - HEAD_SIZE // 64, HEAD_SIZE, False, - ) - else: - sc = _cache_load(cache_ptr, scl_row, offs_full // 64, USE_BUFFER_LOAD) - if ASM_DEQ: - # 2-byte elements so the asm gets <2 x i16> = 4 packed fp8 in one VGPR out - # of a single dword load. Same bytes, same 16-B per-lane run, half the - # columns; cache_bf16_ptr is already the 2-byte view of this cache and the - # RoPE gather below uses the same cs0//2 / *288 addressing. - # The row vector comes from _slots in SliceLayout(1, gather_l); the i16 - # columns live in gather16_l, and a broadcast needs one parent layout. The - # two layouts share their dim-0 tiling, so this convert is a rename. - row16 = gl.convert_layout( - nope_row >> 1, gl.SliceLayout(1, gather16_l) - ) - x_u8 = _cache_load( - cache_bf16_ptr, row16, offs_full16, USE_BUFFER_LOAD - ) - else: - x_u8 = _cache_load(cache_ptr, nope_row, offs_full, USE_BUFFER_LOAD) - bgr, pgr, _ = _slots( - indices_ptr, - seg_start, - k_start + k_rng_rope, - seg_hi, - 0, - BLOCK_SIZE, - False, - HAS_INVALID, - IDX_BUFFER_LOAD, - UNI_TILE, - ) - k_rope = _cache_load( - cache_bf16_ptr, bgr * (cs0 // 2) + pgr * 288 + 224, offs_rope, USE_BUFFER_LOAD - ) - return x_u8, sc, k_rope, valid - - -@gluon.jit -def _qkpv_fp8( - x_u8, - sc, - k_rope, - valid, - q_dot, - m_i, - l_i, - acc, - head_mask, - qk_scale, - kv_smem, - qk_layout: gl.constexpr, - pv_layout: gl.constexpr, - k_layout: gl.constexpr, - v_layout: gl.constexpr, - p_layout: gl.constexpr, - gather_l: gl.constexpr, - NOPE_DIM: gl.constexpr, - ROPE_DIM: gl.constexpr, - HEAD_SIZE: gl.constexpr, - BLOCK_M: gl.constexpr, - BLOCK_K: gl.constexpr, - NOPE_CHUNK: gl.constexpr, - CHUNK_AXIS: gl.constexpr, - ASM_DEQ: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, - UNIFORM: gl.constexpr, - FP8_FNUZ: gl.constexpr, - HAS_INVALID: gl.constexpr, - UNI_TILE: gl.constexpr = False, - k_start=0, - seg_hi=0, -): - """Dequant a prefetched fp8 tile into LDS, then QK -> softmax -> PV. UNIFORM - skips the RoPE slice-store (the whole head is one fp8 tile). When HAS_INVALID, - mask the columns of -1-sentinel slots (``valid`` from _gd_fp8) to -inf.""" - neg_inf = float("-inf") - _deq_store_tile(x_u8, sc, kv_smem, NOPE_CHUNK, CHUNK_AXIS, ASM_DEQ, gather_l, UNIFORM, FP8_FNUZ) - if not UNIFORM: - kv_smem.slice(NOPE_DIM, ROPE_DIM, dim=1).store(k_rope) - k = kv_smem.permute([1, 0]).load(k_layout) - S = gl.amd.cdna4.mfma( - q_dot, k, gl.zeros([BLOCK_M, BLOCK_K], gl.float32, layout=qk_layout) - ) - # UNI_TILE folds the tile's range test into `valid`, so the score mask is no - # longer only about -1 sentinels: it is what makes the partial last tile correct. - COL_MASK: gl.constexpr = HAS_INVALID or UNI_TILE - NEED_MASK: gl.constexpr = COL_MASK or (not HEAD_ALIGNED) - if NEED_MASK: - if COL_MASK: - if UNI_TILE and not HAS_INVALID: - # Build the range mask directly in the MFMA layout. Converting the - # slot-layout `valid` vector instead costs a cross-lane layout change - # per tile, which at 16 tiles ate the whole saving. - col_mask = ( - k_start - + gl.arange(0, BLOCK_K, layout=gl.SliceLayout(0, qk_layout)) - < seg_hi - )[None, :] - else: - col_mask = gl.convert_layout(valid, gl.SliceLayout(0, qk_layout))[None, :] - if not HEAD_ALIGNED: - col_mask = ( - gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[:, None] - & col_mask - ) - else: - col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, qk_layout))[ - :, None - ] - S = gl.where(col_mask, S, neg_inf) - # m_i is carried in base-2 exponent space (see _decode_tile), hence qk_scale. - # max commutes with a positive scale (qk_scale = scale * log2(e); any real attention - # scale is > 0), so scale the row max -- one value per lane -- rather than every - # element of S. That leaves `S * qk_scale - m_new` as the only use of the scaled S, - # which lowers to v_fma_f32 (16 of them) where the split form needed a separate - # multiply over the tile. -inf * positive = -inf, so masked columns stay masked, and - # the row max is the same product bit for bit; only the fused subtract differs, one - # rounding instead of two -- measured <= 0.23 of a bf16 output ulp. - # - # PERFORMANCE-NEUTRAL: 0.989-1.006x across C=32/64/128 x top-k 64/272/1024 and the - # 64-bit path, inside the control arm's own 0.996-1.003x. S is [BLOCK_M, BLOCK_K], - # only 4 f32 per lane per tile, so nothing here can move a 10 us kernel. Kept - # because it is strictly fewer instructions and one fewer rounding. The multiply - # that *would* matter is `acc * alpha_pv` (128 f32 per lane) and MFMA has no - # accumulator-scale operand, so it cannot fuse. - m_block = _rmax(S, 1) * qk_scale - m_new = _max2(m_i, m_block) - m_new = gl.where(m_new > neg_inf, m_new, 0.0) - p = gl.exp2(S * qk_scale - m_new[:, None]) - alpha = gl.exp2(m_i - m_new) - l_new = l_i * alpha + gl.sum(p, axis=1) - v = kv_smem.load(v_layout) - p_dot = gl.convert_layout(p.to(gl.bfloat16), p_layout) - alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, pv_layout)) + v = kv_smem.load(cfg.v_layout) # [BLOCK_K, KV_DIM] + if seg.fmt.KIND == "tensor": + p = p * v_scale # per-tensor V scale on the small side (see _qkpv) + p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) + alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.pv_layout)) acc = acc * alpha_pv[:, None] acc = gl.amd.cdna4.mfma(p_dot, v, acc) return m_new, l_new, acc @@ -829,306 +1277,161 @@ def _qkpv_fp8( @gluon.jit def _process_segment( - UNI_TILE: gl.constexpr, + cfg, + seg, q_dot, - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, + q_rope_dot, lo, hi, - cs0, - num_rows, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, - qk_layout: gl.constexpr, - pv_layout: gl.constexpr, - k_layout: gl.constexpr, - v_layout: gl.constexpr, - p_layout: gl.constexpr, - gather_l: gl.constexpr, - gather_rope_l: gl.constexpr, - gather16_l: gl.constexpr, - scl_l: gl.constexpr, - NARROW_SCALE: gl.constexpr, - slot_l: gl.constexpr, - IS_FP8: gl.constexpr, - BLOCK_SIZE: gl.constexpr, - NOPE_DIM: gl.constexpr, - ROPE_DIM: gl.constexpr, - HEAD_SIZE: gl.constexpr, - BLOCK_M: gl.constexpr, - BLOCK_K: gl.constexpr, - NOPE_CHUNK: gl.constexpr, - CHUNK_AXIS: gl.constexpr, - ASM_DEQ: gl.constexpr, - HEAD_ALIGNED: gl.constexpr, - UNIFORM: gl.constexpr, - USE_BUFFER_LOAD: gl.constexpr, - IDX_BUFFER_LOAD: gl.constexpr, - HAS_INVALID: gl.constexpr, - FP8_FNUZ: gl.constexpr, + rope_smem, ): - offs_full = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, gather_l)) + offs_full = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.gather_l)) # int16 column offsets for the asm dequant's 2-byte gather (half the columns). - offs_full16 = gl.arange(0, HEAD_SIZE // 2, layout=gl.SliceLayout(0, gather16_l)) - offs_rope = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, gather_rope_l)) - k_rng_slot = gl.arange(0, BLOCK_K, layout=slot_l) - k_rng_rope = gl.arange(0, BLOCK_K, layout=gl.SliceLayout(1, gather_rope_l)) + offs_full16 = gl.arange( + 0, cfg.KV_DIM // 2, layout=gl.SliceLayout(0, cfg.gather16_l) + ) + offs_rope = gl.arange(0, cfg.ROPE_DIM, layout=gl.SliceLayout(0, cfg.gather_rope_l)) + k_rng_slot = gl.arange(0, cfg.BLOCK_K, layout=cfg.slot_l) + k_rng_rope = gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(1, cfg.gather_rope_l)) # Peel the (possibly partial) last tile: [lo, hi_full) are full BLOCK_K tiles # whose slots are all valid -> mask-free. Only the peeled tail carries masking. - hi_full = lo + ((hi - lo) // BLOCK_K) * BLOCK_K + hi_full = lo + ((hi - lo) // cfg.BLOCK_K) * cfg.BLOCK_K - if IS_FP8: + if seg.fmt.IS_FP8: # UNI_TILE: the partial tile is just the last iteration -- no peeled masked # copy of the body (RESULTS.md 19). - if UNI_TILE: - n_full = (hi - lo + BLOCK_K - 1) // BLOCK_K + if cfg.UNI_TILE: + n_full = (hi - lo + cfg.BLOCK_K - 1) // cfg.BLOCK_K else: - n_full = (hi_full - lo) // BLOCK_K + n_full = (hi_full - lo) // cfg.BLOCK_K if n_full > 0: - kn, ks, kr, vld = _gd_fp8( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, + kn, ks, kr, vld = _gather_full( + cfg, + seg, lo, - cs0, + hi, offs_full, offs_full16, offs_rope, k_rng_slot, k_rng_rope, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - ASM_DEQ, - BLOCK_SIZE, - HEAD_SIZE, - UNIFORM, - USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, - UNI_TILE, - hi, ) for i in range(1, n_full): - kn2, ks2, kr2, vld2 = _gd_fp8( - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, - lo + i * BLOCK_K, - cs0, + kn2, ks2, kr2, vld2 = _gather_full( + cfg, + seg, + lo + i * cfg.BLOCK_K, + hi, offs_full, offs_full16, offs_rope, k_rng_slot, k_rng_rope, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - ASM_DEQ, - BLOCK_SIZE, - HEAD_SIZE, - UNIFORM, - USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, - UNI_TILE, - hi, ) - m_i, l_i, acc = _qkpv_fp8( + m_i, l_i, acc = _qkpv( + cfg, + seg, kn, ks, kr, vld, q_dot, + q_rope_dot, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, - UNIFORM, - FP8_FNUZ, - HAS_INVALID, - UNI_TILE, - lo + (i - 1) * BLOCK_K, + rope_smem, + lo + (i - 1) * cfg.BLOCK_K, hi, ) kn, ks, kr, vld = kn2, ks2, kr2, vld2 - m_i, l_i, acc = _qkpv_fp8( + m_i, l_i, acc = _qkpv( + cfg, + seg, kn, ks, kr, vld, q_dot, + q_rope_dot, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, - UNIFORM, - FP8_FNUZ, - HAS_INVALID, - UNI_TILE, - lo + (n_full - 1) * BLOCK_K, + rope_smem, + lo + (n_full - 1) * cfg.BLOCK_K, hi, ) else: - for k_start in range(lo, hi_full, BLOCK_K): + for k_start in range(lo, hi_full, cfg.BLOCK_K): m_i, l_i, acc = _decode_tile( + cfg, + seg, q_dot, - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, + q_rope_dot, k_start, hi, - cs0, - num_rows, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, + rope_smem, offs_full, offs_full16, offs_rope, k_rng_slot, k_rng_rope, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - IS_FP8, - BLOCK_SIZE, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, False, - UNIFORM, - USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, - UNI_TILE, - hi, ) - if (not UNI_TILE) or (not IS_FP8): + if (not cfg.UNI_TILE) or (not seg.fmt.IS_FP8): if hi_full < hi: m_i, l_i, acc = _decode_tile( + cfg, + seg, q_dot, - cache_ptr, - cache_bf16_ptr, - indices_ptr, - seg_start, + q_rope_dot, hi_full, hi, - cs0, - num_rows, m_i, l_i, acc, head_mask, qk_scale, + v_scale, kv_smem, + rope_smem, offs_full, offs_full16, offs_rope, k_rng_slot, k_rng_rope, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - IS_FP8, - BLOCK_SIZE, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, True, - UNIFORM, - USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, ) return m_i, l_i, acc _pa_decode_sparse_repr = make_kernel_repr( "_pa_decode_sparse", - ["BLOCK_M", "BLOCK_K", "HEAD_SIZE", "NUM_SPLITS", "UNIFORM", "MAIN_IS_FP8"], + ["BLOCK_M", "BLOCK_K", "HEAD_SIZE", "NUM_SPLITS", "MAIN_FMT", "ROPE_SEPARATE"], ) @@ -1148,6 +1451,11 @@ def _pa_decode_sparse( part_m_ptr, part_l_ptr, part_acc_ptr, + # f32 side-channel per segment: the scalar k_scale ("tensor") or the f32 view + # of the cache ("dsmla"). None -> specialized out of the binary entirely, so + # DSv4 launches keep their exact kernarg layout. + main_scl_ptr, + extra_scl_ptr, scale: gl.constexpr, q_stride0: gl.constexpr, q_stride1: gl.constexpr, @@ -1165,20 +1473,20 @@ def _pa_decode_sparse( num_heads: gl.constexpr, HAS_EXTRA: gl.constexpr, HAS_SINK: gl.constexpr, - MAIN_IS_FP8: gl.constexpr, - EXTRA_IS_FP8: gl.constexpr, + MAIN_FMT: gl.constexpr, + EXTRA_FMT: gl.constexpr, MAIN_BLOCK_SIZE: gl.constexpr, EXTRA_BLOCK_SIZE: gl.constexpr, CS0_ALIGN: gl.constexpr, NOPE_DIM: gl.constexpr, ROPE_DIM: gl.constexpr, HEAD_SIZE: gl.constexpr, + ROPE_SEPARATE: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_K: gl.constexpr, NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, MFMA_K: gl.constexpr, - UNIFORM: gl.constexpr, GATHER_TW1: gl.constexpr, LDS_PAD: gl.constexpr, # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, 1 = @@ -1186,6 +1494,8 @@ def _pa_decode_sparse( NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, PART_STORE_CACHE: gl.constexpr, + UNI_TILE: gl.constexpr, + GRID_ORDER: gl.constexpr, # How many of the NUM_SPLITS programs share the main (SWA) segment. The two # segments have very different shapes -- main is a contiguous window whose # length is fixed by the sliding window, extra is the top-k list -- so one @@ -1194,8 +1504,6 @@ def _pa_decode_sparse( # extra still wants the CTAs. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an # empty main range and contribute an extra-only partial. - UNI_TILE: gl.constexpr, - GRID_ORDER: gl.constexpr, MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. ADAPTIVE_SPLITS: gl.constexpr, @@ -1216,13 +1524,33 @@ def _pa_decode_sparse( un-normalized partials for the reduce kernel. HAS_INVALID gates -1-sentinel handling (clamp + score mask) on the full-tile fast paths.""" NUM_WARPS: gl.constexpr = gl.num_warps() + # The new fp8 formats never trace the peeled masked tail (_decode_tile); + # requiring UNI_TILE keeps it that way instead of silently miscompiling. + gl.static_assert( + UNI_TILE or (MAIN_FMT != "tensor" and MAIN_FMT != "dsmla"), + "tensor/dsmla formats require UNI_TILE=1", + ) + gl.static_assert( + UNI_TILE or (EXTRA_FMT != "tensor" and EXTRA_FMT != "dsmla"), + "tensor/dsmla formats require UNI_TILE=1", + ) + gl.static_assert( + (not ROPE_SEPARATE) or (MAIN_FMT != "dsv4" and MAIN_FMT != "uniform"), + "dsv4/uniform formats carry rope inside plane 0 (ROPE_SEPARATE=False)", + ) + gl.static_assert( + MAIN_FMT != "dsmla" or ROPE_SEPARATE, + "fp8_ds_mla is a separated-rope (GLM) format", + ) + gl.static_assert((not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", + "ASM_DEQ is the dsv4 E8M0 dequant") # Page stride alignment hint. The row base of every cache gather is - # ``block_idx*cs0 + pos*576`` (packed) / ``... + pos*HEAD_SIZE`` (uniform) with - # block_idx/pos gathered at runtime, so the divisibility analysis has to assume - # 1-byte alignment and a contiguous 512-byte row gather lowers to hundreds of - # global_load_ubyte instead of global_load_dwordx4. 576 and HEAD_SIZE are - # literals the compiler can already reason about; cs0 is the one opaque term. - # The driver sets CS0_ALIGN>1 only after checking the strides on the host. + # ``block_idx*cs0 + pos*TOK`` with block_idx/pos gathered at runtime, so the + # divisibility analysis has to assume 1-byte alignment and a contiguous + # 512-byte row gather lowers to hundreds of global_load_ubyte instead of + # global_load_dwordx4. TOK/HEAD_SIZE are literals the compiler can already + # reason about; cs0 is the one opaque term. The driver sets CS0_ALIGN>1 only + # after checking the strides on the host. if CS0_ALIGN > 1: main_cs0 = gl.multiple_of(main_cs0, CS0_ALIGN) extra_cs0 = gl.multiple_of(extra_cs0, CS0_ALIGN) @@ -1232,102 +1560,42 @@ def _pa_decode_sparse( split_id = gl.program_id(GRID_ORDER.index("s")) pid_h = gl.program_id(GRID_ORDER.index("h")) - qk_layout: gl.constexpr = gl.amd.AMDMFMALayout( - version=4, - instr_shape=[16, 16, MFMA_K], - transposed=True, - warps_per_cta=[1, NUM_WARPS], - ) - pv_layout: gl.constexpr = gl.amd.AMDMFMALayout( - version=4, - instr_shape=[16, 16, MFMA_K], - transposed=True, - warps_per_cta=[1, NUM_WARPS], - ) - KW: gl.constexpr = MFMA_K // 2 - q_layout: gl.constexpr = gl.DotOperandLayout(0, qk_layout, KW) - k_layout: gl.constexpr = gl.DotOperandLayout(1, qk_layout, KW) - p_layout: gl.constexpr = gl.DotOperandLayout(0, pv_layout, KW) - v_layout: gl.constexpr = gl.DotOperandLayout(1, pv_layout, KW) - - # 16 uint8 = 128-bit fp8 gather loads - GSPT: gl.constexpr = 16 - # Warps on dim 1 (columns) vs dim 0 (rows). Coalescing is identical either way - # -- a row's 128 contiguous bytes always come from 8 threads -- but the row - # index vector lives in SliceLayout(1, gather_l), so warps-on-dim-1 replicates - # it across all NUM_WARPS and every lane carries BLOCK_K/8 slots, while - # warps-on-dim-0 carries BLOCK_K/(8*NUM_WARPS). Dim 0 also makes the column - # direction a pure per-lane register repeat, so a chunked dequant can split it - # with a trivial (register-renaming) layout convert. - # GATHER_TW1 = threads spent on the head dim. Each thread loads 16 B, so - # TW1=8 requests 128 B of a token row per instruction and TW1=32 requests the - # whole 512 B row -- for a scattered gather that is one request instead of - # four quarters. The cost is that the row-index vector lives in - # SliceLayout(1, gather_l), so a wider TW1 leaves fewer dim-0 thread slots and - # every lane carries more slots (BLOCK_K*TW1/64 of them). Whichever dim keeps - # per-lane register repeats is the one the chunked dequant can split - # (CHUNK_AXIS). - gather_l: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, GSPT], - threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], - warps_per_cta=[NUM_WARPS, 1], - order=[1, 0], - ) - # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting - # them on dim 1 (as the NoPE gather does) overshoots 4x and every warp re-gathers - # the same tile. Worth ~10% on packed fp8. - gather_rope_l: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 8], - threads_per_warp=[8, 8], - warps_per_cta=[NUM_WARPS, 1], - order=[1, 0], - ) - # 3-D companion of gather_l used by _scale_load: dim 1 carries the NG scale - # groups (one thread each) and dim 2 the 64 columns inside a group, so that - # reshaping [BLOCK_K, NG, 64] back to 2-D reproduces gather_l exactly. Legal - # only when the group's columns fill whole threads, i.e. GSPT * (TW1 // NG) is - # the group width; otherwise fall back to the wide (redundant) gather. - NGRP_S: gl.constexpr = HEAD_SIZE // 64 - NARROW_SCALE: gl.constexpr = ( - GATHER_TW1 % NGRP_S == 0 and GSPT * (GATHER_TW1 // NGRP_S) == 64 - ) - # ...and only worth it on the 64-bit path. With buffer_load the wide gather's - # identical 32-bit offsets CSE to the same 48 byte loads, so the narrow form - # only adds a layout convert per tile (+3.4% at extra=1024). Without it they do - # not CSE and the wide form costs 288 loads. - scl_l3: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 1, GSPT], - threads_per_warp=[ - 64 // GATHER_TW1, - NGRP_S if NARROW_SCALE else 1, - (GATHER_TW1 // NGRP_S) if NARROW_SCALE else GATHER_TW1, - ], - warps_per_cta=[NUM_WARPS, 1, 1], - order=[2, 1, 0], - ) - scl_l: gl.constexpr = gl.SliceLayout(2, scl_l3) - # Same tiling as gather_l but 2-byte elements: half the per-thread run, so a lane - # still covers the same 16 bytes and the lo/hi interleave stays lane-local. - gather16_l: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, GSPT // 2], - threads_per_warp=[64 // GATHER_TW1, GATHER_TW1], - warps_per_cta=[NUM_WARPS, 1], - order=[1, 0], + cfg = Cfg( + BLOCK_M, + BLOCK_K, + HEAD_SIZE, + ROPE_DIM, + ROPE_SEPARATE, + MFMA_K, + NUM_WARPS, + GATHER_TW1, + LDS_PAD, + UNI_TILE, + HAS_INVALID, + HEAD_ALIGNED, + IDX_BUFFER_LOAD, ) - slot_l: gl.constexpr = gl.SliceLayout(1, gather_l) - blocked_q: gl.constexpr = gl.BlockedLayout( - size_per_thread=[1, 8], - threads_per_warp=[8, 8], - warps_per_cta=[1, NUM_WARPS], - order=[1, 0], + main_fmt = Fmt( + cfg, + MAIN_FMT, + MAIN_BLOCK_SIZE, + MAIN_USE_BUFFER_LOAD, + FP8_FNUZ, + ASM_DEQ and MAIN_FMT == "dsv4", + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, ) - # LDS pad after every row. Row pitch (512 + LDS_PAD) bf16 sets which banks a - # transposed K read (ds_read_b64_tr_b16 walks down a column) lands on: - # bank = (row * pitch_dwords) mod 32, so pitch_dwords mod 32 == 4 at PAD=8 means - # 32 lanes share 8 banks. Kept as a knob because the conflict share is high - # (46%) but a high sub-metric ratio is not proof of a bottleneck. - kv_shared: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( - [[HEAD_SIZE, LDS_PAD]], [BLOCK_K, HEAD_SIZE], [1, 0] + extra_fmt = Fmt( + cfg, + EXTRA_FMT, + EXTRA_BLOCK_SIZE, + EXTRA_USE_BUFFER_LOAD, + FP8_FNUZ, + ASM_DEQ and EXTRA_FMT == "dsv4", + NOPE_DIM, + NOPE_CHUNK, + CHUNK_AXIS, ) h_off = pid_h * BLOCK_M @@ -1350,9 +1618,34 @@ def _pa_decode_sparse( extra_len = 0 main_len = main_end - main_start - # ---- load Q [BLOCK_M, HEAD_SIZE] ---- - offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, blocked_q)) - offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, blocked_q)) + # ---- per-segment softmax/V scales ---- + # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw `scale` + # for the sink/normalization (sink is a scaled-score-space logit). For the + # per-tensor fp8 format the cache's k_scale joins the fold (max/exp2 commute + # with a positive scale) and the V-side scale is applied to p in the loop -- + # both are exact, so m/l keep the dequantized-score-domain convention the + # reduce kernel expects. + RCP_LN2: gl.constexpr = 1.4426950408889634 + qk_scale = scale * RCP_LN2 + main_qk_scale = qk_scale + main_v_scale = 1.0 + if MAIN_FMT == "tensor": + main_k_scale = gl.load(main_scl_ptr) + main_qk_scale = qk_scale * main_k_scale + main_v_scale = main_k_scale + extra_qk_scale = qk_scale + extra_v_scale = 1.0 + if HAS_EXTRA and EXTRA_FMT == "tensor": + extra_k_scale = gl.load(extra_scl_ptr) + extra_qk_scale = qk_scale * extra_k_scale + extra_v_scale = extra_k_scale + + # ---- load Q [BLOCK_M, QK_DIM] ---- + # ROPE_SEPARATE loads the two q pieces separately (576 is not a pow-2 arange): + # nope columns [0, KV_DIM) and rope columns [KV_DIM, KV_DIM + ROPE_DIM), each + # converted to the dot-operand layout of its own MFMA in the QK chain. + offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.blocked_q)) + offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.blocked_q)) h_q = h_off + offs_m_q h_mask_q = h_q < num_heads q_off = (query_idx * q_stride0 + h_q[:, None] * q_stride1 + offs_d_q[None, :]).to( @@ -1361,26 +1654,43 @@ def _pa_decode_sparse( q = gl.amd.cdna4.buffer_load( ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 ) - q_dot = gl.convert_layout(q, q_layout) + q_dot = gl.convert_layout(q, cfg.q_layout) + if ROPE_SEPARATE: + offs_d_qr = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, cfg.blocked_q)) + qr_off = ( + query_idx * q_stride0 + + h_q[:, None] * q_stride1 + + HEAD_SIZE + + offs_d_qr[None, :] + ).to(gl.int32) + q_rope = gl.amd.cdna4.buffer_load( + ptr=q_ptr, offsets=qr_off, mask=h_mask_q[:, None], other=0.0 + ) + q_rope_dot = gl.convert_layout(q_rope, cfg.q_layout) + else: + q_rope_dot = q_dot # unused (single-plane QK) -> DCE'd # head mask in pv-slice layout (for output / partial masking) - offs_m_pv = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, pv_layout)) + offs_m_pv = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout)) h_pv = h_off + offs_m_pv head_mask_pv = h_pv < num_heads # ---- online-softmax state ---- m_i = gl.full( - [BLOCK_M], float("-inf"), gl.float32, layout=gl.SliceLayout(1, qk_layout) + [BLOCK_M], float("-inf"), gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout) ) - l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, qk_layout)) - acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=pv_layout) + l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout)) + acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=cfg.pv_layout) - kv_smem = gl.allocate_shared_memory(gl.bfloat16, [BLOCK_K, HEAD_SIZE], kv_shared) - - # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw `scale` - # for the sink/normalization (sink is a scaled-score-space logit). - RCP_LN2: gl.constexpr = 1.4426950408889634 - qk_scale = scale * RCP_LN2 + kv_smem = gl.allocate_shared_memory( + gl.bfloat16, [BLOCK_K, HEAD_SIZE], cfg.kv_shared + ) + if ROPE_SEPARATE: + rope_smem = gl.allocate_shared_memory( + gl.bfloat16, [BLOCK_K, ROPE_DIM], cfg.rope_shared + ) + else: + rope_smem = kv_smem # never sliced/read as plane 1 in this geometry # ---- how many of the launched splits this query actually wants ---- # NUM_SPLITS / MAIN_SPLITS come from the host, which only knows the batch's @@ -1395,7 +1705,9 @@ def _pa_decode_sparse( if ADAPTIVE_SPLITS: m_tiles = (main_len + BLOCK_K - 1) // BLOCK_K e_tiles = (extra_len + BLOCK_K - 1) // BLOCK_K - work_splits = gl.minimum(gl.maximum(gl.maximum(m_tiles, e_tiles), 1), NUM_SPLITS) + work_splits = gl.minimum( + gl.maximum(gl.maximum(m_tiles, e_tiles), 1), NUM_SPLITS + ) main_splits = gl.minimum(gl.maximum(m_tiles, 1), work_splits) if split_id >= work_splits: pm_base = query_idx * pm_stride0 + split_id * pm_stride_s @@ -1404,14 +1716,14 @@ def _pa_decode_sparse( [BLOCK_M], float("-inf"), gl.float32, - layout=gl.SliceLayout(1, pv_layout), + layout=gl.SliceLayout(1, cfg.pv_layout), ), ptr=part_m_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv, ) gl.amd.cdna4.buffer_store( - gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, pv_layout)), + gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.pv_layout)), ptr=part_l_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv, @@ -1422,109 +1734,71 @@ def _pa_decode_sparse( main_splits = MAIN_SPLITS # ---- main (SWA) segment ---- + main_seg = Seg( + main_fmt, + main_cache_ptr, + main_cache_bf16_ptr, + main_scl_ptr if (MAIN_FMT == "dsmla") else main_cache_ptr, + main_indices_ptr, + main_start, + main_cs0, + main_num_rows, + ) main_chunk = (main_len + main_splits - 1) // main_splits main_lo = gl.minimum(split_id * main_chunk, main_len) main_hi = gl.minimum(main_lo + main_chunk, main_len) m_i, l_i, acc = _process_segment( - UNI_TILE, + cfg, + main_seg, q_dot, - main_cache_ptr, - main_cache_bf16_ptr, - main_indices_ptr, - main_start, + q_rope_dot, main_lo, main_hi, - main_cs0, - main_num_rows, m_i, l_i, acc, head_mask_pv, - qk_scale, + main_qk_scale, + main_v_scale, kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - slot_l, - MAIN_IS_FP8, - MAIN_BLOCK_SIZE, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, - UNIFORM, - MAIN_USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, + rope_smem, ) if HAS_EXTRA: + extra_seg = Seg( + extra_fmt, + extra_cache_ptr, + extra_cache_bf16_ptr, + extra_scl_ptr if (EXTRA_FMT == "dsmla") else extra_cache_ptr, + extra_indices_ptr, + extra_start, + extra_cs0, + extra_num_rows, + ) extra_chunk = (extra_len + work_splits - 1) // work_splits extra_lo = split_id * extra_chunk extra_hi = gl.minimum(extra_lo + extra_chunk, extra_len) m_i, l_i, acc = _process_segment( - UNI_TILE, + cfg, + extra_seg, q_dot, - extra_cache_ptr, - extra_cache_bf16_ptr, - extra_indices_ptr, - extra_start, + q_rope_dot, extra_lo, extra_hi, - extra_cs0, - extra_num_rows, m_i, l_i, acc, head_mask_pv, - qk_scale, + extra_qk_scale, + extra_v_scale, kv_smem, - qk_layout, - pv_layout, - k_layout, - v_layout, - p_layout, - gather_l, - gather_rope_l, - gather16_l, - scl_l, - NARROW_SCALE, - slot_l, - EXTRA_IS_FP8, - EXTRA_BLOCK_SIZE, - NOPE_DIM, - ROPE_DIM, - HEAD_SIZE, - BLOCK_M, - BLOCK_K, - NOPE_CHUNK, - CHUNK_AXIS, - ASM_DEQ, - HEAD_ALIGNED, - UNIFORM, - EXTRA_USE_BUFFER_LOAD, - IDX_BUFFER_LOAD, - HAS_INVALID, - FP8_FNUZ, + rope_smem, ) # m_i/l_i are in SliceLayout(1, qk_layout); acc in pv_layout. Move the row # reductions into pv-slice space for output/partials. - m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, pv_layout)) - l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, pv_layout)) + m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.pv_layout)) + l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.pv_layout)) if NUM_SPLITS == 1: if HAS_SINK: @@ -1542,7 +1816,7 @@ def _pa_decode_sparse( l_final = l_pv one_over_l = 1.0 / l_final out = acc * one_over_l[:, None] - offs_d_o = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) + offs_d_o = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.pv_layout)) o_off = ( query_idx * out_stride0 + h_pv[:, None] * out_stride1 + offs_d_o[None, :] ).to(gl.int32) @@ -1575,7 +1849,7 @@ def _pa_decode_sparse( mask=head_mask_pv, cache=PART_STORE_CACHE, ) - offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, pv_layout)) + offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.pv_layout)) a_base = query_idx * pa_stride0 + split_id * pa_stride_s a_off = (a_base + h_pv[:, None] * pa_stride_h + offs_d_a[None, :]).to(gl.int32) # Follow part_acc's own dtype: at bf16 this halves both the partial traffic @@ -1622,8 +1896,10 @@ def _pa_decode_sparse_reduce( the final output. Partials store m in the base-2 exponent domain (row-max * softmax_scale * log2e), matching the triton reduce. Grid: (num_queries, heads_blocks). - BLOCK_M is the reduce's own head tile and is deliberately decoupled from the - attention kernel's: the combine is pure bandwidth over + Identical for both geometries: partials are [*, HEAD_SIZE] where HEAD_SIZE is + the V width (512 for DSv4 and GLM-5 alike) -- the rope plane never reaches the + output. BLOCK_M is the reduce's own head tile and is deliberately decoupled + from the attention kernel's: the combine is pure bandwidth over [num_queries, num_heads, HEAD_SIZE] f32, so the only thing that matters is having enough workgroups to cover the machine. At BLOCK_M = num_heads there is one workgroup per query -- 64 of them on a 256-CU part, i.e. 3/4 of the GPU diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 6c764ea7927..3c5a26075a4 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -663,6 +663,15 @@ def _pa_decode_sparse_gfx950_gluon( avg_main = indices.numel() / max(1, num_queries) avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 + # Kernel-side cache-format tags (the kernel is shared with the separated-rope + # GLM/MLA path in sparse_mla_decode.py; DSv4 always runs ROPE_SEPARATE=False). + if UNIFORM: + main_fmt = "uniform" if main_is_fp8 else "bf16" + extra_fmt = main_fmt + else: + main_fmt = "dsv4" if main_is_fp8 else "bf16" + extra_fmt = "dsv4" if extra_is_fp8 else "bf16" + # Largest power-of-2 (<=16) dividing BOTH page strides. Lets the kernel assert # 16B-aligned row bases so a contiguous row gather can vectorize to dwordx4. _s0 = int(cache.stride(0)); _s1 = int(extra_cache.stride(0)) @@ -889,6 +898,11 @@ def _pick_splits(hb): part_m, part_l, part_acc, + # f32 side-channel pointers (per-tensor k_scale / fp8_ds_mla view) -- + # only the separated-rope formats use them; None keeps the DSv4 + # kernarg layout byte-identical (constexpr-None args are elided). + None, + None, scale, q.stride(0), q.stride(1), @@ -906,20 +920,20 @@ def _pick_splits(hb): num_heads, HAS_EXTRA=has_extra, HAS_SINK=has_sink, - MAIN_IS_FP8=main_is_fp8, - EXTRA_IS_FP8=extra_is_fp8, + MAIN_FMT=main_fmt, + EXTRA_FMT=extra_fmt, MAIN_BLOCK_SIZE=main_block, EXTRA_BLOCK_SIZE=extra_block, CS0_ALIGN=cs0_align, NOPE_DIM=nope_dim, ROPE_DIM=ROPE_DIM, HEAD_SIZE=head_dim, + ROPE_SEPARATE=False, BLOCK_M=BLOCK_M, BLOCK_K=BLOCK_K, NUM_SPLITS=num_splits, HEAD_ALIGNED=HEAD_ALIGNED, MFMA_K=MFMA_K, - UNIFORM=UNIFORM, GATHER_TW1=gather_tw1, LDS_PAD=lds_pad, NOPE_CHUNK=nope_chunk, From da5aabbea73e31279ef80e3af82744bfbaf6031a Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 03:13:35 +0000 Subject: [PATCH 25/51] [GLUON][GFX950][MLA] Add sparse_mla_decode_fwd: separated-rope sparse 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 Co-Authored-By: Claude Opus 5 (1M context) --- .../ops/triton/attention/sparse_mla_decode.py | 428 ++++++++++++++++++ .../attention/test_sparse_mla_decode.py | 224 +++++++++ 2 files changed, 652 insertions(+) create mode 100644 aiter/ops/triton/attention/sparse_mla_decode.py create mode 100644 op_tests/triton_tests/attention/test_sparse_mla_decode.py diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py new file mode 100644 index 00000000000..0d6a8d86a2b --- /dev/null +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Sparse MLA decode (gfx950 gluon): separated-rope geometry — GLM-5 / DSA-style +``kv_lora_rank`` latent ⊕ appended decoupled rope, token-granular top-k gather. + +This is the separated-rope sibling of ``pa_decode_sparse``: both drive the same +gluon kernel (``_gluon_kernels/gfx950/attention/pa_decode_sparse.py``), but the +two model families have different enough calling structures that they get +different APIs — DSv4 (rope *inside* the 512-wide row, K == V, optional SWA + +top-k two-loop, attention sink) stays on ``pa_decode_sparse``; MLA-lineage +models (rope *appended*: QK over ``kv_lora_rank + qk_rope_head_dim``, V = the +latent only, single top-k segment, no sink) use this module. + +The calling convention follows the asm path this replaces +(``aiter.mla.mla_decode_fwd`` as invoked by vLLM's ``ROCMAiterMLASparseImpl``), +with two deliberate deviations, both favorable: no head padding (the asm kernel +repeat_interleaves H<16 to 16 and computes the duplicate rows; here H=8 runs +natively at BLOCK_M=8), and q stays bf16 (no q-quant kernel; <1% accuracy +effect). + +Cache formats, inferred from the buffer: + + [nb, block, R] / [slots, 1, 1, R] / [slots, R] bf16 full-row bf16 + same shapes in fp8/uint8 (+ scalar ``kv_scale``) flat per-tensor fp8 + — vLLM's production layout: ``concat_and_cache_mla(...,"fp8", + scale=layer._k_scale)`` quantizes the whole row (rope included) with + one scale. The kernel never dequantizes in the tile loop: the K-side + scale folds into qk_scale and the V-side scale into p (exact; see the + kernel docstring). + [blocks, block_size, 656] uint8 vLLM ``fp8_ds_mla`` + — 512 fp8 | 4×f32 per-128 scales | 64 bf16 rope. + +Indices are GLOBAL slot ids (``triton_convert_req_index_to_global_index`` +output). NOTE: that converter maps -1 / out-of-range entries to slot 0 despite +its docstring, so the production stream carries no -1 sentinels and the default +``has_invalid=False`` matches it (and skips the sentinel clamp+mask). Pass +``has_invalid=True`` for index streams that really carry -1 entries — they are +then masked out, which is *more* correct than the converter's slot-0 +substitution. +""" + +import math +import os as _os + +import torch + +from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( + _pa_decode_sparse as _pa_decode_sparse_gfx950, +) +from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( + _pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950, +) +from aiter.ops.triton.attention.pa_decode_sparse import ( + _as_int32_contiguous_1d, + _decode_num_splits_occ, +) +from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.common_utils import max_addressable_bytes +from aiter.ops.triton.utils.device_info import get_num_sms +from aiter.ops.triton.utils.logger import AiterTritonLogger + +_LOGGER = AiterTritonLogger() + +MAX_BYTES = 2**31 - 1 + + +def _check_out(out, q, kv_lora_rank): + """Caller-supplied output buffer, or a fresh one. The attention output is + [C, H, kv_lora_rank]; a wider buffer (e.g. a q-shaped scratch) is accepted + and written in its leading kv_lora_rank columns.""" + C, H = q.shape[0], q.shape[1] + if out is None: + return torch.empty( + (C, H, kv_lora_rank), dtype=torch.bfloat16, device=q.device + ) + assert out.shape[0] == C and out.shape[1] == H, ( + f"out shape {tuple(out.shape)} != [{C}, {H}, >= {kv_lora_rank}]" + ) + assert out.shape[2] >= kv_lora_rank and out.stride(2) == 1 + assert out.dtype == torch.bfloat16 and out.device == q.device + return out + + +def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): + """-> (fmt, cache, alt_ptr_tensor, scl_ptr_tensor, block_size, fp8_fnuz). + + Pointer roles follow the kernel's Seg contract: alt = the bf16 view + ("dsmla" rope tail) or the cache itself ("bf16"); scl = the f32 scalar + k_scale ("tensor") or the f32 view of the cache ("dsmla").""" + if kv.ndim == 4: # [slots, 1, 1, R] -- the asm mla_decode_fwd view + assert kv.shape[1] == 1 and kv.shape[2] == 1 + kv = kv.reshape(kv.shape[0], kv.shape[3]) + elif kv.ndim == 3 and kv.shape[2] == d_qk: + # the vLLM paged cache [nb, block_size, R]: indices are global slot + # ids and rows are contiguous, so the flat view is stride-identical. + assert kv.stride(2) == 1 and kv.stride(1) == kv.shape[2] * kv.stride(2) + kv = kv.reshape(-1, d_qk) + + fp8_dtypes = (torch.uint8, torch.float8_e4m3fn, torch.float8_e4m3fnuz) + if kv.ndim == 2 and kv.shape[1] == d_qk: + if kv.dtype == torch.bfloat16: + return "bf16", kv, kv, None, 1, False + assert kv.dtype in fp8_dtypes, f"unsupported flat cache dtype {kv.dtype}" + assert kv_scale is not None, ( + "flat fp8 cache needs the per-tensor kv_scale (layer._k_scale)" + ) + assert kv_scale.dtype == torch.float32 + u8 = kv.view(torch.uint8) + fnuz = kv.dtype == torch.float8_e4m3fnuz + return "tensor", u8, u8, kv_scale.reshape(1), 1, fnuz + if ( + 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 + ): + # vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B + u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + assert u8.stride(2) == 1 and u8.stride(1) == u8.shape[2], ( + "fp8_ds_mla rows must be contiguous 656-byte records" + ) + return ( + "dsmla", + u8, + u8.view(torch.bfloat16), + u8.view(torch.float32), + u8.shape[1], + False, + ) + raise ValueError(f"unrecognized sparse-MLA kv cache: {tuple(kv.shape)} {kv.dtype}") + + +def sparse_mla_decode_fwd( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + softmax_scale: float, + kv_scale: torch.Tensor | None = None, + kv_lora_rank: int = 512, + qk_rope_head_dim: int = 64, + kv_splits: int | None = None, + skip_reduce: bool = False, + has_invalid: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse (top-k gathered) MLA decode, separated-rope geometry. + + Args: + q: ``[C, H, kv_lora_rank + qk_rope_head_dim]`` bf16 decode queries. + H = per-rank query heads; 8 (TP8) and 16 (TP4) both run natively. + kv_buffer: the KV pool in one of the formats in the module docstring. + kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts + (vLLM builds ``cumsum(min(seq_len, topk))``; uniform ``C * topk`` + stepping also works). + kv_indices: flat int32 GLOBAL slot ids into the pool. + softmax_scale: the layer's softmax scale. + kv_scale: ``[1]`` f32 per-tensor cache scale — required for the flat + fp8 format (vLLM's ``layer._k_scale``), unused otherwise. + kv_splits: split-K override; default picks by the occupancy policy. + skip_reduce: with split-K active, return ``(part_acc, part_m, + part_l)`` instead of launching the combine. + has_invalid: index stream carries -1 sentinels (masked out). Default + False — the production converter emits none. + out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination (written in + ``[..., :kv_lora_rank]``). + + Returns: + ``[C, H, kv_lora_rank]`` bf16 attention output (the latent V; callers + run ``v_up_proj`` after, exactly as with the asm path). + """ + assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" + assert arch_info.get_arch() == "gfx950", "sparse_mla_decode_fwd is gfx950-only" + assert q.dtype == torch.bfloat16, "q stays bf16 (no q-quant; see module doc)" + num_queries, num_heads, d_qk = q.shape + assert d_qk == kv_lora_rank + qk_rope_head_dim, ( + f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" + ) + _LOGGER.info( + f"SPARSE_MLA_DECODE C={num_queries} H={num_heads} d_qk={d_qk} " + f"nnz={kv_indices.shape[0]}" + ) + + fmt, cache, alt, scl, block_size, fp8_fnuz = _infer_cache_format( + kv_buffer, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale + ) + kv_indices = _as_int32_contiguous_1d(kv_indices) + kv_indptr = _as_int32_contiguous_1d(kv_indptr) + assert kv_indptr.numel() == num_queries + 1 + # HAS_SINK=False: MLA-lineage models carry no attention sink. The kernel + # still wants a live pointer for the elided-at-compile-time argument slot. + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) + + # Tuned launch config (gfx950 / MI355), shared with the DSv4 wrapper; same + # env knobs so interleaved A/B sweeps drive both paths together. + BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 + _bk = _os.environ.get("AITER_PA_DECODE_BLOCK_K") + if _bk: + BLOCK_K = int(_bk) + _bm = _os.environ.get("AITER_PA_DECODE_BLOCK_M") + if _bm: + BLOCK_M = int(_bm) + elif num_heads < 16: + # H=8 (TP8): BLOCK_M=8 computes 8 real rows where a masked BLOCK_M=16 + # pads to 16 -- against the asm kernel's repeat_interleave to >=16 + # heads this is the free half of the measured win. + BLOCK_M = max(8, 1 << (num_heads - 1).bit_length()) + _mk = _os.environ.get("AITER_PA_DECODE_MFMA_K") + if _mk: + MFMA_K = int(_mk) + num_warps = max(1, BLOCK_K // 16) + gather_tw1 = int(_os.environ.get("AITER_PA_DECODE_TW1", "32")) + lds_pad = int(_os.environ.get("AITER_PA_DECODE_LDS_PAD", "8")) + + num_rows = cache.shape[0] * block_size if cache.ndim >= 2 else cache.shape[0] + avg_topk = kv_indices.numel() / max(1, num_queries) + + cs0_align = 1 + for _a in (16, 8, 4, 2): + if int(cache.stride(0)) % _a == 0: + cs0_align = _a + break + _ca = _os.environ.get("AITER_PA_DECODE_CS0_ALIGN") + if _ca: + cs0_align = int(_ca) + + # buffer_load carries a 32-bit offset (2 GB cap); a production pool crosses + # it at ~3.7M cached tokens, so the 64-bit gl.load path is first-class here + # (measured: same speed as the buffer path, +2-3 VGPRs, no spills). + use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES + idx_use_buffer_load = ( + _os.environ.get("AITER_PA_DECODE_IDX_BUF", "1") == "1" + and max_addressable_bytes(kv_indices) < MAX_BYTES + ) + _fg = _os.environ.get("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "") + if _fg in ("1", "main"): + use_buffer_load = False + + HEAD_ALIGNED = num_heads % BLOCK_M == 0 + heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M + out = _check_out(out, q, kv_lora_rank) + + _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") + + def _pick_splits(hb): + if _sp: + return max(1, int(_sp)) + if kv_splits is not None: + return max(1, int(kv_splits)) + return _decode_num_splits_occ(num_queries, hb, avg_topk, 0.0, BLOCK_K) + + num_splits = _pick_splits(heads_blocks) + + # Short top-k BLOCK_M halving gate (real parallelism at small launches). + _tiles = max(1, math.ceil(avg_topk / BLOCK_K)) if avg_topk > 0 else 1 + if ( + _bm is None + and BLOCK_M == 16 + and num_heads % 8 == 0 + and _tiles <= 2 + and num_queries * heads_blocks * num_splits <= get_num_sms() + ): + _hb8 = (num_heads + 7) // 8 + _ns8 = _pick_splits(_hb8) + if _ns8 >= num_splits and num_queries * _hb8 * _ns8 <= get_num_sms(): + BLOCK_M = 8 + HEAD_ALIGNED = num_heads % BLOCK_M == 0 + heads_blocks = _hb8 + num_splits = _ns8 + + if num_splits > 1: + part_m = torch.empty( + (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device + ) + part_l = torch.empty_like(part_m) + _pab = ( + _os.environ.get("AITER_PA_DECODE_PART_BF16", "1") == "1" and not skip_reduce + ) + part_acc = torch.empty( + (num_queries, num_splits, num_heads, kv_lora_rank), + dtype=torch.bfloat16 if _pab else torch.float32, + device=q.device, + ) + pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) + pa_stride0, pa_stride_s, pa_stride_h = ( + part_acc.stride(0), + part_acc.stride(1), + part_acc.stride(2), + ) + else: + part_m = part_l = part_acc = out # unused placeholders (never dereferenced) + pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 + + # Dequant chunking: split the axis that still has per-lane register repeats. + col_reps = kv_lora_rank // (gather_tw1 * 16) + chunk_axis = 1 if col_reps >= 4 else 0 + _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") + if chunk_axis == 0: + nope_chunk = int(_nc) if _nc else max(1, BLOCK_K // 4) + else: + nope_chunk = int(_nc) if _nc else min(128, kv_lora_rank) + + if chunk_axis == 0 or nope_chunk < kv_lora_rank: + waves_per_eu = 2 + one_wg_per_cu = ( + use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms() + ) + if one_wg_per_cu: + waves_per_eu = 1 + _wp = _os.environ.get("AITER_PA_DECODE_WPEU") + if _wp: + waves_per_eu = int(_wp) + + adaptive_splits = ( + num_splits > 1 and _os.environ.get("AITER_PA_DECODE_ADAPTIVE", "1") == "1" + ) + + _go = _os.environ.get("AITER_PA_DECODE_GRID_ORDER", "qsh") + _ax = {"q": num_queries, "s": num_splits, "h": heads_blocks} + assert sorted(_go) == ["h", "q", "s"], f"bad AITER_PA_DECODE_GRID_ORDER {_go!r}" + # The separated-rope fp8 formats never trace the legacy peeled masked tail + # (kernel static_assert); bf16 keeps the knob. + uni_tile = ( + True + if fmt in ("tensor", "dsmla") + else _os.environ.get("AITER_PA_DECODE_UNI_TILE", "1") == "1" + ) + + grid = tuple(_ax[c] for c in _go) + _pa_decode_sparse_gfx950[grid]( + q, + cache, + alt, + kv_indices, + kv_indptr, + cache, # extra_* segment: unread placeholders (HAS_EXTRA=False) + alt, + kv_indices, + kv_indptr, + attn_sink, + out, + part_m, + part_l, + part_acc, + scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") + scl, + float(softmax_scale), + q.stride(0), + q.stride(1), + out.stride(0), + out.stride(1), + cache.stride(0), + cache.stride(0), + num_rows, + num_rows, + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_EXTRA=False, + HAS_SINK=False, + MAIN_FMT=fmt, + EXTRA_FMT=fmt, + MAIN_BLOCK_SIZE=block_size, + EXTRA_BLOCK_SIZE=block_size, + CS0_ALIGN=cs0_align, + NOPE_DIM=kv_lora_rank, + ROPE_DIM=qk_rope_head_dim, + HEAD_SIZE=kv_lora_rank, + ROPE_SEPARATE=True, + BLOCK_M=BLOCK_M, + BLOCK_K=BLOCK_K, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=HEAD_ALIGNED, + MFMA_K=MFMA_K, + GATHER_TW1=gather_tw1, + LDS_PAD=lds_pad, + NOPE_CHUNK=nope_chunk, + CHUNK_AXIS=chunk_axis, + PART_STORE_CACHE=_os.environ.get("AITER_PA_DECODE_PART_ST", ""), + UNI_TILE=uni_tile, + GRID_ORDER=_go, + MAIN_SPLITS=num_splits, + ADAPTIVE_SPLITS=adaptive_splits, + ASM_DEQ=False, + MAIN_USE_BUFFER_LOAD=use_buffer_load, + EXTRA_USE_BUFFER_LOAD=use_buffer_load, + IDX_BUFFER_LOAD=idx_use_buffer_load, + HAS_INVALID=has_invalid, + FP8_FNUZ=fp8_fnuz, + num_warps=num_warps, + waves_per_eu=waves_per_eu, + ) + + if num_splits == 1: + return out + if skip_reduce: + return part_acc, part_m, part_l + + red_block_m = int(_os.environ.get("AITER_PA_DECODE_REDUCE_BLOCK_M", "1")) + red_warps = int(_os.environ.get("AITER_PA_DECODE_REDUCE_WARPS", "1")) + rgrid = (num_queries, (num_heads + red_block_m - 1) // red_block_m) + _pa_decode_sparse_reduce_gfx950[rgrid]( + part_m, + part_l, + part_acc, + attn_sink, + out, + out.stride(0), + out.stride(1), + pm_stride0, + pm_stride_s, + pa_stride0, + pa_stride_s, + pa_stride_h, + num_heads, + HAS_SINK=False, + HEAD_SIZE=kv_lora_rank, + BLOCK_M=red_block_m, + NUM_SPLITS=num_splits, + HEAD_ALIGNED=num_heads % red_block_m == 0, + ADAPTIVE_SPLITS=adaptive_splits, + PART_LOAD_CACHE=_os.environ.get("AITER_PA_DECODE_PART_LD", ""), + num_warps=red_warps, + ) + return out diff --git a/op_tests/triton_tests/attention/test_sparse_mla_decode.py b/op_tests/triton_tests/attention/test_sparse_mla_decode.py new file mode 100644 index 00000000000..54a3984fdf1 --- /dev/null +++ b/op_tests/triton_tests/attention/test_sparse_mla_decode.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Tests for the separated-rope sparse MLA decode API (gfx950 gluon): +``aiter.ops.triton.attention.sparse_mla_decode.sparse_mla_decode_fwd``. + +Covers the three cache formats (bf16, flat per-tensor fp8, vLLM fp8_ds_mla), +both query-head counts GLM-5 ships (8 = TP8, 16 = TP4), ragged and -1-padded +index streams, split-K on/off, the vLLM cache shapes, and the 64-bit (>2 GB +emulated) gather path. The per-tensor scale-folding is additionally pinned by a +bitwise test against the bf16 path. +""" + +import os + +import pytest +import torch + +from aiter.ops.triton.utils._triton import arch_info + +if arch_info.get_arch() == "gfx950": + from aiter.ops.triton.attention.sparse_mla_decode import sparse_mla_decode_fwd + +FP8_MAX = 448.0 +TINY = 1.1754944e-38 +KV_LORA, ROPE = 512, 64 +D_QK = KV_LORA + ROPE + + +def _skip_unless_gfx950(): + if arch_info.get_arch() != "gfx950": + pytest.skip("sparse_mla_decode_fwd is gfx950-only") + + +# ---------------------------------------------------------------- packers + + +def quantize_flat_fp8(kv): + """GLM-5 production layout: whole row (rope included) fp8 with ONE + per-tensor scale (what concat_and_cache_mla writes with layer._k_scale).""" + scale = (kv.float().abs().amax() / FP8_MAX).clamp_min(TINY).reshape(1) + q8 = (kv.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + return q8.view(torch.uint8), scale.to(torch.float32) + + +def pack_ds_mla(kv, block_size=64): + """vLLM ``fp8_ds_mla``: 656-byte rows [512 fp8 | 4 x f32 per-128 | 64 bf16].""" + T = kv.shape[0] + nb = (T + block_size - 1) // block_size + cache = torch.zeros(nb * block_size, 656, dtype=torch.uint8, device=kv.device) + x = kv[:, :KV_LORA].float().reshape(T, 4, 128) + scale = (x.abs().amax(-1) / FP8_MAX).clamp_min(TINY) + q8 = (x / scale[..., None]).to(torch.float8_e4m3fn).view(torch.uint8) + cache[:T, :KV_LORA] = q8.reshape(T, KV_LORA) + cache[:T, 512:528].view(torch.float32).copy_(scale) + cache[:T, 528:].view(torch.bfloat16).copy_(kv[:, KV_LORA:]) + return cache.view(nb, block_size, 656) + + +def dequant_flat_fp8(u8, scale): + return u8.view(torch.float8_e4m3fn).float() * scale + + +def dequant_ds_mla(cache): + flat = cache.reshape(-1, 656) + x = flat[:, :KV_LORA].view(torch.float8_e4m3fn).float().reshape(-1, 4, 128) + sc = flat[:, 512:528].view(torch.float32) + nope = (x * sc[..., None]).reshape(-1, KV_LORA) + rope = flat[:, 528:].view(torch.bfloat16).float() + return torch.cat([nope, rope], dim=-1) + + +def ragged_indices(C, pool, topk, device, seed, mode): + """mode: 'uniform' (full topk), 'ragged' (true prefix sum, no sentinels -- + the production shape), 'padded' (-1 tail padding, needs has_invalid).""" + g = torch.Generator(device="cpu").manual_seed(seed) + lens = ( + torch.full((C,), topk) + if mode == "uniform" + else torch.randint(1, topk + 1, (C,), generator=g) + ) + rows, flat_lens = [], [] + for c in range(C): + k = int(lens[c]) + sel = torch.randperm(pool, generator=g)[:k] + if mode == "padded": + row = torch.full((topk,), -1, dtype=torch.int64) + row[:k] = sel + rows.append(row) + flat_lens.append(topk) + else: + rows.append(sel) + flat_lens.append(k) + indices = torch.cat(rows) + indptr = torch.zeros(C + 1, dtype=torch.int64) + indptr[1:] = torch.tensor(flat_lens).cumsum(0) + return ( + indices.to(torch.int32).to(device), + indptr.to(torch.int32).to(device), + ) + + +def reference(q, kv_truth, indices, indptr, sm_scale): + """f32 attention over the dequantized cache (tests the kernel, not fp8).""" + C = q.shape[0] + outs = [] + for c in range(C): + sel = indices[indptr[c] : indptr[c + 1]].long() + sel = sel[sel >= 0] + kvs = kv_truth[sel].float() + s = torch.einsum("hd,td->ht", q[c].float(), kvs) * sm_scale + p = torch.softmax(s, dim=-1) + outs.append(torch.einsum("ht,td->hd", p, kvs[:, :KV_LORA])) + return torch.stack(outs) + + +def rel_err(out, ref): + return ((out.float() - ref).abs().max() / ref.abs().max().clamp_min(1e-20)).item() + + +def _build(fmt, C, H, topk, pool, mode, seed=0): + torch.manual_seed(seed) + q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 + kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 + kv[:: max(1, pool // 97)] *= 7.0 # a few heavy-tail tokens + idx, ptr = ragged_indices(C, pool, topk, "cuda", seed + 7, mode) + ks = None + if fmt == "bf16": + cache, truth = kv, kv.float() + elif fmt == "tensor": + cache, ks = quantize_flat_fp8(kv) + truth = dequant_flat_fp8(cache, ks) + else: + cache = pack_ds_mla(kv) + truth = dequant_ds_mla(cache)[:pool] + return q, cache, ks, idx, ptr, truth + + +# ---------------------------------------------------------------- tests + + +@pytest.mark.parametrize("fmt", ["bf16", "tensor", "dsmla"]) +@pytest.mark.parametrize("H", [8, 16]) +@pytest.mark.parametrize( + "topk,mode", + [(2048, "uniform"), (500, "ragged"), (2048, "padded")], + ids=["topk2048", "ragged500", "padded2048"], +) +def test_sparse_mla_decode(fmt, H, topk, mode): + _skip_unless_gfx950() + C, pool = 8, 1 << 16 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, mode) + ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) + out = sparse_mla_decode_fwd( + q, cache, ptr, idx, sm, kv_scale=ks, has_invalid=(mode == "padded") + ) + e = rel_err(out, ref) + assert e < 2e-2, f"{fmt} H={H} topk={topk} {mode}: rel-err {e:.3e}" + + +@pytest.mark.parametrize("kv_splits", [1, None], ids=["splits1", "splitsauto"]) +def test_split_k(kv_splits): + _skip_unless_gfx950() + C, pool, topk = 4, 1 << 15, 2048 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build("tensor", C, 16, topk, pool, "uniform") + ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) + out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks, + kv_splits=kv_splits) + e = rel_err(out, ref) + assert e < 2e-2, f"splits={kv_splits}: rel-err {e:.3e}" + + +def test_per_tensor_scale_folding_bitwise(): + """kv_scale=1.0 on fp8-exact values must equal the bf16 path bit-for-bit: + the K-scale fold into qk_scale and the V-scale fold into p are exact, and + fp8 -> bf16 staging is exact.""" + _skip_unless_gfx950() + torch.manual_seed(3) + C, H, topk, pool = 4, 16, 512, 1 << 15 + sm = D_QK**-0.5 + raw = torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25 + exact = raw.float().to(torch.float8_e4m3fn) + kv_bf16 = exact.to(torch.bfloat16) + cache_u8 = exact.view(torch.uint8) + ks = torch.ones(1, dtype=torch.float32, device="cuda") + q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 + idx, ptr = ragged_indices(C, pool, topk, "cuda", 11, "ragged") + o_t = sparse_mla_decode_fwd(q, cache_u8, ptr, idx, sm, kv_scale=ks) + o_b = sparse_mla_decode_fwd(q, kv_bf16, ptr, idx, sm) + assert torch.equal(o_t.view(torch.int16), o_b.view(torch.int16)) + + +def test_vllm_cache_shapes(): + """The [slots,1,1,R] (asm mla_decode_fwd) and [nb,block,R] (vLLM paged) + views must produce identical results to the flat [slots,R] form.""" + _skip_unless_gfx950() + C, pool, topk = 4, 1 << 14, 512 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, "uniform") + o_flat = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks) + o_4d = sparse_mla_decode_fwd( + q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks + ) + o_paged = sparse_mla_decode_fwd( + q, cache.view(pool // 64, 64, D_QK), ptr, idx, sm, kv_scale=ks + ) + assert torch.equal(o_flat, o_4d) and torch.equal(o_flat, o_paged) + + +@pytest.mark.parametrize("fmt", ["tensor", "dsmla"]) +def test_global_load_path(fmt, monkeypatch): + """Force the 64-bit gather path (production: pools cross buffer_load's 2 GB + offset limit at ~3.7M tokens) and check correctness is unchanged.""" + _skip_unless_gfx950() + monkeypatch.setenv("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "1") + C, pool, topk = 8, 1 << 16, 2048 + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, 16, topk, pool, "uniform") + ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) + out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks) + e = rel_err(out, ref) + assert e < 2e-2, f"{fmt} 64-bit path: rel-err {e:.3e}" From f89f5e25ba0249709b2cd3290df5dd167d6456ae Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 05:42:56 +0000 Subject: [PATCH 26/51] [GLUON][GFX950][MLA] Simplify: concise comments, no env knobs, leaner 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 --- .../gfx950/attention/pa_decode_sparse.py | 621 ++++++------------ .../ops/triton/attention/pa_decode_sparse.py | 8 +- .../ops/triton/attention/sparse_mla_decode.py | 277 +++----- .../attention/test_sparse_mla_decode.py | 152 ++--- 4 files changed, 350 insertions(+), 708 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index dad942a4b9a..1735714ce20 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1,60 +1,34 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Gluon (gfx950) sparse-MLA decode for DeepSeek-V4 *and* GLM-5 (MLA) geometries. - -Restructuring of the previous DSv4-only kernel (last DSv4-only head: -``645a30b0e``); for the existing DSv4/uniform configurations it compiles to -byte-identical .amdgcn (gate: ``compile_check.py`` in the port workspace, -``/app/scripts/glm_gluon_port_8_19``, which keeps the frozen reference). -Two changes: - -1. **Aggregate config objects** (the ``unified_attention_2d.py`` pattern): the - ~40-entry constexpr parameter lists threaded through every helper are replaced by - three aggregates — - - Cfg compile-time geometry, layouts, and behavior knobs (one per launch) - Fmt compile-time cache-format description (one per segment) - Seg runtime segment context: pointers + bounds (one per segment) - - Helpers take ``(cfg, seg, ...)``; the flat constexpr list exists only at the - kernel entry, where the aggregates are built. - -2. **Two geometries under one constexpr, ``ROPE_SEPARATE``**: - - ROPE_SEPARATE=False (DSv4): rope lives *inside* the 512-wide row. One LDS - plane ``kv_smem [BLOCK_K, KV_DIM]``, one QK MFMA chain over KV_DIM, - V = the same plane. Bit-for-bit today's kernel (compile_check.py - diffs the .amdgcn against the aiter reference to prove it). - ROPE_SEPARATE=True (GLM-5): rope is *appended* (512 latent + 64 rope = 576 - QK; V = the 512 latent only). Two LDS planes, both pow-2: - ``kv_smem [BLOCK_K, KV_DIM]`` + ``rope_smem [BLOCK_K, ROPE_DIM]``. - QK = mfma(q_rope, k_rope, mfma(q_nope, k_nope, 0)) — MFMA accumulates - natively so the 576-wide contraction is two chained dots. The PV dot - and the accumulator are unchanged: V is always exactly plane 0. - -Cache formats (``Fmt.KIND``), all sharing one gather/stage/dot pipeline: - - KIND row layout (per token) scales rope - "bf16" [KV_DIM (+ROPE_DIM if separate)] bf16 — in-row - "dsv4" 448 fp8 | 64 bf16 rope (576 B), 8 B E8M0/64 E8M0, block in-plane0 - trailer per 64-token block (aiter fp8_ds_mla-packed DSv4) - "uniform" [KV_DIM] fp8 + separate [tokens, NG] f32 f32 per-64 in-row - "tensor" [KV_DIM (+ROPE_DIM if separate)] fp8, one scalar fp8 tail - per-tensor f32 scale (GLM-5's production asm format; also the - DSv4 per-tensor experiment) - "dsmla" 512 fp8 | 4 f32 | 64 bf16 rope (656 B) f32 per-128 bf16 tail - (vLLM CacheDType "fp8_ds_mla"; requires ROPE_SEPARATE) - -Per-tensor scale handling ("tensor") never touches the tile loop: the K-side scale -folds into the segment's qk_scale (max/exp2 commute with a positive scale) and the -V-side scale multiplies ``p`` — 4 f32 per lane — right before the PV dot, so the -staged tile is a raw fp8→bf16 convert with no multiplies at all. - -Two-loop (SWA + top-k) or a single segment; 2D and 3D (split-K + reduce) share one -kernel. Launchers: ``aiter/ops/triton/attention/pa_decode_sparse.py`` (DSv4 / -uniform-pool, geometry unchanged) and -``aiter/ops/triton/attention/sparse_mla_decode.py`` (GLM-5 / separated-rope MLA). +"""Gluon (gfx950) sparse decode for DeepSeek-V4 and separated-rope MLA (GLM-5). + +One kernel serves two geometries, selected by the ``ROPE_SEPARATE`` constexpr: + + False (DSv4): rope lives inside the pow-2 row. One LDS plane + ``kv_smem [BLOCK_K, HEAD_SIZE]``, one QK MFMA chain, V = the same plane. + True (MLA): rope is appended (``kv_lora_rank`` latent + rope = QK width; + V = the latent only). A second pow-2 LDS plane ``rope_smem`` holds the + K-only rope and the QK contraction chains a second MFMA into the first. + Plane 0 is the entire V in both geometries. + +Cache formats (``Fmt.KIND``), per segment: + + "bf16" full-row bf16 + "dsv4" 448 fp8 | 64 bf16 rope (576 B rows) + 8 B UE8M0-per-64 trailer + "uniform" whole-row fp8 + separate f32 per-64 kv_scales + "tensor" whole-row fp8 + ONE per-tensor f32 scale (k_scale). The scale + never touches the tile loop: K-side folds into qk_scale, + V-side into p. + "dsmla" vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope + (656 B rows); requires ROPE_SEPARATE. + +Compile-time state is bundled into aggregates instead of long parameter lists: +``Cfg`` (geometry, layouts, behavior knobs), ``Fmt`` (per-segment cache +format), ``Seg`` (runtime pointers/bounds). Two-loop (SWA + top-k) or a single +segment; 2D and 3D (split-K + reduce) share one kernel. Launchers: +``aiter/ops/triton/attention/pa_decode_sparse.py`` (DSv4 / uniform pool) and +``aiter/ops/triton/attention/sparse_mla_decode.py`` (separated-rope MLA). """ from triton.experimental import gluon @@ -65,10 +39,9 @@ from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr from aiter.ops.triton.utils.common_utils import strip_annotate -# Triton's default max ignores NaN, which on AMD costs a v_max_f32 x, x, x -# canonicalize per operand before the real compare -- 60 of the 96 v_max in this -# kernel were those no-ops. Nothing here produces NaN (masked lanes are -inf and -# the all-masked row is guarded explicitly), so propagate instead. +# Triton's default max ignores NaN, which on AMD costs a canonicalize per +# operand. Nothing here produces NaN (masked lanes are -inf and the all-masked +# row is guarded), so propagate instead. _MAX_PROP_NAN: gl.constexpr = gl.constexpr(PropagateNan.ALL) @@ -84,17 +57,10 @@ def _rmax(x, axis): @gluon.jit def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): - """Gather rows[i] + col[j] out of a cache. - - ``row`` is the per-token offset in ptr's element units ([BLOCK_K]); ``col`` is - the in-row offset ([W]), always a small compile-time arange. Keeping them apart - is what makes the >2 GB path affordable: buffer_load carries a 32-bit offset - (2 GB cap), so a bigger cache has to gather through 64-bit addresses -- and - adding the pointer to a fully materialized [BLOCK_K, W] offset tensor makes - that a 64-bit add *per element*. Resolving one pointer per token instead costs - BLOCK_K of them, and the leftover column offset is a constant the load can fold - into its immediate field. - """ + """Gather rows[i] + col[j]. ``row`` is the per-token offset in ptr's element + units; ``col`` a small compile-time arange. Keeping them apart resolves one + pointer per token on the 64-bit path (the column offset folds into the + load's immediate) instead of a 64-bit add per element.""" if USE_BUFFER_LOAD: return gl.amd.cdna4.buffer_load( ptr=ptr, @@ -114,9 +80,8 @@ def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=N @gluon.jit def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): - # gfx950's native e4m3 cvt is OCP (float8e4nv). fnuz (float8e4b8) -> f32 has no - # native cvt and lowers to a ~5x software unpack that spills; but fnuz -> bf16 is - # cheap and fp8 -> bf16 is exact (3 mantissa bits), so route fnuz through bf16. + # fnuz -> f32 has no native cvt on gfx950 (software unpack); fnuz -> bf16 is + # cheap and exact (3 mantissa bits), so route fnuz through bf16. if FP8_FNUZ: return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16).to(gl.float32) return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) @@ -124,8 +89,7 @@ def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): @gluon.jit def _fp8_to_bf16(x_u8, FP8_FNUZ: gl.constexpr): - # fp8 -> bf16 is exact (3 mantissa bits into 8), so the per-tensor mode can - # stage the raw code points without any scale arithmetic at all. + # Exact: fp8's 3 mantissa bits fit bf16's 8. if FP8_FNUZ: return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16) return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.bfloat16) @@ -144,22 +108,13 @@ def _scale_load( MASKED: gl.constexpr, OTHER: gl.constexpr, ): - """Gather the NG per-group scale entries of each token and broadcast them out - to W_FULL columns in registers. - - The obvious spelling -- index the full row with ``offs_full // GROUP`` -- builds - a [BLOCK_K, W_FULL] pointer tensor for a value that has only NG distinct entries - per row. On the buffer_load path the identical 32-bit offsets CSE away; on the - 64-bit path they do not, and the wide form costs 6x the loads. Gather NG wide - instead and broadcast: ``scl_l`` is the dim-2 slice of a 3-D layout picked so - that reshaping [BLOCK_K, NG, GROUP] back to [BLOCK_K, W_FULL] lands exactly on - ``gather_l`` -- so the broadcast is a register rename (assert_trivial proves it - at compile time; GLM5_PORT_PLAN flagged re-verifying this for NG != 8). - - Element type follows ``ptr``: u8 E8M0 bytes for the DSv4 packed cache - (OTHER=127 -> 2^0), f32 group scales for fp8_ds_mla (OTHER=0.0 -> masked - lanes dequant to 0). - """ + """Gather the NG per-group scales of each token and broadcast to W_FULL + columns. Indexing the full row with ``offs // GROUP`` would build a + [BLOCK_K, W_FULL] pointer tensor for NG distinct values — redundancy only + buffer_load's 32-bit offsets CSE away. ``scl_l`` is chosen so the broadcast + lands on ``gather_l`` as a register rename (assert_trivial proves it). + Element type follows ``ptr``: u8 E8M0 bytes (OTHER=127 -> 2^0) or f32 + (OTHER=0.0 -> masked lanes dequant to 0).""" cols = gl.arange(0, NG, layout=gl.SliceLayout(0, scl_l)) rows = gl.convert_layout(row, gl.SliceLayout(1, scl_l)) if MASKED: @@ -199,14 +154,9 @@ def _scale_load( @gluon.jit def _split2(x): - """Contiguous register split along dim 1: [A, B] -> two [A, B//2]. - - x0 takes columns [0, B//2), x1 takes [B//2, B). Both halves come back in the - input's own layout, so with warps tiling dim 0 (see ``Cfg.gather_l``) the column - direction is a pure per-lane register repeat and the split is a rename -- no - cross-lane traffic. ``assert_trivial`` makes a non-free split a compile error - rather than a silent LDS round-trip. - """ + """Register split along dim 1: [A, B] -> two [A, B//2] in the input's own + layout. Free only while the column direction is a per-lane register repeat; + ``assert_trivial`` turns anything else into a compile error.""" layout: gl.constexpr = x.type.layout x_r = x.reshape(x.shape[0], 2, x.shape[1] // 2).permute(0, 2, 1) x0, x1 = gl.split(x_r) @@ -217,12 +167,8 @@ def _split2(x): @gluon.jit def _split2_dim0(x): - """Contiguous register split along dim 0: [A, B] -> two [A//2, B]. - - The dim-0 counterpart of _split2, for the row-wide gather layout where dim 1 - is spent entirely on threads (one instruction = whole token rows) and the - per-lane register repeats live on dim 0 instead. - """ + """Dim-0 counterpart of _split2, for the row-wide gather layout where the + per-lane register repeats live on dim 0.""" layout: gl.constexpr = x.type.layout x_r = x.reshape(2, x.shape[0] // 2, x.shape[1]).permute(1, 2, 0) x0, x1 = gl.split(x_r) @@ -241,16 +187,9 @@ def _split_ax(x, AXIS: gl.constexpr): # fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. Emitted as two -# separate single-output asm calls, NOT as one two-output blob: a single-instruction -# asm reads all its sources before writing dst, so no output can clobber an input and -# neither needs early clobber. The blob form does -- without `=&v` on its first output -# the allocator may give it the scale's register and the second convert then reads a -# clobbered scale, wrong in exactly the elements it produces -- and that forced -# liveness costs spills: 21 st/28 ld against 11/12 here, on a kernel with zero -# headroom at the 256-VGPR cap. The split form costs +184 instructions (each call -# re-reads the sources) and is worth it: 0.970-0.977x against the blob's 0.983-1.017x -# across C=16/64/128 at short top-k, and it removes the reduce knock-on that the -# blob's scratch traffic caused by evicting the split-K partials from L2. +# single-output asm calls, not one two-output blob: a single instruction reads +# all sources before writing dst, so no output needs early clobber — the blob's +# forced `=&v` liveness spills on a kernel with no VGPR headroom. _DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") _DEQ_HI: gl.constexpr = gl.constexpr( "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" @@ -260,15 +199,12 @@ def _split_ax(x, AXIS: gl.constexpr): @gluon.jit def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): - """[BLOCK_K, W8/2] int16 (4 packed fp8 per 32 bits) + raw E8M0 byte -> [BLOCK_K, W8] bf16. - - The scale operand is read as E8M0 -- the hardware takes bits [30:23] and nothing - else. Those are bits [14:7] of the operand's high half, so a 16-bit `e << 7` in - both halves of the register puts the exponent exactly where an f32 `e << 23` - would, and an i16 operand costs ONE register at pack=2 where f32 cost two. - Converting to f32 with exp2 also works but is a longer chain for nothing. Because the scale is a power of two this is bit-identical to the f32 - chain: fp8 -> bf16 is exact (3 mantissa bits into 8) and scaling by 2^k stays exact. - """ + """[BLOCK_K, W8/2] int16 (4 packed fp8) + raw E8M0 byte -> [BLOCK_K, W8] bf16. + + The hardware reads only bits [30:23] of the scale operand — bits [14:7] of + its high half — so a 16-bit ``e << 7`` lands the exponent where an f32 + ``e << 23`` would, in ONE register. Power-of-two scale keeps this + bit-identical to the unfused convert+multiply chain.""" sc16 = e_u8.to(gl.uint16) << 7 lo = gl.inline_asm_elementwise( _DEQ_LO, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 @@ -276,22 +212,19 @@ def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): hi = gl.inline_asm_elementwise( _DEQ_HI, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 ) - # lo carries fp8 elements 4i+{0,1} and hi 4i+{2,3}, so the flat order is - # 4i + 2*lohi + s. A lane's 8 int16 are the same 16-byte run the u8 gather used, - # so this whole interleave stays inside one lane -- assert_trivial proves it. + # Flat order is 4i + 2*lohi + s; a lane's int16s are the same 16-byte run + # the gather used, so the interleave stays lane-local (assert_trivial). W16: gl.constexpr = W8 // 2 lo3 = lo.reshape(lo.shape[0], W16 // 2, 2) hi3 = hi.reshape(hi.shape[0], W16 // 2, 2) both = gl.join(lo3, hi3).permute(0, 1, 3, 2).reshape(lo.shape[0], W8) - # Back to the byte-gather layout, or the kv_smem store sees the join/permute layout - # and lowers to narrow ds_writes (112 ds_write_b128 -> 28, and ~1.5x slower). + # Back to the gather layout, or the kv_smem store lowers to narrow ds_writes. return gl.convert_layout(both, out_l, assert_trivial=True) # --------------------------------------------------------------------------- -# Aggregates. Compile-time state lives in Cfg (shared) and Fmt (per segment); -# runtime pointers/bounds live in Seg. Built once at the kernel entry, passed -# in place of the former ~40-parameter lists. +# Aggregates: Cfg (shared compile-time), Fmt (per-segment compile-time), +# Seg (per-segment runtime). Built once at the kernel entry. # --------------------------------------------------------------------------- @@ -306,9 +239,8 @@ class Cfg: KV_DIM: gl.constexpr # plane-0 width = LDS tile width = V width (pow-2) ROPE_DIM: gl.constexpr # plane-1 width when ROPE_SEPARATE, else the bf16 # tail width inside plane 0 (DSv4 packed) - ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0 (DSv4). True: rope is - # a K-only second plane (GLM-5); QK contracts over - # KV_DIM + ROPE_DIM via two chained MFMAs. + ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0. True: K-only second + # plane; QK contracts over KV_DIM + ROPE_DIM. QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) MFMA_K: gl.constexpr NUM_WARPS: gl.constexpr @@ -389,23 +321,12 @@ def __init__( self.p_layout = gl.constexpr(gl.DotOperandLayout(0, self.pv_layout, KW)) self.v_layout = gl.constexpr(gl.DotOperandLayout(1, self.pv_layout, KW)) - # 16 uint8 = 128-bit fp8 gather loads + # 16 uint8 per thread = 128-bit gather loads. Warps tile dim 0 so the + # per-lane slot vector stays short and the column direction keeps + # register repeats (what the chunked dequant splits for free). + # GATHER_TW1 = threads spent on the head dim: 32 requests a whole + # 512 B token row per instruction, at the cost of a longer slot vector. GSPT = 16 - # Warps on dim 1 (columns) vs dim 0 (rows). Coalescing is identical either way - # -- a row's 128 contiguous bytes always come from 8 threads -- but the row - # index vector lives in SliceLayout(1, gather_l), so warps-on-dim-1 replicates - # it across all NUM_WARPS and every lane carries BLOCK_K/8 slots, while - # warps-on-dim-0 carries BLOCK_K/(8*NUM_WARPS). Dim 0 also makes the column - # direction a pure per-lane register repeat, so a chunked dequant can split it - # with a trivial (register-renaming) layout convert. - # GATHER_TW1 = threads spent on the head dim. Each thread loads 16 B, so - # TW1=8 requests 128 B of a token row per instruction and TW1=32 requests the - # whole 512 B row -- for a scattered gather that is one request instead of - # four quarters. The cost is that the row-index vector lives in - # SliceLayout(1, gather_l), so a wider TW1 leaves fewer dim-0 thread slots and - # every lane carries more slots (BLOCK_K*TW1/64 of them). Whichever dim keeps - # per-lane register repeats is the one the chunked dequant can split - # (CHUNK_AXIS). self.gather_l = gl.constexpr( gl.BlockedLayout( size_per_thread=[1, GSPT], @@ -414,9 +335,8 @@ def __init__( order=[1, 0], ) ) - # Warps tile dim 0 here, one warp already covers all 64 RoPE columns, so putting - # them on dim 1 (as the NoPE gather does) overshoots 4x and every warp re-gathers - # the same tile. Worth ~10% on packed fp8. + # One warp already covers all rope columns; warps on dim 1 would + # re-gather the same tile NUM_WARPS times. self.gather_rope_l = gl.constexpr( gl.BlockedLayout( size_per_thread=[1, 8], @@ -425,8 +345,8 @@ def __init__( order=[1, 0], ) ) - # Same tiling as gather_l but 2-byte elements: half the per-thread run, so a lane - # still covers the same 16 bytes and the lo/hi interleave stays lane-local. + # gather_l with 2-byte elements: same 16-byte per-lane run, half the + # columns (for the fused-dequant int16 view). self.gather16_l = gl.constexpr( gl.BlockedLayout( size_per_thread=[1, GSPT // 2], @@ -444,19 +364,14 @@ def __init__( order=[1, 0], ) ) - # LDS pad after every row. Row pitch (512 + LDS_PAD) bf16 sets which banks a - # transposed K read (ds_read_b64_tr_b16 walks down a column) lands on: - # bank = (row * pitch_dwords) mod 32, so pitch_dwords mod 32 == 4 at PAD=8 means - # 32 lanes share 8 banks. Kept as a knob because the conflict share is high - # (46%) but a high sub-metric ratio is not proof of a bottleneck. + # Row pitch (KV_DIM + LDS_PAD) decides which banks the transposed K + # read (walks down a column) lands on. self.kv_shared = gl.constexpr( gl.PaddedSharedLayout.with_identity_for( [[KV_DIM, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] ) ) - # Plane 1 (K-only rope) when ROPE_SEPARATE; same padding rule as plane 0 - # (its K read is the same transposed ds_read pattern). Built regardless -- - # a dead constexpr field costs nothing when the plane is never allocated. + # Plane 1 (K-only rope) when ROPE_SEPARATE; dead constexpr otherwise. self.rope_shared = gl.constexpr( gl.PaddedSharedLayout.with_identity_for( [[ROPE_DIM, LDS_PAD]], [BLOCK_K, ROPE_DIM], [1, 0] @@ -488,7 +403,7 @@ class Fmt: ROPE_U16_OFF: gl.constexpr # bf16-view offset of the rope tail in a row SCL_TRAILER_U8: gl.constexpr # dsv4: scale bytes per token in the block trailer TOK_F32: gl.constexpr # dsmla: f32-view stride per token row - SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the 4 group scales + SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the group scales TOK_EL: gl.constexpr # flat formats: cache elements per token row @gluon.constexpr_function @@ -520,20 +435,14 @@ def __init__( NG = KV_DIM // GROUP self.GROUP = gl.constexpr(GROUP) self.NG = gl.constexpr(NG) - # 3-D companion of gather_l used by _scale_load: dim 1 carries the NG scale - # groups (one thread each) and dim 2 the GROUP columns inside a group, so - # that reshaping [BLOCK_K, NG, GROUP] back to 2-D reproduces gather_l - # exactly. Legal only when the group's columns fill whole threads, i.e. - # GSPT * (TW1 // NG) is the group width; otherwise fall back to the wide - # (redundant) gather. + # 3-D companion of gather_l for _scale_load: dim 1 carries the NG scale + # groups, dim 2 the columns inside a group, so the [BLOCK_K, NG, GROUP] + # -> 2-D reshape reproduces gather_l exactly. Legal only when a group's + # columns fill whole threads. GSPT = 16 TW1 = cfg.GATHER_TW1.value NARROW_SCALE = TW1 % NG == 0 and GSPT * (TW1 // NG) == GROUP self.NARROW_SCALE = gl.constexpr(NARROW_SCALE) - # ...and only worth it on the 64-bit path. With buffer_load the wide gather's - # identical 32-bit offsets CSE to the same few loads, so the narrow form - # only adds a layout convert per tile (+3.4% at extra=1024). Without it they - # do not CSE and the wide form costs 6x the loads. scl_l3 = gl.BlockedLayout( size_per_thread=[1, 1, GSPT], threads_per_warp=[ @@ -546,9 +455,8 @@ def __init__( ) self.scl_l = gl.constexpr(gl.SliceLayout(2, scl_l3)) - # Packed-row constants. dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] with an - # 8-byte-per-token E8M0 trailer after the block; dsmla row: - # [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] all inline. + # dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] + 8 B UE8M0 per token after + # the block; dsmla row: [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] inline. if KIND == "dsv4": TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 elif KIND == "dsmla": @@ -563,8 +471,8 @@ def __init__( self.SCL_TRAILER_U8 = gl.constexpr(8) self.TOK_F32 = gl.constexpr(TOK_U8 // 4) self.SCL_F32_OFF = gl.constexpr(KV_DIM // 4) - # Flat formats gather in cache elements (bf16 or fp8-byte): the row is - # KV_DIM wide, plus the appended rope when the geometry separates it. + # Flat formats gather in cache elements: KV_DIM wide, plus the appended + # rope when the geometry separates it. self.TOK_EL = gl.constexpr( KV_DIM + (ROPE_DIM if cfg.ROPE_SEPARATE.value and KIND != "uniform" else 0) ) @@ -573,16 +481,15 @@ def __init__( @aggregate @strip_annotate class Seg: - """Runtime context of one segment: cache pointers, index list, bounds. - - Pointer roles by format ("--" = unused, carries a harmless duplicate): + """Runtime context of one segment. Pointer roles by format ("--" = unused + duplicate): KIND cache_ptr alt_ptr scl_ptr bf16 -- bf16 cache -- dsv4 u8 cache bf16 view of the cache -- uniform u8 cache f32 per-64 kv_scales -- tensor u8 cache -- f32 scalar k_scale - dsmla u8 cache bf16 view (rope tail) f32 view (group scales) + dsmla u8 cache bf16 view (rope tail) f32 view (scales) """ fmt: Fmt @@ -615,22 +522,15 @@ def __init__( @gluon.jit def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): - """Dequant one fp8 slab and write it straight to kv_smem[:, off:off+W]. - - Keep the dequant in f32: gfx950 has no bf16 multiply, so a bf16 path lowers to - an emulated mul and doubles the loop. ``sc`` is the raw per-group exponent - byte (dsv4 packed, UE8M0), an f32 scale (uniform pool / fp8_ds_mla), or unused - ("tensor": the per-tensor scale is folded outside the tile loop, so staging is - a bare fp8 -> bf16 convert). - """ + """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 + (gfx950 has no bf16 multiply). ``sc``: raw UE8M0 byte (dsv4), f32 scale + (uniform/dsmla), or unused ("tensor": bare fp8 -> bf16 convert).""" if fmt.ASM_DEQ: # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. W8: gl.constexpr = x_u8.shape[1] * 2 - # Adjacent fp8 columns share a scale (groups are 64 wide, so even), which makes - # dropping every other broadcast column exact. + # Adjacent fp8 columns share a scale (groups are even), so dropping + # every other broadcast column is exact. s_even, _ = gl.split(sc.reshape(sc.shape[0], sc.shape[1] // 2, 2)) - # split leaves a SliceLayout of a 3-D parent; the asm needs both operands in - # the int16 gather's own layout. s_even = gl.convert_layout(s_even, x_u8.type.layout) val = _deq_asm(x_u8.to(gl.int16, bitcast=True), s_even, W8, cfg.gather_l) if AXIS == 1: @@ -654,15 +554,10 @@ def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): @gluon.jit def _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt): - """Dequant a whole gathered fp8 tile into kv_smem in NOPE_CHUNK-sized pieces - along CHUNK_AXIS (0 = rows, 1 = columns). - - The f32 expansion is 4x the fp8 tile (a [64, 512] tile is 128 f32 VGPRs/lane), - so materializing it in one shot is what pins 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 -- so only a 1/pieces fraction of it is ever - live. The splits themselves are register renames (see ``_split2``). - """ + """Dequant a gathered fp8 tile into kv_smem in NOPE_CHUNK-sized pieces + along CHUNK_AXIS (0 = rows, 1 = columns). The f32 expansion is 4x the fp8 + tile; chunking keeps only 1/pieces of it live (each piece's converts feed + its ds_writes and die). The splits are register renames (see _split2).""" AXIS: gl.constexpr = fmt.CHUNK_AXIS NOPE_CHUNK: gl.constexpr = fmt.NOPE_CHUNK W: gl.constexpr = x_u8.shape[1] if AXIS == 1 else x_u8.shape[0] @@ -716,23 +611,16 @@ def _slots( MASKED: gl.constexpr, UNI_TILE: gl.constexpr = False, ): - # Returns in whatever layout k_pos carries. Called once per gather layout: NoPE - # and RoPE tile their warps differently, and re-loading this tiny broadcast - # vector is cheaper than a cross-lane convert between the two. - # 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 is also the largest single source - # of exec-mask branching in the kernel (20 of 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. + """Index-list read -> (block, pos, valid), in whatever layout k_pos carries. + A masked gl.load predicates on exec while a masked buffer_load folds the + mask into the offset, so the unmasked paths clamp the read in-range and + mask the *score* instead (UNI_TILE); -1 sentinels are handled the same way.""" indices_ptr = seg.indices_ptr seg_start = seg.seg_start BLOCK_SIZE: gl.constexpr = seg.fmt.BLOCK_SIZE HAS_INVALID: gl.constexpr = cfg.HAS_INVALID IDX_BUFFER_LOAD: gl.constexpr = cfg.IDX_BUFFER_LOAD if MASKED: - # Legacy peeled tail (UNI_TILE=0, and the bf16 path). Predicating the read is - # what makes this expensive: a masked gl.load branches on exec. in_range = k_pos < hi if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load( @@ -743,11 +631,7 @@ def _slots( valid = in_range & (slot >= 0) & (slot < num_rows) slot = gl.where(valid, slot, 0) else: - # UNI_TILE lets the partial tile ride this same body: clamp the index-list read - # into the segment rather than predicating it, so every lane reads a genuine - # in-segment index, and add the range test to `valid` -- which becomes a -inf - # score mask, exactly how -1 sentinels are already handled here. hi >= 1 - # whenever this runs (_gather_full is guarded by n_full > 0, so hi > lo >= 0). + # hi >= 1 whenever UNI_TILE runs (guarded by n_full > 0). off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off) @@ -763,9 +647,8 @@ def _slots( @gluon.jit def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): - """QK scores for one tile. ROPE_SEPARATE=False is a single MFMA chain over - plane 0; ROPE_SEPARATE=True chains a second MFMA over the rope plane (MFMA - accumulates natively, so the KV_DIM + ROPE_DIM contraction is two dots).""" + """QK scores for one tile; ROPE_SEPARATE chains a second MFMA over the rope + plane (MFMA accumulates natively, so KV_DIM + ROPE_DIM is two dots).""" k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] S = gl.amd.cdna4.mfma( q_dot, @@ -779,8 +662,8 @@ def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): # --------------------------------------------------------------------------- -# The prefetched fp8 pipeline: _gather_full issues tile N+1's loads while -# _qkpv stages + dots tile N. +# Prefetched fp8 pipeline: _gather_full issues tile N+1's loads while _qkpv +# stages and dots tile N. # --------------------------------------------------------------------------- @@ -796,21 +679,14 @@ def _gather_full( k_rng_slot, k_rng_rope, ): - """Gather one full fp8 tile. Split from the LDS-write/MFMA so the gather issues - an iteration early. - - The prefetch is carried across the MFMA in **raw fp8**, not dequantized bf16: - the loop keeps two tiles in flight, and a [BLOCK_K, KV_DIM] tile is 32 - VGPRs/lane as u8 but 64 as bf16, so dequantizing here would burn an extra 64 - VGPRs on loop-carried state alone. The consumer dequants (chunked) instead. - - Returns (x, sc, k_rope, valid); unused slots carry a harmless duplicate that - DCE removes ("tensor" has no scale vector, "uniform" no rope side-channel). - """ + """Gather one full fp8 tile, split from the LDS-write/MFMA so it issues an + iteration early. The prefetch is carried in raw fp8 — dequantized here it + would double the loop-carried registers; the consumer dequants chunked. + Returns (x, sc, k_rope, valid); unused slots carry a duplicate DCE removes.""" fmt = seg.fmt cs0 = seg.cs0 if not fmt.USE_BUFFER_LOAD: - cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets bg, pg, valid = _slots( cfg, seg, @@ -828,17 +704,15 @@ def _gather_full( sc = _cache_load( seg.alt_ptr, bg * NGRP, offs_full // 64, fmt.USE_BUFFER_LOAD ) - k_rope = x_u8 # unused for "uniform" (rope slice-store skipped) -> DCE'd + k_rope = x_u8 # no rope side-channel -> DCE'd elif fmt.KIND == "tensor": - # Flat fp8 rows, one per-tensor scale for the whole cache. No scale work - # here at all: the K-side scale is folded into this segment's qk_scale and - # the V-side scale into p (see _qkpv), so staging is a bare convert. + # Per-tensor scale is folded outside the loop (qk_scale / p), so this + # is a bare gather; the K-only rope tail follows when separated. x_u8 = _cache_load( seg.cache_ptr, bg * cs0 + pg * fmt.TOK_EL, offs_full, fmt.USE_BUFFER_LOAD ) sc = x_u8 # no scale vector -> DCE'd if cfg.ROPE_SEPARATE: - # K-only rope tail: fp8 columns [KV_DIM, KV_DIM + ROPE_DIM) of the row. bgr, pgr, _ = _slots( cfg, seg, @@ -857,10 +731,11 @@ def _gather_full( else: k_rope = x_u8 # rope lives inside plane 0 -> DCE'd elif fmt.KIND == "dsmla": - # vLLM fp8_ds_mla row: [KV_DIM fp8 | NG f32 scales | ROPE_DIM bf16]. nope_row = bg * cs0 + pg * fmt.TOK_U8 scl_row = bg * (cs0 // 4) + pg * fmt.TOK_F32 + fmt.SCL_F32_OFF - # f32 group scales first, bulk fp8 after (same vmcnt-FIFO argument as dsv4). + # Scales issue before the bulk fp8: vmcnt is one in-order FIFO, so + # issued last they would stall the first dequant piece behind every + # data load as well. if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: sc = _scale_load( seg.scl_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, @@ -889,11 +764,7 @@ def _gather_full( else: # "dsv4" nope_row = bg * cs0 + pg * fmt.TOK_U8 scl_row = bg * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pg * fmt.SCL_TRAILER_U8 - # UE8M0 scales first, bulk fp8 after. vmcnt is one in-order FIFO, so a - # wait can only name "at most N outstanding", never a specific load: - # issuing the scales last would make the first dequant piece wait behind - # every data load as well. Issued first, the scales are covered by the - # wait that the first piece's own data needs anyway. + # Scales first (vmcnt FIFO; see the dsmla branch). if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: sc = _scale_load( seg.cache_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, @@ -904,13 +775,9 @@ def _gather_full( seg.cache_ptr, scl_row, offs_full // 64, fmt.USE_BUFFER_LOAD ) if fmt.ASM_DEQ: - # 2-byte elements so the asm gets <2 x i16> = 4 packed fp8 in one VGPR out - # of a single dword load. Same bytes, same 16-B per-lane run, half the - # columns; alt_ptr is already the 2-byte view of this cache and the - # RoPE gather below uses the same cs0//2 / TOK_U16 addressing. - # The row vector comes from _slots in SliceLayout(1, gather_l); the i16 - # columns live in gather16_l, and a broadcast needs one parent layout. The - # two layouts share their dim-0 tiling, so this convert is a rename. + # 2-byte elements: <2 x i16> = 4 packed fp8 per VGPR out of one + # dword load. Same byte as nope_row (both even), addressed through + # the bf16 view; the layout convert is a rename (shared dim-0 tiling). row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) @@ -939,10 +806,10 @@ def _gather_full( @gluon.jit def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): - """Write one prefetched tile into the LDS plane(s). Plane 0 is always the full - KV_DIM-wide dequant ("dsv4" gathers KV_DIM bytes too: the last 64 are the bf16 - rope read as garbage fp8, dequanted, and then overwritten by the real rope - slice-store below -- that keeps the gather pow-2 wide).""" + """Write one prefetched tile into the LDS plane(s). Plane 0 is always the + full KV_DIM-wide dequant — "dsv4" gathers KV_DIM bytes too (the last 64 are + bf16 rope read as garbage fp8, overwritten by the slice-store below), which + keeps the gather pow-2 wide.""" fmt = seg.fmt _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) if fmt.KIND == "dsv4": @@ -976,22 +843,20 @@ def _qkpv( k_start=0, seg_hi=0, ): - """Dequant a prefetched fp8 tile into LDS, then QK -> softmax -> PV. When - HAS_INVALID, mask the columns of -1-sentinel slots (``valid`` from - _gather_full) to -inf.""" + """Stage a prefetched fp8 tile into LDS, then QK -> softmax -> PV.""" neg_inf = float("-inf") _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem) S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) - # UNI_TILE folds the tile's range test into `valid`, so the score mask is no - # longer only about -1 sentinels: it is what makes the partial last tile correct. + # UNI_TILE folds the tile's range test into the score mask, which is what + # makes the partial last tile correct without a peeled masked body. COL_MASK: gl.constexpr = cfg.HAS_INVALID or cfg.UNI_TILE NEED_MASK: gl.constexpr = COL_MASK or (not cfg.HEAD_ALIGNED) if NEED_MASK: if COL_MASK: if cfg.UNI_TILE and not cfg.HAS_INVALID: - # Build the range mask directly in the MFMA layout. Converting the - # slot-layout `valid` vector instead costs a cross-lane layout change - # per tile, which at 16 tiles ate the whole saving. + # Build the range mask directly in the MFMA layout; converting + # the slot-layout `valid` vector costs a cross-lane convert per + # tile. col_mask = ( k_start + gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(0, cfg.qk_layout)) @@ -1013,14 +878,10 @@ def _qkpv( :, None ] S = gl.where(col_mask, S, neg_inf) - # m_i is carried in base-2 exponent space, hence qk_scale. max commutes with a - # positive scale (qk_scale = scale * log2(e) [* k_scale for "tensor"]; any real - # attention scale is > 0), so scale the row max -- one value per lane -- rather - # than every element of S. That leaves `S * qk_scale - m_new` as the only use of - # the scaled S, which lowers to v_fma_f32 where the split form needed a separate - # multiply over the tile. -inf * positive = -inf, so masked columns stay masked, - # and the row max is the same product bit for bit; only the fused subtract - # differs, one rounding instead of two -- measured <= 0.23 of a bf16 output ulp. + # Online softmax in the base-2 exponent domain; m_i carries qk_scale + # (= scale * log2e [* k_scale for "tensor"]). max commutes with a positive + # scale, so scale the row max instead of every element — the remaining + # `S * qk_scale - m_new` lowers to FMA, and -inf columns stay -inf. m_block = _rmax(S, 1) * qk_scale m_new = _max2(m_i, m_block) m_new = gl.where(m_new > neg_inf, m_new, 0.0) @@ -1028,10 +889,8 @@ def _qkpv( alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(cfg.v_layout) - # "tensor": V was staged as raw fp8 code points, so apply the per-tensor scale - # on the small side -- p is [BLOCK_M, BLOCK_K], 4 f32 per lane, vs 32 for the - # accumulator. l_new stays scale-free (it normalizes p, not p*V), so - # out = sum(p * v_scale * V_raw) / l is exactly sum(p * (v_scale * V_raw)) / l. + # "tensor": V was staged as raw fp8 code points; apply the per-tensor scale + # on the small side (p), leaving l scale-free — out = sum(p*s*V)/l exactly. if seg.fmt.KIND == "tensor": p = p * v_scale p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) @@ -1042,8 +901,8 @@ def _qkpv( # --------------------------------------------------------------------------- -# The non-prefetched path: bf16 segments, and the peeled masked tail when -# UNI_TILE is off (dsv4/uniform only; the new fp8 formats require UNI_TILE). +# Non-prefetched path: bf16 segments, and the peeled masked tail when UNI_TILE +# is off (dsv4/uniform only; tensor/dsmla require UNI_TILE). # --------------------------------------------------------------------------- @@ -1070,15 +929,14 @@ def _decode_tile( k_rng_rope, MASKED: gl.constexpr, ): - """One KV tile -> online-softmax update. MASKED=False (peeled full tiles) drops - the in-range / gather / score masking; MASKED=True (the tail) keeps it. When - HAS_INVALID, full tiles also clamp -1 sentinels in-bounds for the gather and - mask their scores to -inf (matching the tail's slot-validity handling).""" + """One KV tile -> online-softmax update. MASKED=True is the peeled tail + (fully predicated); full tiles clamp -1 sentinels and mask scores when + HAS_INVALID.""" neg_inf = float("-inf") fmt = seg.fmt cs0 = seg.cs0 if not fmt.USE_BUFFER_LOAD: - cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets (see _cache_load) + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets block_idx, pos, valid1d = _slots( cfg, seg, @@ -1093,7 +951,6 @@ def _decode_tile( valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather_l)) if fmt.KIND == "uniform": - # uniform pool: one fp8 gather over the whole head + separate fp32 scales. NGRP: gl.constexpr = cfg.KV_DIM // 64 kv_row = block_idx_g * cs0 + pos_g * cfg.KV_DIM scl_row = block_idx_g * NGRP @@ -1110,13 +967,12 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, - ) # uniform pool: fp32 scales, left wide (see _scale_load) + ) else: x_u8 = _cache_load(seg.cache_ptr, kv_row, offs_full, fmt.USE_BUFFER_LOAD) sc = _cache_load(seg.alt_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) elif fmt.KIND == "dsv4": - # DSv4 packed fp8_ds_mla: NoPE fp8 + embedded UE8M0 + separate RoPE-bf16. nope_row = block_idx_g * cs0 + pos_g * fmt.TOK_U8 scl_row = ( block_idx_g * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pos_g * fmt.SCL_TRAILER_U8 @@ -1134,9 +990,6 @@ def _decode_tile( mask=valid_g[:, None], other=127, ) if fmt.ASM_DEQ: - # Same byte as nope_row, addressed in 2-byte units: one shift on a - # live vector, and block_idx_g * cs0 stays common with the scale - # gather. Exact -- cs0 and TOK_U8 are both even. row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) @@ -1161,9 +1014,6 @@ def _decode_tile( else: exps = _cache_load(seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) if fmt.ASM_DEQ: - # Same byte as nope_row, addressed in 2-byte units: one shift on a - # live vector, and block_idx_g * cs0 stays common with the scale - # gather. Exact -- cs0 and TOK_U8 are both even. row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) @@ -1196,7 +1046,7 @@ def _decode_tile( else: k_rope = _cache_load(seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD) kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) - else: # "bf16" (the new fp8 formats require UNI_TILE, so they never come here) + else: # "bf16" (tensor/dsmla require UNI_TILE, so they never come here) kv_row2 = block_idx_g * cs0 + pos_g * fmt.TOK_EL if MASKED: kv = _cache_load( @@ -1207,7 +1057,6 @@ def _decode_tile( kv = _cache_load(seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD) kv_smem.store(kv) if cfg.ROPE_SEPARATE: - # Appended bf16 rope columns [KV_DIM, KV_DIM + ROPE_DIM) -> plane 1. block_idx_gr, pos_gr, valid_gr = _slots( cfg, seg, @@ -1229,9 +1078,7 @@ def _decode_tile( rope_smem.store(k_rope) S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) - # exp2 softmax: qk_scale folds in log2(e) so we hit the HW exp2 directly. - # Running max stays in raw-score space; masked cols (-inf) give exp2=0. - COL_VALID: gl.constexpr = MASKED or cfg.HAS_INVALID # valid1d defined in both + COL_VALID: gl.constexpr = MASKED or cfg.HAS_INVALID NEED_MASK: gl.constexpr = COL_VALID or (not cfg.HEAD_ALIGNED) if NEED_MASK: if COL_VALID: @@ -1245,18 +1092,13 @@ def _decode_tile( ] & col_mask ) - else: # not HEAD_ALIGNED and no col invalidity -> head mask only + else: col_mask = gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ :, None ] S = gl.where(col_mask, S, neg_inf) - # Online softmax in the base-2 exponent domain: fold qk_scale*log2(e) into S - # once, right out of the MFMA, and carry m_i already scaled. The alternative - # (raw running max, scale at use) needs qk_scale on both the S tile and the two - # m vectors every iteration; here the per-element work in the loop is one - # subtract feeding exp2, and the whole epilogue -- sink combine and partial - # store, both of which want base-2 m -- stops converting. + # exp2 softmax with qk_scale folded in; masked cols (-inf) give exp2 = 0. S = S * qk_scale m_block = _rmax(S, 1) m_new = _max2(m_i, m_block) @@ -1293,7 +1135,6 @@ def _process_segment( rope_smem, ): offs_full = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.gather_l)) - # int16 column offsets for the asm dequant's 2-byte gather (half the columns). offs_full16 = gl.arange( 0, cfg.KV_DIM // 2, layout=gl.SliceLayout(0, cfg.gather16_l) ) @@ -1301,13 +1142,12 @@ def _process_segment( k_rng_slot = gl.arange(0, cfg.BLOCK_K, layout=cfg.slot_l) k_rng_rope = gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(1, cfg.gather_rope_l)) - # Peel the (possibly partial) last tile: [lo, hi_full) are full BLOCK_K tiles - # whose slots are all valid -> mask-free. Only the peeled tail carries masking. + # [lo, hi_full) are full mask-free tiles; only the peeled tail is masked. hi_full = lo + ((hi - lo) // cfg.BLOCK_K) * cfg.BLOCK_K if seg.fmt.IS_FP8: - # UNI_TILE: the partial tile is just the last iteration -- no peeled masked - # copy of the body (RESULTS.md 19). + # UNI_TILE: the partial tile is just the last iteration (no peeled + # masked copy of the body). if cfg.UNI_TILE: n_full = (hi - lo + cfg.BLOCK_K - 1) // cfg.BLOCK_K else: @@ -1451,9 +1291,9 @@ def _pa_decode_sparse( part_m_ptr, part_l_ptr, part_acc_ptr, - # f32 side-channel per segment: the scalar k_scale ("tensor") or the f32 view - # of the cache ("dsmla"). None -> specialized out of the binary entirely, so - # DSv4 launches keep their exact kernarg layout. + # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view + # ("dsmla"). None elides the argument, keeping other formats' kernarg + # layouts unchanged. main_scl_ptr, extra_scl_ptr, scale: gl.constexpr, @@ -1489,43 +1329,32 @@ def _pa_decode_sparse( MFMA_K: gl.constexpr, GATHER_TW1: gl.constexpr, LDS_PAD: gl.constexpr, - # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, 1 = - # columns); >= the tile's extent on that axis means one shot. + # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, + # 1 = columns); >= the tile's extent means one shot. NOPE_CHUNK: gl.constexpr, CHUNK_AXIS: gl.constexpr, PART_STORE_CACHE: gl.constexpr, UNI_TILE: gl.constexpr, GRID_ORDER: gl.constexpr, - # How many of the NUM_SPLITS programs share the main (SWA) segment. The two - # segments have very different shapes -- main is a contiguous window whose - # length is fixed by the sliding window, extra is the top-k list -- so one - # split count has to be wrong for one of them. Splitting main past - # main_len/BLOCK_K turns every full tile into a half-empty masked one, while - # extra still wants the CTAs. MAIN_SPLITS <= NUM_SPLITS lets main stop at whole - # tiles and extra keep going; programs with split_id >= MAIN_SPLITS get an - # empty main range and contribute an extra-only partial. + # MAIN_SPLITS <= NUM_SPLITS: splitting the SWA window past its tile count + # only manufactures masked partial tiles, so main stops early and extra + # keeps all programs (surplus ones get an empty main range). MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. ADAPTIVE_SPLITS: gl.constexpr, ASM_DEQ: gl.constexpr, - # Per-cache buffer/global gate. buffer_load carries a 32-bit offset, so a cache - # whose span exceeds that must gather via 64-bit gl.load -- but the two caches - # are sized independently (SWA window vs full compressed history), so gating - # them together would drop the fast path on a small cache just because its - # partner is large. At tiny top-k the main/SWA gather is ~94% of the tokens. + # Per-cache buffer/global gate: buffer_load carries a 32-bit offset (2 GB), + # and the two caches are sized independently. MAIN_USE_BUFFER_LOAD: gl.constexpr, EXTRA_USE_BUFFER_LOAD: gl.constexpr, IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, ): - """One program = (query t, split, head-block). Two-loop: main (SWA) then - extra (top-k). NUM_SPLITS==1 writes the output directly; NUM_SPLITS>1 stores - un-normalized partials for the reduce kernel. HAS_INVALID gates -1-sentinel - handling (clamp + score mask) on the full-tile fast paths.""" + """One program = (query, split, head-block). Two-loop: main (SWA) then + extra (top-k). NUM_SPLITS==1 writes the output directly; otherwise stores + un-normalized partials for the reduce kernel.""" NUM_WARPS: gl.constexpr = gl.num_warps() - # The new fp8 formats never trace the peeled masked tail (_decode_tile); - # requiring UNI_TILE keeps it that way instead of silently miscompiling. gl.static_assert( UNI_TILE or (MAIN_FMT != "tensor" and MAIN_FMT != "dsmla"), "tensor/dsmla formats require UNI_TILE=1", @@ -1540,22 +1369,17 @@ def _pa_decode_sparse( ) gl.static_assert( MAIN_FMT != "dsmla" or ROPE_SEPARATE, - "fp8_ds_mla is a separated-rope (GLM) format", + "fp8_ds_mla is a separated-rope (MLA) format", ) gl.static_assert((not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", "ASM_DEQ is the dsv4 E8M0 dequant") - # Page stride alignment hint. The row base of every cache gather is - # ``block_idx*cs0 + pos*TOK`` with block_idx/pos gathered at runtime, so the - # divisibility analysis has to assume 1-byte alignment and a contiguous - # 512-byte row gather lowers to hundreds of global_load_ubyte instead of - # global_load_dwordx4. TOK/HEAD_SIZE are literals the compiler can already - # reason about; cs0 is the one opaque term. The driver sets CS0_ALIGN>1 only - # after checking the strides on the host. + # Row bases are block*cs0 + pos*TOK with runtime block/pos, so divisibility + # analysis sees 1-byte alignment unless the driver vouches for cs0. if CS0_ALIGN > 1: main_cs0 = gl.multiple_of(main_cs0, CS0_ALIGN) extra_cs0 = gl.multiple_of(extra_cs0, CS0_ALIGN) - # GRID_ORDER names the launch axes in grid-dim order (see the driver): dim 0 - # varies fastest, which is what decides XCD/L2 sharing. + # GRID_ORDER names the launch axes in grid-dim order; dim 0 varies fastest, + # which decides XCD/L2 sharing. query_idx = gl.program_id(GRID_ORDER.index("q")) split_id = gl.program_id(GRID_ORDER.index("s")) pid_h = gl.program_id(GRID_ORDER.index("h")) @@ -1600,13 +1424,8 @@ def _pa_decode_sparse( h_off = pid_h * BLOCK_M - # ---- segment lengths, issued first ---- - # These gate the whole KV chain: indptr -> segment range -> indices gather -> - # cache addresses -> cache gather, three dependent memory round trips deep. - # Q is independent of all of it, so issue the indptr loads before Q rather - # than after: the scalar readfirstlane they feed needs s_waitcnt vmcnt(0), and - # behind Q's load plus its through-LDS layout conversion (two barriers) that - # wait lands at the end of a long serial prologue instead of overlapping it. + # Segment lengths issue first: they gate the three-deep memory chain + # (indptr -> indices -> cache), and Q is independent of all of it. main_start = gl.load(main_indptr_ptr + query_idx) main_end = gl.load(main_indptr_ptr + query_idx + 1) if HAS_EXTRA: @@ -1618,13 +1437,10 @@ def _pa_decode_sparse( extra_len = 0 main_len = main_end - main_start - # ---- per-segment softmax/V scales ---- # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw `scale` - # for the sink/normalization (sink is a scaled-score-space logit). For the - # per-tensor fp8 format the cache's k_scale joins the fold (max/exp2 commute - # with a positive scale) and the V-side scale is applied to p in the loop -- - # both are exact, so m/l keep the dequantized-score-domain convention the - # reduce kernel expects. + # for the sink (a scaled-score-space logit). For "tensor" the cache's + # k_scale joins the per-segment fold (max/exp2 commute with a positive + # scale) and the V-side scale is applied to p in the loop — both exact. RCP_LN2: gl.constexpr = 1.4426950408889634 qk_scale = scale * RCP_LN2 main_qk_scale = qk_scale @@ -1640,10 +1456,8 @@ def _pa_decode_sparse( extra_qk_scale = qk_scale * extra_k_scale extra_v_scale = extra_k_scale - # ---- load Q [BLOCK_M, QK_DIM] ---- - # ROPE_SEPARATE loads the two q pieces separately (576 is not a pow-2 arange): - # nope columns [0, KV_DIM) and rope columns [KV_DIM, KV_DIM + ROPE_DIM), each - # converted to the dot-operand layout of its own MFMA in the QK chain. + # Load Q. ROPE_SEPARATE loads the two pieces separately (the combined + # width is not a pow-2 arange), each converted to its dot layout. offs_m_q = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.blocked_q)) offs_d_q = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.blocked_q)) h_q = h_off + offs_m_q @@ -1675,7 +1489,7 @@ def _pa_decode_sparse( h_pv = h_off + offs_m_pv head_mask_pv = h_pv < num_heads - # ---- online-softmax state ---- + # online-softmax state m_i = gl.full( [BLOCK_M], float("-inf"), gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout) ) @@ -1690,18 +1504,13 @@ def _pa_decode_sparse( gl.bfloat16, [BLOCK_K, ROPE_DIM], cfg.rope_shared ) else: - rope_smem = kv_smem # never sliced/read as plane 1 in this geometry - - # ---- how many of the launched splits this query actually wants ---- - # NUM_SPLITS / MAIN_SPLITS come from the host, which only knows the batch's - # AVERAGE segment lengths. In a ragged batch (the SWA window saturates at - # sliding_window while top-k keeps growing with context) the per-query lengths - # differ a lot, and a split count sized for the average over-splits the short - # queries -- each surplus program still gathers a mostly-masked tile and writes - # a full [BLOCK_M, HEAD_SIZE] f32 partial. Recompute the useful count from this - # query's own lengths and let the surplus programs write a neutral partial and - # leave. The reduce skips a split whose m is -inf, so their part_acc is never - # read and does not have to be written. + rope_smem = kv_smem # never read as plane 1 in this geometry + + # ADAPTIVE_SPLITS: the host split count is sized for batch averages; in a + # ragged batch surplus programs would each gather a mostly-masked tile and + # write a full partial. Recompute from this query's own lengths and let + # surplus programs write a neutral partial (m = -inf) and leave — the + # reduce skips those, so their part_acc is never read. if ADAPTIVE_SPLITS: m_tiles = (main_len + BLOCK_K - 1) // BLOCK_K e_tiles = (extra_len + BLOCK_K - 1) // BLOCK_K @@ -1733,7 +1542,7 @@ def _pa_decode_sparse( work_splits = NUM_SPLITS main_splits = MAIN_SPLITS - # ---- main (SWA) segment ---- + # main (SWA) segment main_seg = Seg( main_fmt, main_cache_ptr, @@ -1795,16 +1604,13 @@ def _pa_decode_sparse( rope_smem, ) - # m_i/l_i are in SliceLayout(1, qk_layout); acc in pv_layout. Move the row - # reductions into pv-slice space for output/partials. + # Move the row reductions into pv-slice space for output/partials. m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.pv_layout)) l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.pv_layout)) if NUM_SPLITS == 1: if HAS_SINK: - # m_pv is already the row-max in the base-2 exponent domain - # (row_max * softmax_scale * log2e); lift the sink -- a scaled-score - # logit -- into the same domain and combine there. + # m_pv is in the base-2 exponent domain; lift the sink into it. sink = gl.amd.cdna4.buffer_load( ptr=attn_sink_ptr, offsets=h_pv, mask=head_mask_pv, other=float("-inf") ).to(gl.float32) * RCP_LN2 @@ -1827,14 +1633,9 @@ def _pa_decode_sparse( mask=head_mask_pv[:, None], ) else: - # store un-normalized partials for the reduce kernel. m is already in the - # base-2 exponent domain (row-max * softmax_scale * log2e), which is the - # reduce/skip_reduce convention the triton kernel also uses. + # Un-normalized partials for the reduce kernel; m stays in the base-2 + # exponent domain (the triton reduce's convention too). pm_base = query_idx * pm_stride0 + split_id * pm_stride_s - # The reduce kernel reads these back immediately, so they want to stay - # resident rather than stream to memory. ATT puts 8.5% of all stall cycles - # on the 68 buffer_store_dwordx4 of the accumulator alone (~154 cycles - # each), and the partials are ~31% of the kernel's HBM traffic. gl.amd.cdna4.buffer_store( m_pv, ptr=part_m_ptr + pm_base, @@ -1852,8 +1653,7 @@ def _pa_decode_sparse( offs_d_a = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, cfg.pv_layout)) a_base = query_idx * pa_stride0 + split_id * pa_stride_s a_off = (a_base + h_pv[:, None] * pa_stride_h + offs_d_a[None, :]).to(gl.int32) - # Follow part_acc's own dtype: at bf16 this halves both the partial traffic - # (~31% of the kernel's HBM bytes) and the store instruction count. + # Follow part_acc's own dtype: bf16 halves the partial HBM traffic. gl.amd.cdna4.buffer_store( acc.to(part_acc_ptr.dtype.element_ty), ptr=part_acc_ptr, @@ -1892,26 +1692,16 @@ def _pa_decode_sparse_reduce( ADAPTIVE_SPLITS: gl.constexpr, PART_LOAD_CACHE: gl.constexpr, ): - """Split-KV combine: merge the per-split partials, fold the attn sink, and write - the final output. Partials store m in the base-2 exponent domain (row-max * - softmax_scale * log2e), matching the triton reduce. Grid: (num_queries, heads_blocks). - - Identical for both geometries: partials are [*, HEAD_SIZE] where HEAD_SIZE is - the V width (512 for DSv4 and GLM-5 alike) -- the rope plane never reaches the - output. BLOCK_M is the reduce's own head tile and is deliberately decoupled - from the attention kernel's: the combine is pure bandwidth over - [num_queries, num_heads, HEAD_SIZE] f32, so the only thing that matters is - having enough workgroups to cover the machine. At BLOCK_M = num_heads there is - one workgroup per query -- 64 of them on a 256-CU part, i.e. 3/4 of the GPU - idle and one wave to hide every partial-load latency behind. - """ + """Split-KV combine: merge per-split partials, fold the sink, write the + output. Identical for both geometries (partials are HEAD_SIZE = V wide). + Grid: (num_queries, heads_blocks); the combine is pure bandwidth, so + BLOCK_M is sized for workgroup count, not for the attention kernel's tile.""" NUM_WARPS: gl.constexpr = gl.num_warps() RCP_LN2: gl.constexpr = 1.4426950408889634 query_idx = gl.program_id(0) pid_h = gl.program_id(1) - # One warp covers [BLOCK_M, 64//BLOCK_M * 8] -- lay the 64 lanes out so a - # small BLOCK_M spends them on the head dim instead of idling them on rows. + # Lay the 64 lanes out so a small BLOCK_M spends them on the head dim. TPW0: gl.constexpr = BLOCK_M if BLOCK_M < 8 else 8 TPW1: gl.constexpr = 64 // TPW0 BLK: gl.constexpr = gl.BlockedLayout( @@ -1929,11 +1719,9 @@ def _pa_decode_sparse_reduce( offs_d = gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, BLK)) neg_inf = float("-inf") - # NOT made adaptive. Bounding these loops by the per-query split count instead - # of NUM_SPLITS 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 the - # launch over-splits. The acc load is masked on m > -inf instead, which is - # what actually has to be right (a split that bowed out wrote no accumulator). + # Deliberately NOT bounded by the per-query split count: that would make + # these dynamic loops and lose the static unroll. A split that bowed out + # wrote m = -inf, which is what actually has to be handled (below). m_final = gl.full([BLOCK_M], neg_inf, gl.float32, layout=row_l) # pass 1: global max over splits for s in range(NUM_SPLITS): @@ -1966,9 +1754,8 @@ def _pa_decode_sparse_reduce( l_final = l_final + w * l_s a_base = query_idx * pa_stride0 + s * pa_stride_s a_off = (a_base + h[:, None] * pa_stride_h + offs_d[None, :]).to(gl.int32) - # An adaptive-split program that bowed out wrote m = -inf and no - # accumulator, so its part_acc slot is uninitialized -- mask the load - # rather than relying on w == 0 (0 * NaN is NaN). + # A bowed-out split's part_acc is uninitialized — mask the load rather + # than relying on w == 0 (0 * NaN is NaN). if ADAPTIVE_SPLITS: acc_mask = head_mask[:, None] & (m_s > neg_inf)[:, None] else: @@ -1985,9 +1772,7 @@ def _pa_decode_sparse_reduce( ).to(gl.float32) l_final = l_final + gl.exp2(sink * RCP_LN2 - m_final) - # One reciprocal per row, then a broadcast multiply: an f32 divide lowers to a - # ~10-instruction sequence, so dividing the whole [BLOCK_M, HEAD_SIZE] tile costs - # that per element. + # One reciprocal per row instead of a per-element f32 divide. one_over_l = 1.0 / l_final out = acc * one_over_l[:, None] o_off = (query_idx * out_stride0 + h[:, None] * out_stride1 + offs_d[None, :]).to( diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 3c5a26075a4..7e796ad7b5a 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -663,8 +663,7 @@ def _pa_decode_sparse_gfx950_gluon( avg_main = indices.numel() / max(1, num_queries) avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 - # Kernel-side cache-format tags (the kernel is shared with the separated-rope - # GLM/MLA path in sparse_mla_decode.py; DSv4 always runs ROPE_SEPARATE=False). + # Kernel-side cache-format tags (kernel shared with sparse_mla_decode.py). if UNIFORM: main_fmt = "uniform" if main_is_fp8 else "bf16" extra_fmt = main_fmt @@ -898,9 +897,8 @@ def _pick_splits(hb): part_m, part_l, part_acc, - # f32 side-channel pointers (per-tensor k_scale / fp8_ds_mla view) -- - # only the separated-rope formats use them; None keeps the DSv4 - # kernarg layout byte-identical (constexpr-None args are elided). + # f32 scale pointers (separated-rope formats only); None is elided, + # keeping the DSv4 kernarg layout unchanged. None, None, scale, diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index 0d6a8d86a2b..0edc949d72a 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -1,48 +1,21 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Sparse MLA decode (gfx950 gluon): separated-rope geometry — GLM-5 / DSA-style -``kv_lora_rank`` latent ⊕ appended decoupled rope, token-granular top-k gather. - -This is the separated-rope sibling of ``pa_decode_sparse``: both drive the same -gluon kernel (``_gluon_kernels/gfx950/attention/pa_decode_sparse.py``), but the -two model families have different enough calling structures that they get -different APIs — DSv4 (rope *inside* the 512-wide row, K == V, optional SWA + -top-k two-loop, attention sink) stays on ``pa_decode_sparse``; MLA-lineage -models (rope *appended*: QK over ``kv_lora_rank + qk_rope_head_dim``, V = the -latent only, single top-k segment, no sink) use this module. - -The calling convention follows the asm path this replaces -(``aiter.mla.mla_decode_fwd`` as invoked by vLLM's ``ROCMAiterMLASparseImpl``), -with two deliberate deviations, both favorable: no head padding (the asm kernel -repeat_interleaves H<16 to 16 and computes the duplicate rows; here H=8 runs -natively at BLOCK_M=8), and q stays bf16 (no q-quant kernel; <1% accuracy -effect). - -Cache formats, inferred from the buffer: - - [nb, block, R] / [slots, 1, 1, R] / [slots, R] bf16 full-row bf16 - same shapes in fp8/uint8 (+ scalar ``kv_scale``) flat per-tensor fp8 - — vLLM's production layout: ``concat_and_cache_mla(...,"fp8", - scale=layer._k_scale)`` quantizes the whole row (rope included) with - one scale. The kernel never dequantizes in the tile loop: the K-side - scale folds into qk_scale and the V-side scale into p (exact; see the - kernel docstring). - [blocks, block_size, 656] uint8 vLLM ``fp8_ds_mla`` - — 512 fp8 | 4×f32 per-128 scales | 64 bf16 rope. - -Indices are GLOBAL slot ids (``triton_convert_req_index_to_global_index`` -output). NOTE: that converter maps -1 / out-of-range entries to slot 0 despite -its docstring, so the production stream carries no -1 sentinels and the default -``has_invalid=False`` matches it (and skips the sentinel clamp+mask). Pass -``has_invalid=True`` for index streams that really carry -1 entries — they are -then masked out, which is *more* correct than the converter's slot-0 -substitution. +"""Sparse MLA decode (gfx950 gluon): ``kv_lora_rank`` latent + appended +decoupled rope, token-granular top-k gather. + +Separated-rope sibling of ``pa_decode_sparse`` — both drive the same gluon +kernel, but the calling structures differ enough to warrant separate APIs: +DSv4 (rope inside the row, K == V, optional SWA + top-k two-loop, attention +sink) stays on ``pa_decode_sparse``; MLA-lineage models (single top-k segment, +no sink, QK over ``kv_lora_rank + qk_rope_head_dim``, V = the latent) use this. + +Follows the ``aiter.mla.mla_decode_fwd`` calling convention as vLLM's +``ROCMAiterMLASparseImpl`` uses it, with two deviations: no head padding +(H < 16 runs natively at BLOCK_M = next_pow2(H)) and q stays bf16 (no q-quant +kernel; <1% accuracy effect). """ -import math -import os as _os - import torch from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( @@ -62,13 +35,22 @@ _LOGGER = AiterTritonLogger() +# buffer_load carries a signed 32-bit offset; caches past this span gather +# through 64-bit addresses (a production MLA pool crosses it at ~3.7M tokens). MAX_BYTES = 2**31 - 1 +# Tuned launch config (gfx950 / MI355). num_warps = BLOCK_K // 16 (warps tile +# the dot-N, MFMA N=16); GATHER_TW1=32 requests a whole 512 B token row per +# gather instruction. +_BLOCK_K = 64 +_MFMA_K = 16 +_GATHER_TW1 = 32 +_LDS_PAD = 8 + def _check_out(out, q, kv_lora_rank): - """Caller-supplied output buffer, or a fresh one. The attention output is - [C, H, kv_lora_rank]; a wider buffer (e.g. a q-shaped scratch) is accepted - and written in its leading kv_lora_rank columns.""" + """Caller-supplied output buffer, or a fresh one. A buffer wider than + kv_lora_rank is accepted and written in its leading columns.""" C, H = q.shape[0], q.shape[1] if out is None: return torch.empty( @@ -83,17 +65,14 @@ def _check_out(out, q, kv_lora_rank): def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): - """-> (fmt, cache, alt_ptr_tensor, scl_ptr_tensor, block_size, fp8_fnuz). - - Pointer roles follow the kernel's Seg contract: alt = the bf16 view - ("dsmla" rope tail) or the cache itself ("bf16"); scl = the f32 scalar - k_scale ("tensor") or the f32 view of the cache ("dsmla").""" + """-> (fmt, cache, alt_ptr, scl_ptr, block_size, fp8_fnuz). Pointer roles + follow the kernel's Seg contract (see its docstring).""" if kv.ndim == 4: # [slots, 1, 1, R] -- the asm mla_decode_fwd view assert kv.shape[1] == 1 and kv.shape[2] == 1 kv = kv.reshape(kv.shape[0], kv.shape[3]) elif kv.ndim == 3 and kv.shape[2] == d_qk: - # the vLLM paged cache [nb, block_size, R]: indices are global slot - # ids and rows are contiguous, so the flat view is stride-identical. + # vLLM paged cache [nb, block_size, R]: indices are global slot ids and + # rows are contiguous, so the flat view is stride-identical. assert kv.stride(2) == 1 and kv.stride(1) == kv.shape[2] * kv.stride(2) kv = kv.reshape(-1, d_qk) @@ -148,30 +127,27 @@ def sparse_mla_decode_fwd( Args: q: ``[C, H, kv_lora_rank + qk_rope_head_dim]`` bf16 decode queries. - H = per-rank query heads; 8 (TP8) and 16 (TP4) both run natively. - kv_buffer: the KV pool in one of the formats in the module docstring. - kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts - (vLLM builds ``cumsum(min(seq_len, topk))``; uniform ``C * topk`` - stepping also works). + kv_buffer: the KV pool — ``[nb, block, R]`` / ``[slots, 1, 1, R]`` / + ``[slots, R]`` in bf16, the same shapes in fp8 (+ scalar + ``kv_scale``), or ``[nb, block, 656]`` uint8 (vLLM ``fp8_ds_mla``). + kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts. kv_indices: flat int32 GLOBAL slot ids into the pool. softmax_scale: the layer's softmax scale. - kv_scale: ``[1]`` f32 per-tensor cache scale — required for the flat - fp8 format (vLLM's ``layer._k_scale``), unused otherwise. - kv_splits: split-K override; default picks by the occupancy policy. - skip_reduce: with split-K active, return ``(part_acc, part_m, - part_l)`` instead of launching the combine. + kv_scale: ``[1]`` f32 per-tensor cache scale (``layer._k_scale``); + required for the flat fp8 format. + kv_splits: split-K override; default follows the occupancy policy. + skip_reduce: with split-K active, return ``(part_acc, part_m, part_l)`` + instead of launching the combine. has_invalid: index stream carries -1 sentinels (masked out). Default - False — the production converter emits none. - out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination (written in - ``[..., :kv_lora_rank]``). + False: vLLM's index converter emits none (it maps -1 to slot 0). + out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. Returns: - ``[C, H, kv_lora_rank]`` bf16 attention output (the latent V; callers - run ``v_up_proj`` after, exactly as with the asm path). + ``[C, H, kv_lora_rank]`` bf16 attention output (the latent V). """ assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert arch_info.get_arch() == "gfx950", "sparse_mla_decode_fwd is gfx950-only" - assert q.dtype == torch.bfloat16, "q stays bf16 (no q-quant; see module doc)" + assert q.dtype == torch.bfloat16 num_queries, num_heads, d_qk = q.shape assert d_qk == kv_lora_rank + qk_rope_head_dim, ( f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" @@ -187,98 +163,48 @@ def sparse_mla_decode_fwd( kv_indices = _as_int32_contiguous_1d(kv_indices) kv_indptr = _as_int32_contiguous_1d(kv_indptr) assert kv_indptr.numel() == num_queries + 1 - # HAS_SINK=False: MLA-lineage models carry no attention sink. The kernel - # still wants a live pointer for the elided-at-compile-time argument slot. + # No attention sink in MLA models; the kernel still wants a live pointer + # for the compile-time-elided argument slot. attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) - # Tuned launch config (gfx950 / MI355), shared with the DSv4 wrapper; same - # env knobs so interleaved A/B sweeps drive both paths together. - BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 - _bk = _os.environ.get("AITER_PA_DECODE_BLOCK_K") - if _bk: - BLOCK_K = int(_bk) - _bm = _os.environ.get("AITER_PA_DECODE_BLOCK_M") - if _bm: - BLOCK_M = int(_bm) - elif num_heads < 16: - # H=8 (TP8): BLOCK_M=8 computes 8 real rows where a masked BLOCK_M=16 - # pads to 16 -- against the asm kernel's repeat_interleave to >=16 - # heads this is the free half of the measured win. - BLOCK_M = max(8, 1 << (num_heads - 1).bit_length()) - _mk = _os.environ.get("AITER_PA_DECODE_MFMA_K") - if _mk: - MFMA_K = int(_mk) - num_warps = max(1, BLOCK_K // 16) - gather_tw1 = int(_os.environ.get("AITER_PA_DECODE_TW1", "32")) - lds_pad = int(_os.environ.get("AITER_PA_DECODE_LDS_PAD", "8")) + # H < 16 runs natively at BLOCK_M = next_pow2(H) instead of padding heads. + block_m = 16 if num_heads >= 16 else max(8, 1 << (num_heads - 1).bit_length()) + num_warps = _BLOCK_K // 16 num_rows = cache.shape[0] * block_size if cache.ndim >= 2 else cache.shape[0] avg_topk = kv_indices.numel() / max(1, num_queries) + # Alignment hint for cs0 so row gathers can vectorize. cs0_align = 1 - for _a in (16, 8, 4, 2): - if int(cache.stride(0)) % _a == 0: - cs0_align = _a + for a in (16, 8, 4, 2): + if int(cache.stride(0)) % a == 0: + cs0_align = a break - _ca = _os.environ.get("AITER_PA_DECODE_CS0_ALIGN") - if _ca: - cs0_align = int(_ca) - # buffer_load carries a 32-bit offset (2 GB cap); a production pool crosses - # it at ~3.7M cached tokens, so the 64-bit gl.load path is first-class here - # (measured: same speed as the buffer path, +2-3 VGPRs, no spills). use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES - idx_use_buffer_load = ( - _os.environ.get("AITER_PA_DECODE_IDX_BUF", "1") == "1" - and max_addressable_bytes(kv_indices) < MAX_BYTES - ) - _fg = _os.environ.get("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "") - if _fg in ("1", "main"): - use_buffer_load = False + idx_use_buffer_load = max_addressable_bytes(kv_indices) < MAX_BYTES - HEAD_ALIGNED = num_heads % BLOCK_M == 0 - heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M + head_aligned = num_heads % block_m == 0 + heads_blocks = (num_heads + block_m - 1) // block_m out = _check_out(out, q, kv_lora_rank) - _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") - - def _pick_splits(hb): - if _sp: - return max(1, int(_sp)) - if kv_splits is not None: - return max(1, int(kv_splits)) - return _decode_num_splits_occ(num_queries, hb, avg_topk, 0.0, BLOCK_K) - - num_splits = _pick_splits(heads_blocks) - - # Short top-k BLOCK_M halving gate (real parallelism at small launches). - _tiles = max(1, math.ceil(avg_topk / BLOCK_K)) if avg_topk > 0 else 1 - if ( - _bm is None - and BLOCK_M == 16 - and num_heads % 8 == 0 - and _tiles <= 2 - and num_queries * heads_blocks * num_splits <= get_num_sms() - ): - _hb8 = (num_heads + 7) // 8 - _ns8 = _pick_splits(_hb8) - if _ns8 >= num_splits and num_queries * _hb8 * _ns8 <= get_num_sms(): - BLOCK_M = 8 - HEAD_ALIGNED = num_heads % BLOCK_M == 0 - heads_blocks = _hb8 - num_splits = _ns8 + if kv_splits is not None: + num_splits = max(1, int(kv_splits)) + else: + num_splits = _decode_num_splits_occ( + num_queries, heads_blocks, avg_topk, 0.0, _BLOCK_K + ) if num_splits > 1: part_m = torch.empty( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device ) part_l = torch.empty_like(part_m) - _pab = ( - _os.environ.get("AITER_PA_DECODE_PART_BF16", "1") == "1" and not skip_reduce - ) + # bf16 partials halve the split-K HBM traffic; skip_reduce hands the + # partials back to the caller and keeps f32. part_acc = torch.empty( (num_queries, num_splits, num_heads, kv_lora_rank), - dtype=torch.bfloat16 if _pab else torch.float32, + dtype=torch.float32 if skip_reduce else torch.bfloat16, device=q.device, ) pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) @@ -291,42 +217,17 @@ def _pick_splits(hb): part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # Dequant chunking: split the axis that still has per-lane register repeats. - col_reps = kv_lora_rank // (gather_tw1 * 16) + # Dequant chunking: split the axis that keeps per-lane register repeats + # (rows at GATHER_TW1=32); waves_per_eu=2 caps the allocator at 256 VGPRs + # for the second occupancy slot, except when the launch cannot fill it. + col_reps = kv_lora_rank // (_GATHER_TW1 * 16) chunk_axis = 1 if col_reps >= 4 else 0 - _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") - if chunk_axis == 0: - nope_chunk = int(_nc) if _nc else max(1, BLOCK_K // 4) - else: - nope_chunk = int(_nc) if _nc else min(128, kv_lora_rank) - - if chunk_axis == 0 or nope_chunk < kv_lora_rank: - waves_per_eu = 2 - one_wg_per_cu = ( - use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms() - ) - if one_wg_per_cu: + nope_chunk = max(1, _BLOCK_K // 4) if chunk_axis == 0 else min(128, kv_lora_rank) + waves_per_eu = 2 + if use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms(): waves_per_eu = 1 - _wp = _os.environ.get("AITER_PA_DECODE_WPEU") - if _wp: - waves_per_eu = int(_wp) - - adaptive_splits = ( - num_splits > 1 and _os.environ.get("AITER_PA_DECODE_ADAPTIVE", "1") == "1" - ) - - _go = _os.environ.get("AITER_PA_DECODE_GRID_ORDER", "qsh") - _ax = {"q": num_queries, "s": num_splits, "h": heads_blocks} - assert sorted(_go) == ["h", "q", "s"], f"bad AITER_PA_DECODE_GRID_ORDER {_go!r}" - # The separated-rope fp8 formats never trace the legacy peeled masked tail - # (kernel static_assert); bf16 keeps the knob. - uni_tile = ( - True - if fmt in ("tensor", "dsmla") - else _os.environ.get("AITER_PA_DECODE_UNI_TILE", "1") == "1" - ) - grid = tuple(_ax[c] for c in _go) + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, cache, @@ -370,20 +271,20 @@ def _pick_splits(hb): ROPE_DIM=qk_rope_head_dim, HEAD_SIZE=kv_lora_rank, ROPE_SEPARATE=True, - BLOCK_M=BLOCK_M, - BLOCK_K=BLOCK_K, + BLOCK_M=block_m, + BLOCK_K=_BLOCK_K, NUM_SPLITS=num_splits, - HEAD_ALIGNED=HEAD_ALIGNED, - MFMA_K=MFMA_K, - GATHER_TW1=gather_tw1, - LDS_PAD=lds_pad, + HEAD_ALIGNED=head_aligned, + MFMA_K=_MFMA_K, + GATHER_TW1=_GATHER_TW1, + LDS_PAD=_LDS_PAD, NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, - PART_STORE_CACHE=_os.environ.get("AITER_PA_DECODE_PART_ST", ""), - UNI_TILE=uni_tile, - GRID_ORDER=_go, + PART_STORE_CACHE="", + UNI_TILE=True, + GRID_ORDER="qsh", MAIN_SPLITS=num_splits, - ADAPTIVE_SPLITS=adaptive_splits, + ADAPTIVE_SPLITS=num_splits > 1, ASM_DEQ=False, MAIN_USE_BUFFER_LOAD=use_buffer_load, EXTRA_USE_BUFFER_LOAD=use_buffer_load, @@ -399,9 +300,9 @@ def _pick_splits(hb): if skip_reduce: return part_acc, part_m, part_l - red_block_m = int(_os.environ.get("AITER_PA_DECODE_REDUCE_BLOCK_M", "1")) - red_warps = int(_os.environ.get("AITER_PA_DECODE_REDUCE_WARPS", "1")) - rgrid = (num_queries, (num_heads + red_block_m - 1) // red_block_m) + # One head per reduce workgroup: the combine is pure bandwidth, so size the + # grid for coverage, not for the attention kernel's tile. + rgrid = (num_queries, num_heads) _pa_decode_sparse_reduce_gfx950[rgrid]( part_m, part_l, @@ -418,11 +319,11 @@ def _pick_splits(hb): num_heads, HAS_SINK=False, HEAD_SIZE=kv_lora_rank, - BLOCK_M=red_block_m, + BLOCK_M=1, NUM_SPLITS=num_splits, - HEAD_ALIGNED=num_heads % red_block_m == 0, - ADAPTIVE_SPLITS=adaptive_splits, - PART_LOAD_CACHE=_os.environ.get("AITER_PA_DECODE_PART_LD", ""), - num_warps=red_warps, + HEAD_ALIGNED=True, + ADAPTIVE_SPLITS=num_splits > 1, + PART_LOAD_CACHE="", + num_warps=1, ) return out diff --git a/op_tests/triton_tests/attention/test_sparse_mla_decode.py b/op_tests/triton_tests/attention/test_sparse_mla_decode.py index 54a3984fdf1..89d7d554e3a 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla_decode.py +++ b/op_tests/triton_tests/attention/test_sparse_mla_decode.py @@ -1,17 +1,12 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Tests for the separated-rope sparse MLA decode API (gfx950 gluon): -``aiter.ops.triton.attention.sparse_mla_decode.sparse_mla_decode_fwd``. - -Covers the three cache formats (bf16, flat per-tensor fp8, vLLM fp8_ds_mla), -both query-head counts GLM-5 ships (8 = TP8, 16 = TP4), ragged and -1-padded -index streams, split-K on/off, the vLLM cache shapes, and the 64-bit (>2 GB -emulated) gather path. The per-tensor scale-folding is additionally pinned by a -bitwise test against the bf16 path. -""" +"""Tests for ``sparse_mla_decode_fwd`` (gfx950 gluon, separated-rope MLA). -import os +Target-oriented: the production formats (bf16 and flat per-tensor fp8) get the +shape matrix; fp8_ds_mla gets a single smoke case. Index streams are uniform or +ragged with no -1 sentinels, matching what vLLM's converter emits. +""" import pytest import torch @@ -19,10 +14,10 @@ from aiter.ops.triton.utils._triton import arch_info if arch_info.get_arch() == "gfx950": + import aiter.ops.triton.attention.sparse_mla_decode as smd from aiter.ops.triton.attention.sparse_mla_decode import sparse_mla_decode_fwd FP8_MAX = 448.0 -TINY = 1.1754944e-38 KV_LORA, ROPE = 512, 64 D_QK = KV_LORA + ROPE @@ -32,24 +27,20 @@ def _skip_unless_gfx950(): pytest.skip("sparse_mla_decode_fwd is gfx950-only") -# ---------------------------------------------------------------- packers - - def quantize_flat_fp8(kv): - """GLM-5 production layout: whole row (rope included) fp8 with ONE - per-tensor scale (what concat_and_cache_mla writes with layer._k_scale).""" - scale = (kv.float().abs().amax() / FP8_MAX).clamp_min(TINY).reshape(1) + """vLLM's production layout: whole row fp8 with one per-tensor scale.""" + scale = (kv.float().abs().amax() / FP8_MAX).clamp_min(1e-30).reshape(1) q8 = (kv.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) return q8.view(torch.uint8), scale.to(torch.float32) def pack_ds_mla(kv, block_size=64): - """vLLM ``fp8_ds_mla``: 656-byte rows [512 fp8 | 4 x f32 per-128 | 64 bf16].""" + """vLLM fp8_ds_mla: 656 B rows [512 fp8 | 4 x f32 per-128 | 64 bf16 rope].""" T = kv.shape[0] nb = (T + block_size - 1) // block_size cache = torch.zeros(nb * block_size, 656, dtype=torch.uint8, device=kv.device) x = kv[:, :KV_LORA].float().reshape(T, 4, 128) - scale = (x.abs().amax(-1) / FP8_MAX).clamp_min(TINY) + scale = (x.abs().amax(-1) / FP8_MAX).clamp_min(1e-30) q8 = (x / scale[..., None]).to(torch.float8_e4m3fn).view(torch.uint8) cache[:T, :KV_LORA] = q8.reshape(T, KV_LORA) cache[:T, 512:528].view(torch.float32).copy_(scale) @@ -70,44 +61,27 @@ def dequant_ds_mla(cache): return torch.cat([nope, rope], dim=-1) -def ragged_indices(C, pool, topk, device, seed, mode): - """mode: 'uniform' (full topk), 'ragged' (true prefix sum, no sentinels -- - the production shape), 'padded' (-1 tail padding, needs has_invalid).""" +def make_indices(C, pool, topk, device, seed, ragged): g = torch.Generator(device="cpu").manual_seed(seed) lens = ( - torch.full((C,), topk) - if mode == "uniform" - else torch.randint(1, topk + 1, (C,), generator=g) + torch.randint(1, topk + 1, (C,), generator=g) + if ragged + else torch.full((C,), topk) ) - rows, flat_lens = [], [] - for c in range(C): - k = int(lens[c]) - sel = torch.randperm(pool, generator=g)[:k] - if mode == "padded": - row = torch.full((topk,), -1, dtype=torch.int64) - row[:k] = sel - rows.append(row) - flat_lens.append(topk) - else: - rows.append(sel) - flat_lens.append(k) - indices = torch.cat(rows) + rows = [torch.randperm(pool, generator=g)[: int(lens[c])] for c in range(C)] indptr = torch.zeros(C + 1, dtype=torch.int64) - indptr[1:] = torch.tensor(flat_lens).cumsum(0) + indptr[1:] = lens.cumsum(0) return ( - indices.to(torch.int32).to(device), + torch.cat(rows).to(torch.int32).to(device), indptr.to(torch.int32).to(device), ) def reference(q, kv_truth, indices, indptr, sm_scale): """f32 attention over the dequantized cache (tests the kernel, not fp8).""" - C = q.shape[0] outs = [] - for c in range(C): - sel = indices[indptr[c] : indptr[c + 1]].long() - sel = sel[sel >= 0] - kvs = kv_truth[sel].float() + for c in range(q.shape[0]): + kvs = kv_truth[indices[indptr[c] : indptr[c + 1]].long()].float() s = torch.einsum("hd,td->ht", q[c].float(), kvs) * sm_scale p = torch.softmax(s, dim=-1) outs.append(torch.einsum("ht,td->hd", p, kvs[:, :KV_LORA])) @@ -115,15 +89,14 @@ def reference(q, kv_truth, indices, indptr, sm_scale): def rel_err(out, ref): - return ((out.float() - ref).abs().max() / ref.abs().max().clamp_min(1e-20)).item() + return ((out.float() - ref).abs().max() / ref.abs().max()).item() -def _build(fmt, C, H, topk, pool, mode, seed=0): +def _build(fmt, C, H, topk, pool, ragged, seed=0): torch.manual_seed(seed) q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 - kv[:: max(1, pool // 97)] *= 7.0 # a few heavy-tail tokens - idx, ptr = ragged_indices(C, pool, topk, "cuda", seed + 7, mode) + idx, ptr = make_indices(C, pool, topk, "cuda", seed + 7, ragged) ks = None if fmt == "bf16": cache, truth = kv, kv.float() @@ -136,69 +109,61 @@ def _build(fmt, C, H, topk, pool, mode, seed=0): return q, cache, ks, idx, ptr, truth -# ---------------------------------------------------------------- tests +def _run_and_check(fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, **kwargs): + sm = D_QK**-0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) + ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) + out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + e = rel_err(out, ref) + assert e < tol, f"{fmt} C={C} H={H} topk={topk} ragged={ragged}: rel-err {e:.3e}" -@pytest.mark.parametrize("fmt", ["bf16", "tensor", "dsmla"]) +@pytest.mark.parametrize("fmt", ["bf16", "tensor"]) @pytest.mark.parametrize("H", [8, 16]) @pytest.mark.parametrize( - "topk,mode", - [(2048, "uniform"), (500, "ragged"), (2048, "padded")], - ids=["topk2048", "ragged500", "padded2048"], + "topk,ragged", [(2048, False), (500, True)], ids=["topk2048", "ragged500"] ) -def test_sparse_mla_decode(fmt, H, topk, mode): +def test_sparse_mla_decode(fmt, H, topk, ragged): _skip_unless_gfx950() - C, pool = 8, 1 << 16 - sm = D_QK**-0.5 - q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, mode) - ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out = sparse_mla_decode_fwd( - q, cache, ptr, idx, sm, kv_scale=ks, has_invalid=(mode == "padded") - ) - e = rel_err(out, ref) - assert e < 2e-2, f"{fmt} H={H} topk={topk} {mode}: rel-err {e:.3e}" + _run_and_check(fmt, C=8, H=H, topk=topk, ragged=ragged) + + +def test_ds_mla_format(): + _skip_unless_gfx950() + _run_and_check("dsmla", C=8, H=16, topk=2048, ragged=True) @pytest.mark.parametrize("kv_splits", [1, None], ids=["splits1", "splitsauto"]) def test_split_k(kv_splits): _skip_unless_gfx950() - C, pool, topk = 4, 1 << 15, 2048 - sm = D_QK**-0.5 - q, cache, ks, idx, ptr, truth = _build("tensor", C, 16, topk, pool, "uniform") - ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks, - kv_splits=kv_splits) - e = rel_err(out, ref) - assert e < 2e-2, f"splits={kv_splits}: rel-err {e:.3e}" + _run_and_check("tensor", C=4, H=16, topk=2048, ragged=False, + pool=1 << 15, kv_splits=kv_splits) def test_per_tensor_scale_folding_bitwise(): - """kv_scale=1.0 on fp8-exact values must equal the bf16 path bit-for-bit: - the K-scale fold into qk_scale and the V-scale fold into p are exact, and - fp8 -> bf16 staging is exact.""" + """kv_scale=1.0 on fp8-exact values must equal the bf16 path bit-for-bit + (K-scale folds into qk_scale, V-scale into p; fp8 -> bf16 is exact).""" _skip_unless_gfx950() torch.manual_seed(3) C, H, topk, pool = 4, 16, 512, 1 << 15 sm = D_QK**-0.5 - raw = torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25 - exact = raw.float().to(torch.float8_e4m3fn) - kv_bf16 = exact.to(torch.bfloat16) - cache_u8 = exact.view(torch.uint8) + exact = (torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25 + ).float().to(torch.float8_e4m3fn) ks = torch.ones(1, dtype=torch.float32, device="cuda") q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 - idx, ptr = ragged_indices(C, pool, topk, "cuda", 11, "ragged") - o_t = sparse_mla_decode_fwd(q, cache_u8, ptr, idx, sm, kv_scale=ks) - o_b = sparse_mla_decode_fwd(q, kv_bf16, ptr, idx, sm) + idx, ptr = make_indices(C, pool, topk, "cuda", 11, ragged=True) + o_t = sparse_mla_decode_fwd(q, exact.view(torch.uint8), ptr, idx, sm, kv_scale=ks) + o_b = sparse_mla_decode_fwd(q, exact.to(torch.bfloat16), ptr, idx, sm) assert torch.equal(o_t.view(torch.int16), o_b.view(torch.int16)) def test_vllm_cache_shapes(): - """The [slots,1,1,R] (asm mla_decode_fwd) and [nb,block,R] (vLLM paged) - views must produce identical results to the flat [slots,R] form.""" + """[slots,1,1,R] (asm view) and [nb,block,R] (vLLM paged) must match the + flat [slots,R] form exactly.""" _skip_unless_gfx950() C, pool, topk = 4, 1 << 14, 512 sm = D_QK**-0.5 - q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, "uniform") + q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, ragged=False) o_flat = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks) o_4d = sparse_mla_decode_fwd( q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks @@ -209,16 +174,9 @@ def test_vllm_cache_shapes(): assert torch.equal(o_flat, o_4d) and torch.equal(o_flat, o_paged) -@pytest.mark.parametrize("fmt", ["tensor", "dsmla"]) -def test_global_load_path(fmt, monkeypatch): - """Force the 64-bit gather path (production: pools cross buffer_load's 2 GB - offset limit at ~3.7M tokens) and check correctness is unchanged.""" +def test_global_load_path(monkeypatch): + """Force the 64-bit gather path (production pools cross buffer_load's 2 GB + offset limit at ~3.7M tokens).""" _skip_unless_gfx950() - monkeypatch.setenv("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "1") - C, pool, topk = 8, 1 << 16, 2048 - sm = D_QK**-0.5 - q, cache, ks, idx, ptr, truth = _build(fmt, C, 16, topk, pool, "uniform") - ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks) - e = rel_err(out, ref) - assert e < 2e-2, f"{fmt} 64-bit path: rel-err {e:.3e}" + monkeypatch.setattr(smd, "MAX_BYTES", 0) + _run_and_check("tensor", C=8, H=16, topk=2048, ragged=False) From 4f4fb36b934a415314fe6e5a58233a27932108e5 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 06:10:34 +0000 Subject: [PATCH 27/51] [GLUON][GFX950][DSV4] Drop the gluon driver's env knobs; leaner reduce _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) --- .../gfx950/attention/pa_decode_sparse.py | 144 ++++----- .../ops/triton/attention/pa_decode_sparse.py | 292 +++++------------- .../ops/triton/attention/sparse_mla_decode.py | 1 - 3 files changed, 160 insertions(+), 277 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 1735714ce20..fa8ddbde32b 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -3,32 +3,32 @@ """Gluon (gfx950) sparse decode for DeepSeek-V4 and separated-rope MLA (GLM-5). -One kernel serves two geometries, selected by the ``ROPE_SEPARATE`` constexpr: +One kernel serves two geometries, selected by the ROPE_SEPARATE constexpr: False (DSv4): rope lives inside the pow-2 row. One LDS plane - ``kv_smem [BLOCK_K, HEAD_SIZE]``, one QK MFMA chain, V = the same plane. - True (MLA): rope is appended (``kv_lora_rank`` latent + rope = QK width; - V = the latent only). A second pow-2 LDS plane ``rope_smem`` holds the + kv_smem [BLOCK_K, HEAD_SIZE], one QK MFMA chain, V = the same plane. + True (MLA): rope is appended (kv_lora_rank latent + rope = QK width; + V = the latent only). A second pow-2 LDS plane rope_smem holds the K-only rope and the QK contraction chains a second MFMA into the first. Plane 0 is the entire V in both geometries. -Cache formats (``Fmt.KIND``), per segment: +Cache formats (Fmt.KIND), per segment: "bf16" full-row bf16 "dsv4" 448 fp8 | 64 bf16 rope (576 B rows) + 8 B UE8M0-per-64 trailer "uniform" whole-row fp8 + separate f32 per-64 kv_scales - "tensor" whole-row fp8 + ONE per-tensor f32 scale (k_scale). The scale - never touches the tile loop: K-side folds into qk_scale, + "tensor" whole-row fp8 + a single per-tensor f32 scale (k_scale). The + scale never touches the tile loop: K-side folds into qk_scale, V-side into p. "dsmla" vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope (656 B rows); requires ROPE_SEPARATE. Compile-time state is bundled into aggregates instead of long parameter lists: -``Cfg`` (geometry, layouts, behavior knobs), ``Fmt`` (per-segment cache -format), ``Seg`` (runtime pointers/bounds). Two-loop (SWA + top-k) or a single -segment; 2D and 3D (split-K + reduce) share one kernel. Launchers: -``aiter/ops/triton/attention/pa_decode_sparse.py`` (DSv4 / uniform pool) and -``aiter/ops/triton/attention/sparse_mla_decode.py`` (separated-rope MLA). +Cfg (geometry, layouts, behavior knobs), Fmt (per-segment cache format), Seg +(runtime pointers/bounds). Two-loop (SWA + top-k) or a single segment; 2D and +3D (split-K + reduce) share one kernel. Launchers: +aiter/ops/triton/attention/pa_decode_sparse.py (DSv4 / uniform pool) and +aiter/ops/triton/attention/sparse_mla_decode.py (separated-rope MLA). """ from triton.experimental import gluon @@ -57,8 +57,8 @@ def _rmax(x, axis): @gluon.jit def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): - """Gather rows[i] + col[j]. ``row`` is the per-token offset in ptr's element - units; ``col`` a small compile-time arange. Keeping them apart resolves one + """Gather rows[i] + col[j]. row is the per-token offset in ptr's element + units; col a small compile-time arange. Keeping them apart resolves one pointer per token on the 64-bit path (the column offset folds into the load's immediate) instead of a 64-bit add per element.""" if USE_BUFFER_LOAD: @@ -109,11 +109,12 @@ def _scale_load( OTHER: gl.constexpr, ): """Gather the NG per-group scales of each token and broadcast to W_FULL - columns. Indexing the full row with ``offs // GROUP`` would build a - [BLOCK_K, W_FULL] pointer tensor for NG distinct values — redundancy only - buffer_load's 32-bit offsets CSE away. ``scl_l`` is chosen so the broadcast - lands on ``gather_l`` as a register rename (assert_trivial proves it). - Element type follows ``ptr``: u8 E8M0 bytes (OTHER=127 -> 2^0) or f32 + columns. Indexing the full row with offs // GROUP would build a + [BLOCK_K, W_FULL] pointer tensor for NG distinct values, and only + buffer_load's 32-bit offsets CSE that redundancy away. scl_l is chosen so + the broadcast lands on gather_l as a register rename (assert_trivial + proves it). + Element type follows ptr: u8 E8M0 bytes (OTHER=127 -> 2^0) or f32 (OTHER=0.0 -> masked lanes dequant to 0).""" cols = gl.arange(0, NG, layout=gl.SliceLayout(0, scl_l)) rows = gl.convert_layout(row, gl.SliceLayout(1, scl_l)) @@ -156,7 +157,7 @@ def _scale_load( def _split2(x): """Register split along dim 1: [A, B] -> two [A, B//2] in the input's own layout. Free only while the column direction is a per-lane register repeat; - ``assert_trivial`` turns anything else into a compile error.""" + assert_trivial turns anything else into a compile error.""" layout: gl.constexpr = x.type.layout x_r = x.reshape(x.shape[0], 2, x.shape[1] // 2).permute(0, 2, 1) x0, x1 = gl.split(x_r) @@ -188,8 +189,8 @@ def _split_ax(x, AXIS: gl.constexpr): # fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. Emitted as two # single-output asm calls, not one two-output blob: a single instruction reads -# all sources before writing dst, so no output needs early clobber — the blob's -# forced `=&v` liveness spills on a kernel with no VGPR headroom. +# all its sources before writing dst, so no output needs early clobber. The +# blob form does, and that forced =&v liveness spills a kernel with no headroom. _DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") _DEQ_HI: gl.constexpr = gl.constexpr( "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" @@ -201,10 +202,10 @@ def _split_ax(x, AXIS: gl.constexpr): def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): """[BLOCK_K, W8/2] int16 (4 packed fp8) + raw E8M0 byte -> [BLOCK_K, W8] bf16. - The hardware reads only bits [30:23] of the scale operand — bits [14:7] of - its high half — so a 16-bit ``e << 7`` lands the exponent where an f32 - ``e << 23`` would, in ONE register. Power-of-two scale keeps this - bit-identical to the unfused convert+multiply chain.""" + The hardware reads only bits [30:23] of the scale operand, i.e. bits [14:7] + of its high half, so a 16-bit e << 7 lands the exponent where an f32 + e << 23 would and costs one register instead of two. The scale is a power of + two, so this stays bit-identical to the unfused convert + multiply.""" sc16 = e_u8.to(gl.uint16) << 7 lo = gl.inline_asm_elementwise( _DEQ_LO, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 @@ -523,7 +524,7 @@ def __init__( @gluon.jit def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 - (gfx950 has no bf16 multiply). ``sc``: raw UE8M0 byte (dsv4), f32 scale + (gfx950 has no bf16 multiply). sc: raw UE8M0 byte (dsv4), f32 scale (uniform/dsmla), or unused ("tensor": bare fp8 -> bf16 convert).""" if fmt.ASM_DEQ: # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. @@ -614,7 +615,7 @@ def _slots( """Index-list read -> (block, pos, valid), in whatever layout k_pos carries. A masked gl.load predicates on exec while a masked buffer_load folds the mask into the offset, so the unmasked paths clamp the read in-range and - mask the *score* instead (UNI_TILE); -1 sentinels are handled the same way.""" + mask the score instead (UNI_TILE); -1 sentinels are handled the same way.""" indices_ptr = seg.indices_ptr seg_start = seg.seg_start BLOCK_SIZE: gl.constexpr = seg.fmt.BLOCK_SIZE @@ -680,8 +681,8 @@ def _gather_full( k_rng_rope, ): """Gather one full fp8 tile, split from the LDS-write/MFMA so it issues an - iteration early. The prefetch is carried in raw fp8 — dequantized here it - would double the loop-carried registers; the consumer dequants chunked. + iteration early. The prefetch stays in raw fp8: dequantizing here would + double the loop-carried registers, so the consumer dequants in chunks. Returns (x, sc, k_rope, valid); unused slots carry a duplicate DCE removes.""" fmt = seg.fmt cs0 = seg.cs0 @@ -807,9 +808,9 @@ def _gather_full( @gluon.jit def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): """Write one prefetched tile into the LDS plane(s). Plane 0 is always the - full KV_DIM-wide dequant — "dsv4" gathers KV_DIM bytes too (the last 64 are - bf16 rope read as garbage fp8, overwritten by the slice-store below), which - keeps the gather pow-2 wide.""" + full KV_DIM-wide dequant: "dsv4" gathers KV_DIM bytes too, the last 64 being + bf16 rope read as garbage fp8 and overwritten by the slice-store below, so + the gather stays pow-2 wide.""" fmt = seg.fmt _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) if fmt.KIND == "dsv4": @@ -855,7 +856,7 @@ def _qkpv( if COL_MASK: if cfg.UNI_TILE and not cfg.HAS_INVALID: # Build the range mask directly in the MFMA layout; converting - # the slot-layout `valid` vector costs a cross-lane convert per + # the slot-layout valid vector costs a cross-lane convert per # tile. col_mask = ( k_start @@ -880,8 +881,8 @@ def _qkpv( S = gl.where(col_mask, S, neg_inf) # Online softmax in the base-2 exponent domain; m_i carries qk_scale # (= scale * log2e [* k_scale for "tensor"]). max commutes with a positive - # scale, so scale the row max instead of every element — the remaining - # `S * qk_scale - m_new` lowers to FMA, and -inf columns stay -inf. + # scale, so scale the row max instead of every element of S; what is left, + # S * qk_scale - m_new, lowers to one FMA, and -inf columns stay -inf. m_block = _rmax(S, 1) * qk_scale m_new = _max2(m_i, m_block) m_new = gl.where(m_new > neg_inf, m_new, 0.0) @@ -890,7 +891,7 @@ def _qkpv( l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(cfg.v_layout) # "tensor": V was staged as raw fp8 code points; apply the per-tensor scale - # on the small side (p), leaving l scale-free — out = sum(p*s*V)/l exactly. + # on the small side (p) and leave l scale-free: out = sum(p*s*V)/l exactly. if seg.fmt.KIND == "tensor": p = p * v_scale p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) @@ -1277,20 +1278,29 @@ def _process_segment( @gluon.jit(repr=_pa_decode_sparse_repr) def _pa_decode_sparse( - q_ptr, - main_cache_ptr, - main_cache_bf16_ptr, - main_indices_ptr, - main_indptr_ptr, - extra_cache_ptr, + # Shapes below: C = queries, H = num_heads, S = HEAD_SIZE (the V width), + # R = ROPE_DIM, nnz = total gathered tokens in a segment's index list. + q_ptr, # [C, H, S (+R when ROPE_SEPARATE)] bf16 + # One segment = a cache plus its index list. The cache is paged + # [num_blocks, BLOCK_SIZE, row] (BLOCK_SIZE = 1 for a flat pool), and + # indices[indptr[t]:indptr[t + 1]] are the rows query t attends to. The two + # cache pointers are the same allocation under different element types; + # which of them is live depends on the format (see Seg). + main_cache_ptr, # main (SWA) cache, u8 view + main_cache_bf16_ptr, # bf16 view of it, or the f32 scale pool ("uniform") + main_indices_ptr, # [nnz_main] int32 row ids + main_indptr_ptr, # [C + 1] int32 + extra_cache_ptr, # top-k segment; aliases main when HAS_EXTRA=False extra_cache_bf16_ptr, - extra_indices_ptr, - extra_indptr_ptr, - attn_sink_ptr, - out_ptr, - part_m_ptr, - part_l_ptr, - part_acc_ptr, + extra_indices_ptr, # [nnz_extra] int32 + extra_indptr_ptr, # [C + 1] int32 + attn_sink_ptr, # [H] f32, HAS_SINK only + out_ptr, # [C, H, S] bf16, written when NUM_SPLITS == 1 + # Split-K partials, written instead of out_ptr when NUM_SPLITS > 1 (unused + # placeholders otherwise). + part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain + part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum + part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view # ("dsmla"). None elides the argument, keeping other formats' kernarg # layouts unchanged. @@ -1437,10 +1447,10 @@ def _pa_decode_sparse( extra_len = 0 main_len = main_end - main_start - # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw `scale` + # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw scale # for the sink (a scaled-score-space logit). For "tensor" the cache's # k_scale joins the per-segment fold (max/exp2 commute with a positive - # scale) and the V-side scale is applied to p in the loop — both exact. + # scale) and the V-side scale hits p in the loop. Both are exact. RCP_LN2: gl.constexpr = 1.4426950408889634 qk_scale = scale * RCP_LN2 main_qk_scale = qk_scale @@ -1507,10 +1517,10 @@ def _pa_decode_sparse( rope_smem = kv_smem # never read as plane 1 in this geometry # ADAPTIVE_SPLITS: the host split count is sized for batch averages; in a - # ragged batch surplus programs would each gather a mostly-masked tile and - # write a full partial. Recompute from this query's own lengths and let - # surplus programs write a neutral partial (m = -inf) and leave — the - # reduce skips those, so their part_acc is never read. + # ragged batch the surplus programs would each gather a mostly-masked tile + # and write a full partial. Recompute from this query's own lengths and let + # those programs write a neutral partial (m = -inf) and leave. The reduce + # skips them, so their part_acc never has to be written. if ADAPTIVE_SPLITS: m_tiles = (main_len + BLOCK_K - 1) // BLOCK_K e_tiles = (extra_len + BLOCK_K - 1) // BLOCK_K @@ -1690,7 +1700,6 @@ def _pa_decode_sparse_reduce( NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, ADAPTIVE_SPLITS: gl.constexpr, - PART_LOAD_CACHE: gl.constexpr, ): """Split-KV combine: merge per-split partials, fold the sink, write the output. Identical for both geometries (partials are HEAD_SIZE = V wide). @@ -1728,14 +1737,14 @@ def _pa_decode_sparse_reduce( base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, - cache=PART_LOAD_CACHE, ) m_final = _max2(m_final, m_s) # m_s already in base-2 exponent domain if HAS_SINK: sink = gl.amd.cdna4.buffer_load( - ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf + ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf, cache=".cg", ).to(gl.float32) - m_final = _max2(m_final, sink * RCP_LN2) # lift sink to base-2 + scaled_sink = sink * RCP_LN2 + m_final = _max2(m_final, scaled_sink) # lift sink to base-2 # pass 2: weighted sums l_final = gl.zeros([BLOCK_M], gl.float32, layout=row_l) @@ -1744,33 +1753,30 @@ def _pa_decode_sparse_reduce( base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, - cache=PART_LOAD_CACHE, + cache=".cg", ) l_s = gl.amd.cdna4.buffer_load( ptr=part_l_ptr + base, offsets=h, mask=head_mask, other=0.0, - cache=PART_LOAD_CACHE, + cache=".cg", ) w = gl.exp2(m_s - m_final) l_final = l_final + w * l_s a_base = query_idx * pa_stride0 + s * pa_stride_s a_off = (a_base + h[:, None] * pa_stride_h + offs_d[None, :]).to(gl.int32) - # A bowed-out split's part_acc is uninitialized — mask the load rather - # than relying on w == 0 (0 * NaN is NaN). + # A bowed-out split's part_acc is uninitialized, so mask the load + # rather than relying on w == 0 (0 * NaN is NaN). if ADAPTIVE_SPLITS: acc_mask = head_mask[:, None] & (m_s > neg_inf)[:, None] else: acc_mask = head_mask[:, None] acc_s = gl.amd.cdna4.buffer_load( ptr=part_acc_ptr, offsets=a_off, mask=acc_mask, other=0.0, - cache=PART_LOAD_CACHE, + cache=".cg", ) acc = acc + w[:, None] * acc_s.to(gl.float32) if HAS_SINK: - sink = gl.amd.cdna4.buffer_load( - ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf - ).to(gl.float32) - l_final = l_final + gl.exp2(sink * RCP_LN2 - m_final) + l_final = l_final + gl.exp2(scaled_sink - m_final) # One reciprocal per row instead of a per-element f32 divide. one_over_l = 1.0 / l_final diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 7e796ad7b5a..6768325276f 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -558,46 +558,27 @@ def _pa_decode_sparse_gfx950_gluon( has_invalid=False, fp8_fnuz=False, ): - """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from ``cache.ndim``: + """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from cache.ndim: 3D [nb, block, ...] -> packed fp8_ds_mla (uint8: 448 NoPE fp8 e4m3 OCP + embedded UE8M0 per-64 scale + 64 RoPE bf16) or a bf16 - block cache; pass ``extra_*`` for the SWA+top-k two-loop, + block cache; pass extra_* for the SWA+top-k two-loop, else a single segment. - 2D [pages, D] -> uniform pool: fp8 (uint8) + ``cache_scales`` - [pages, D//64] fp32, or bf16 (``cache_scales`` None). + 2D [pages, D] -> uniform pool: fp8 (uint8) + cache_scales + [pages, D//64] fp32, or bf16 (cache_scales None). """ assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert DEVICE_ARCH == "gfx950", "gluon DSv4 decode kernel is gfx950-only" - # Tuned launch config (gfx950 / MI355), inlined. BLOCK_M = heads per MFMA M-tile; - # BLOCK_K = KV tile; num_warps = BLOCK_K // 16 (warps tile the dot-N, MFMA N=16). - BLOCK_M, BLOCK_K, MFMA_K, waves_per_eu = 16, 64, 16, 0 - # AITER_PA_DECODE_BLOCK_K: experiment override for the KV tile width. - # num_warps stays BLOCK_K//16 (warps tile the dot-N, MFMA N=16). - import os as _os - _bk = _os.environ.get("AITER_PA_DECODE_BLOCK_K") - if _bk: - BLOCK_K = int(_bk) - # AITER_PA_DECODE_BLOCK_M: heads per MFMA M-tile. 16 is both the MFMA M and the - # DSv4 head count; larger pads and quadruples the accumulator (see RESULTS.md). - _bm = _os.environ.get("AITER_PA_DECODE_BLOCK_M") - if _bm: - BLOCK_M = int(_bm) - # AITER_PA_DECODE_MFMA_K: 32 selects CDNA4's double-rate v_mfma_f32_16x16x32_bf16 - # (16 does 16x16x16, i.e. half the K per instruction). - _mk = _os.environ.get("AITER_PA_DECODE_MFMA_K") - if _mk: - MFMA_K = int(_mk) + # Tuned launch config (gfx950 / MI355). BLOCK_M = heads per MFMA M-tile, and 16 + # is both the MFMA M and the DSv4 head count; BLOCK_K = KV tile; num_warps = + # BLOCK_K // 16, because warps tile the dot-N and MFMA N = 16. + BLOCK_M, BLOCK_K, MFMA_K = 16, 64, 16 num_warps = max(1, BLOCK_K // 16) - # AITER_PA_DECODE_TW1: threads the NoPE gather spends on the head dim (16 B - # each). 32 issues a whole 512 B token row per instruction instead of 8 x 128 B - # quarters, which is what a scattered top-k gather wants; the cost is a longer - # per-lane slot vector. Measured best at 32 for every extra >= 128 (-10% at - # extra=128) and within noise at extra=8. - gather_tw1 = int(_os.environ.get("AITER_PA_DECODE_TW1", "32")) - # AITER_PA_DECODE_LDS_PAD: bf16 elements of padding after each kv_smem row. - lds_pad = int(_os.environ.get("AITER_PA_DECODE_LDS_PAD", "8")) - # AITER_PA_DECODE_NOPE_CHUNK: width of one fp8 dequant piece (0/unset = one shot). + # Threads the NoPE gather spends on the head dim, 16 B each: 32 requests a whole + # 512 B token row per instruction instead of four 128 B quarters, which is what + # a scattered top-k gather wants. The cost is a longer per-lane slot vector. + gather_tw1 = 32 + lds_pad = 8 # bf16 elements of padding after each kv_smem row NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -633,7 +614,6 @@ def _pa_decode_sparse_gfx950_gluon( main_block, extra_block = 1, 1 nope_dim = head_dim main_num_rows = extra_num_rows = cache.shape[0] - cache_bytes = max_addressable_bytes(cache) avg_main = indices.numel() / max(1, num_queries) # one segment; no extra avg_extra = 0.0 else: @@ -657,9 +637,6 @@ def _pa_decode_sparse_gfx950_gluon( nope_dim = NOPE_DIM main_num_rows = cache.shape[0] * cache.shape[1] extra_num_rows = extra_cache.shape[0] * extra_cache.shape[1] - cache_bytes = max( - max_addressable_bytes(cache), max_addressable_bytes(extra_cache) - ) avg_main = indices.numel() / max(1, num_queries) avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 @@ -671,115 +648,73 @@ def _pa_decode_sparse_gfx950_gluon( main_fmt = "dsv4" if main_is_fp8 else "bf16" extra_fmt = "dsv4" if extra_is_fp8 else "bf16" - # Largest power-of-2 (<=16) dividing BOTH page strides. Lets the kernel assert - # 16B-aligned row bases so a contiguous row gather can vectorize to dwordx4. - _s0 = int(cache.stride(0)); _s1 = int(extra_cache.stride(0)) + # Alignment hint for the page strides so row gathers can vectorize: the largest + # power of 2 (<= 16) dividing both. + s0, s1 = int(cache.stride(0)), int(extra_cache.stride(0)) cs0_align = 1 - for _a in (16, 8, 4, 2): - if _s0 % _a == 0 and _s1 % _a == 0: - cs0_align = _a + for a in (16, 8, 4, 2): + if s0 % a == 0 and s1 % a == 0: + cs0_align = a break - # AITER_PA_DECODE_CS0_ALIGN: override for re-measuring the hint (1 = off). As of - # the row/column gather split the hint is codegen-neutral -- the column offsets are - # compile-time constants, so vectorization no longer depends on knowing cs0's - # divisibility, and the emitted load mix is identical with it on and off. - _ca = _os.environ.get("AITER_PA_DECODE_CS0_ALIGN") - if _ca: - cs0_align = int(_ca) - # Gate each cache on its OWN span: one oversized cache must not disable the - # fast path for the other. (cache_bytes, the max of the two, is kept for the - # arch/format asserts below.) + + # Gate each cache on its own span: buffer_load carries a 32-bit offset, and one + # oversized cache must not drop the fast path for the other. The index lists are + # one int32 per gathered token, far under the limit even for a full batch. main_use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES extra_use_buffer_load = max_addressable_bytes(extra_cache) < MAX_BYTES - # The index lists are one int32 per gathered token, orders of magnitude under - # the 2 GB buffer offset limit even for a full batch, so they keep the fast path - # regardless of how big the KV caches are. idx_use_buffer_load = ( - _os.environ.get("AITER_PA_DECODE_IDX_BUF", "1") == "1" - and max_addressable_bytes(indices) < MAX_BYTES + max_addressable_bytes(indices) < MAX_BYTES and max_addressable_bytes(extra_indices) < MAX_BYTES ) use_buffer_load = main_use_buffer_load and extra_use_buffer_load - # AITER_PA_DECODE_FORCE_GLOBAL_LOAD emulates a >2GB production cache without - # allocating one: 1 = both caches, main = main only, extra = extra only. - _fg = _os.environ.get("AITER_PA_DECODE_FORCE_GLOBAL_LOAD", "") - if _fg == "1": - main_use_buffer_load = extra_use_buffer_load = False - elif _fg == "main": - main_use_buffer_load = False - elif _fg == "extra": - extra_use_buffer_load = False - use_buffer_load = main_use_buffer_load and extra_use_buffer_load HEAD_ALIGNED = num_heads % BLOCK_M == 0 heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M out = _check_out(out, q, torch.bfloat16) - # AITER_PA_DECODE_SPLITS: experiment override for the split-K count. - _sp = _os.environ.get("AITER_PA_DECODE_SPLITS") - def _pick_splits(hb): - if _sp: - return max(1, int(_sp)) if kv_splits is not None: return max(1, int(kv_splits)) - if _os.environ.get("AITER_PA_DECODE_SPLIT_POLICY", "occ") == "occ": - return _decode_num_splits_occ( - num_queries, hb, avg_main, avg_extra, BLOCK_K - ) - return _decode_num_splits(num_queries, hb, avg_main, avg_extra, BLOCK_K) + return _decode_num_splits_occ(num_queries, hb, avg_main, avg_extra, BLOCK_K) num_splits = _pick_splits(heads_blocks) # Short top-k is latency- and coverage-bound rather than throughput-bound: at - # <= 2 KV tiles per query the whole launch is a few dozen workgroups on a - # 256-CU part, so the win is MORE programs, not fatter ones. Halving BLOCK_M - # doubles heads_blocks, which is real parallelism -- unlike a finer split-K, - # which hands every split a partial tile (1.14-1.36x worse when forced). Both - # gates matter: measured 2-5% better at every concurrency on top-k 8/64, and - # 1.06-1.22x WORSE on top-k 272/1024, where the halved MFMA M dominates. - _tiles = max( + # <= 2 KV tiles per query the whole launch is a few dozen workgroups, so what + # helps is more programs, not fatter ones. Halving BLOCK_M doubles heads_blocks, + # which is real parallelism, unlike a finer split-K that hands every split a + # partial tile. It costs MFMA M, though, and doubling heads_blocks can push the + # launch past one workgroup per CU, which makes the split policy answer with + # fewer splits: take the halving only when it costs neither a split nor the + # one-workgroup-per-CU property. + tiles = max( math.ceil(avg_main / BLOCK_K) if avg_main > 0 else 1, math.ceil(avg_extra / BLOCK_K) if avg_extra > 0 else 1, ) if ( - _bm is None - and BLOCK_M == 16 - and num_heads % 8 == 0 - and _tiles <= 2 + num_heads % 8 == 0 + and tiles <= 2 and num_queries * heads_blocks * num_splits <= get_num_sms() ): - _hb8 = (num_heads + 7) // 8 - _ns8 = _pick_splits(_hb8) - # Doubling heads_blocks can push the launch past one workgroup per CU, and - # the split policy then answers with fewer splits. At C=128 short top-k that - # is 2 -> 1: split-K disappears and the shape goes 0.94x -> 1.11x. Take the - # halving only when it costs neither a split nor the occupancy property. - if _ns8 >= num_splits and num_queries * _hb8 * _ns8 <= get_num_sms(): + hb8 = (num_heads + 7) // 8 + ns8 = _pick_splits(hb8) + if ns8 >= num_splits and num_queries * hb8 * ns8 <= get_num_sms(): BLOCK_M = 8 HEAD_ALIGNED = num_heads % BLOCK_M == 0 - heads_blocks = _hb8 - num_splits = _ns8 + heads_blocks = hb8 + num_splits = ns8 if num_splits > 1: part_m = torch.empty( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device ) part_l = torch.empty_like(part_m) - # Split-K accumulator partials in bf16. They are ~31% of the kernel's HBM - # traffic and ATT puts 8.5% of all stall cycles on the 68 buffer_store of - # them, so halving both is worth 6-23% depending on split count. The - # mantissa bits it costs are far inside the error the fp8 KV format already - # carries: against an fp64 reference, max|delta|/max|ref| goes 3.517e-2 -> - # 3.567e-2 at C128A, and the bf16-vs-f32 delta (4-6e-3) is ~6x smaller than - # the error already present. Not used on the skip_reduce path, which hands - # part_acc back to the caller and must keep its dtype. - _pab = ( - _os.environ.get("AITER_PA_DECODE_PART_BF16", "1") == "1" - and not skip_reduce - ) + # bf16 partials halve both the split-K HBM traffic (~31% of the kernel's + # bytes) and the store count, and the mantissa bits that costs are well + # inside the error the fp8 KV format already carries. skip_reduce hands the + # partials back to the caller, so it keeps their dtype. part_acc = torch.empty( (num_queries, num_splits, num_heads, head_dim), - dtype=torch.bfloat16 if _pab else torch.float32, + dtype=torch.float32 if skip_reduce else torch.bfloat16, device=q.device, ) pm_stride0, pm_stride_s = part_m.stride(0), part_m.stride(1) @@ -792,96 +727,47 @@ def _pick_splits(hb): part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # 128 is the narrowest piece the gather layout can split to for free (below - # that a split falls inside a thread's 16-byte run, which Triton rejects), and - # the split is only trivial when the gather's warps tile dim 0. - # Row-wide gather chunks rows (4 pieces of BLOCK_K/4) instead of 128-wide - # column pieces; both leave a quarter of the tile's f32 expansion live. - # Split the axis that still has per-lane register repeats: columns when the - # gather leaves >=2 column repeats (TW1=8), rows otherwise. + # Dequant chunking: split the axis that still has per-lane register repeats, + # rows at gather_tw1=32 and 128-wide column pieces otherwise (128 is the + # narrowest free split -- below that a piece falls inside a thread's 16-byte + # run). Either way a quarter of the tile's f32 expansion is live at a time. col_reps = head_dim // (gather_tw1 * 16) chunk_axis = 1 if col_reps >= 4 else 0 - _nc = _os.environ.get("AITER_PA_DECODE_NOPE_CHUNK") - if chunk_axis == 0: - nope_chunk = int(_nc) if _nc else max(1, BLOCK_K // 4) - else: - nope_chunk = int(_nc) if _nc else min(128, head_dim) + nope_chunk = max(1, BLOCK_K // 4) if chunk_axis == 0 else min(128, head_dim) # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD - # threshold on gfx950's 512-VGPR file. Asked for on both load paths. - # - # This used to be gated on buffer_load, because a cache past buffer_load's 2 GB - # offset gathers through 64-bit addresses and used to land at ~321 VGPRs, where - # the cap bought nothing and cost ~150 scratch stores (+40% at C4A). Narrowing - # the scale gather brought that path to 283, close enough that the cap now - # reaches 256 with 17 spill slots instead of 148 -- worth 0.745x at extra=272 - # and 0.833x at extra=1024 on the 64-bit path. Re-measure this gate whenever the - # register picture moves. - if chunk_axis == 0 or nope_chunk < head_dim: - waves_per_eu = 2 - # num_warps is 4 -- one wave per SIMD per workgroup -- so the second occupancy - # slot exists only if a CU gets a second workgroup. When the entire launch is at - # most one workgroup per CU it cannot, and the 2-waves/EU register cap is then - # pure loss: it spills ~29 dwords across the per-tile softmax row-max reduction - # and buys no overlap. Measured on the main kernel at C in {4,16,32,64,128} x - # top-k in {64,272,1024}: 1 wave/EU + ASM_DEQ is 0.927-0.953x at every shape with - # launch <= CU, and 1.19-1.43x *worse* at every shape above it. + # threshold on gfx950's 512-VGPR file, so a CU can hold a second workgroup. + # num_warps is 4, one wave per SIMD, so that second slot only exists if the + # launch has more than one workgroup per CU; below that the cap buys no overlap + # and only spills. + waves_per_eu = 2 one_wg_per_cu = ( - use_buffer_load - and num_queries * heads_blocks * num_splits <= get_num_sms() + use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms() ) if one_wg_per_cu: waves_per_eu = 1 - _wp = _os.environ.get("AITER_PA_DECODE_WPEU") - if _wp: - waves_per_eu = int(_wp) # Cap the main segment's split count at whole BLOCK_K tiles: past that every # split's main range is a partial tile, which costs a full masked gather for # half the tokens. extra keeps all num_splits programs. main_splits = num_splits if has_extra and avg_main > 0: - main_splits = max(1, min(num_splits, int(math.ceil(avg_main / BLOCK_K)))) - _ms = _os.environ.get("AITER_PA_DECODE_SPLITS_MAIN") - if _ms: - main_splits = max(1, min(num_splits, int(_ms))) - - # AITER_PA_DECODE_ADAPTIVE: per-query split count decided in-kernel. Only - # meaningful when there is more than one split to give back. - adaptive_splits = ( - num_splits > 1 - and _os.environ.get("AITER_PA_DECODE_ADAPTIVE", "1") == "1" - ) - - # AITER_PA_DECODE_ASM_DEQ: fuse the fp8 x E8M0 dequant into - # v_cvt_scalef32_pk_bf16_fp8 via inline asm. gfx950 only, packed OCP fp8 only - # (UNIFORM scales are f32 and FNUZ is a different encoding), and bit-identical - # there because the scale is a power of two. - # Default follows the occupancy decision: the fused convert is a shorter - # dependent chain but a wider live range, so it pays only where the extra - # registers are free (see one_wg_per_cu). 0/1 forces it either way. - asm_deq = ( - _os.environ.get("AITER_PA_DECODE_ASM_DEQ", "1" if one_wg_per_cu else "0") == "1" - and not UNIFORM - and not fp8_fnuz - and main_is_fp8 - ) - - # AITER_PA_DECODE_GRID_ORDER: which launch axis varies fastest. Dim 0 is the - # fastest-varying and XCD assignment is round-robin over the linear workgroup id, - # so this decides what shares an XCD's L2. "qsh" is the historical order. - _go = _os.environ.get("AITER_PA_DECODE_GRID_ORDER", "qsh") - _ax = {"q": num_queries, "s": num_splits, "h": heads_blocks} - assert sorted(_go) == ["h", "q", "s"], f"bad AITER_PA_DECODE_GRID_ORDER {_go!r}" - # AITER_PA_DECODE_UNI_TILE: run the partial last tile through the same body as the - # full ones instead of peeling a masked copy of it. Gluon inlines, so the peeled - # copy is a second ~1000-instruction gather+dequant+MFMA body whose register demand - # spills the tile loop: worth 0.900-0.907x on the buffer path and 0.943-0.972x on - # the 64-bit path at >=2 tiles/split, scaling with tiles per split. 0 restores the - # peeled tail. - uni_tile = _os.environ.get("AITER_PA_DECODE_UNI_TILE", "1") == "1" - - grid = tuple(_ax[c] for c in _go) + main_splits = max(1, min(num_splits, math.ceil(avg_main / BLOCK_K))) + + # Per-query split count decided in-kernel; only meaningful when there is more + # than one split to give back. + adaptive_splits = num_splits > 1 + + # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm: + # packed OCP fp8 only (uniform scales are f32, fnuz is a different encoding), + # bit-identical there because the scale is a power of two. It trades a shorter + # dependent chain for a wider live range, so it pays only where the extra + # registers are free -- the same condition that lifts the occupancy cap. + asm_deq = one_wg_per_cu and not UNIFORM and not fp8_fnuz and main_is_fp8 + + # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear + # workgroup id, so the axis order decides what shares an XCD's L2. + grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, cache, @@ -936,9 +822,12 @@ def _pick_splits(hb): LDS_PAD=lds_pad, NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, - PART_STORE_CACHE=_os.environ.get('AITER_PA_DECODE_PART_ST', ''), - GRID_ORDER=_go, - UNI_TILE=uni_tile, + PART_STORE_CACHE="", + GRID_ORDER="qsh", + # The partial last tile rides the full-tile body. Gluon inlines, so a peeled + # masked copy would be a second gather+dequant+MFMA body, and its register + # demand spills the tile loop. + UNI_TILE=True, MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, ASM_DEQ=asm_deq, @@ -956,20 +845,10 @@ def _pick_splits(hb): if skip_reduce: return part_acc, part_m, part_l - # The combine is pure bandwidth over [num_queries, num_heads, head_dim] f32, - # so size its head tile for workgroup count, not for the attention kernel's - # BLOCK_M: at BLOCK_M=num_heads there is one workgroup per query (64 on a - # 256-CU part). One head per workgroup gives num_queries*num_heads of them and - # enough concurrent waves to hide the partial-load latency. - red_block_m = 1 - _rbm = _os.environ.get("AITER_PA_DECODE_REDUCE_BLOCK_M") - if _rbm: - red_block_m = int(_rbm) - red_warps = 1 - _rw = _os.environ.get("AITER_PA_DECODE_REDUCE_WARPS") - if _rw: - red_warps = int(_rw) - rgrid = (num_queries, (num_heads + red_block_m - 1) // red_block_m) + # One head per reduce workgroup: the combine is pure bandwidth, so size the grid + # for coverage rather than for the attention kernel's tile. num_queries*num_heads + # workgroups give enough concurrent waves to hide the partial loads. + rgrid = (num_queries, num_heads) _pa_decode_sparse_reduce_gfx950[rgrid]( part_m, part_l, @@ -986,11 +865,10 @@ def _pick_splits(hb): num_heads, HAS_SINK=has_sink, HEAD_SIZE=head_dim, - BLOCK_M=red_block_m, + BLOCK_M=1, NUM_SPLITS=num_splits, - HEAD_ALIGNED=num_heads % red_block_m == 0, + HEAD_ALIGNED=True, ADAPTIVE_SPLITS=adaptive_splits, - PART_LOAD_CACHE=_os.environ.get("AITER_PA_DECODE_PART_LD", ""), - num_warps=red_warps, + num_warps=1, ) return out diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index 0edc949d72a..282920706c5 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -323,7 +323,6 @@ def sparse_mla_decode_fwd( NUM_SPLITS=num_splits, HEAD_ALIGNED=True, ADAPTIVE_SPLITS=num_splits > 1, - PART_LOAD_CACHE="", num_warps=1, ) return out From f5e0f14092e1079470759bc19b51a65e8c7ec058 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Wed, 19 Aug 2026 06:23:20 +0000 Subject: [PATCH 28/51] remove leftover stuff --- aiter/ops/triton/attention/pa_decode_sparse.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 6768325276f..06f32f1ea89 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -282,9 +282,6 @@ def pa_decode_sparse( # mismatches between 2D (m/l) and 3D (acc) loads. reduce_num_warps = 1 if use_gluon else 4 reduce_waves_per_eu = 4 if use_gluon else 1 - _rw = _os.environ.get("AITER_PA_DECODE_REDUCE_WARPS") - if _rw and use_gluon: - reduce_num_warps = int(_rw) USE_EXP2 = True # Infer KV_SPLITS from inputs when caller doesn't override. From c3b541d4408a364152156ab53c675fd7f00f8b2a Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 20 Aug 2026 03:53:09 +0000 Subject: [PATCH 29/51] [GLUON][GFX950][MLA] fp8_mfma: keep the tile in fp8 through both dots 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) --- .../gfx950/attention/pa_decode_sparse.py | 100 +++++++++++++++--- .../ops/triton/attention/sparse_mla_decode.py | 26 ++++- 2 files changed, 112 insertions(+), 14 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index fa8ddbde32b..b05d2c65d95 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -252,6 +252,8 @@ class Cfg: HAS_INVALID: gl.constexpr HEAD_ALIGNED: gl.constexpr IDX_BUFFER_LOAD: gl.constexpr + FP8_MFMA: gl.constexpr # "tensor" only: keep the tile fp8 all the way into + # the matrix core instead of dequantizing to bf16 # operator layouts qk_layout: gl.constexpr pv_layout: gl.constexpr @@ -284,6 +286,7 @@ def __init__( HAS_INVALID, HEAD_ALIGNED, IDX_BUFFER_LOAD, + FP8_MFMA=False, ): self.BLOCK_M = gl.constexpr(BLOCK_M) self.BLOCK_K = gl.constexpr(BLOCK_K) @@ -291,7 +294,6 @@ def __init__( self.ROPE_DIM = gl.constexpr(ROPE_DIM) self.ROPE_SEPARATE = gl.constexpr(ROPE_SEPARATE) self.QK_DIM = gl.constexpr(KV_DIM + (ROPE_DIM if ROPE_SEPARATE else 0)) - self.MFMA_K = gl.constexpr(MFMA_K) self.NUM_WARPS = gl.constexpr(NUM_WARPS) self.GATHER_TW1 = gl.constexpr(GATHER_TW1) self.LDS_PAD = gl.constexpr(LDS_PAD) @@ -299,6 +301,13 @@ def __init__( self.HAS_INVALID = gl.constexpr(HAS_INVALID) self.HEAD_ALIGNED = gl.constexpr(HEAD_ALIGNED) self.IDX_BUFFER_LOAD = gl.constexpr(IDX_BUFFER_LOAD) + self.FP8_MFMA = gl.constexpr(FP8_MFMA) + # fp8 operands need v_mfma_f32_16x16x32_fp8_fp8: K=32 is the only shape + # the backend has for plain fp8, and it divides KV_DIM, ROPE_DIM and + # BLOCK_K alike, so one layout still covers both dots. + if FP8_MFMA: + MFMA_K = 32 + self.MFMA_K = gl.constexpr(MFMA_K) self.qk_layout = gl.constexpr( gl.amd.AMDMFMALayout( @@ -812,15 +821,23 @@ def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): bf16 rope read as garbage fp8 and overwritten by the slice-store below, so the gather stays pow-2 wide.""" fmt = seg.fmt - _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) - if fmt.KIND == "dsv4": - kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) - elif fmt.KIND == "dsmla": - rope_smem.store(k_rope) - elif fmt.KIND == "tensor": + if cfg.FP8_MFMA: + # No dequant at all: the matrix core reads the cache's own code points. + # The per-tensor scale is folded outside the loop (qk_scale / the output), + # so what lands in LDS is exactly what the gather returned. + kv_smem.store(x_u8.to(gl.float8e4nv, bitcast=True)) if cfg.ROPE_SEPARATE: - rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) - # "uniform": the whole head is one fp8 tile; nothing else to store. + rope_smem.store(k_rope.to(gl.float8e4nv, bitcast=True)) + else: + _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) + if fmt.KIND == "dsv4": + kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) + elif fmt.KIND == "dsmla": + rope_smem.store(k_rope) + elif fmt.KIND == "tensor": + if cfg.ROPE_SEPARATE: + rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) + # "uniform": the whole head is one fp8 tile; nothing else to store. @gluon.jit @@ -892,9 +909,17 @@ def _qkpv( v = kv_smem.load(cfg.v_layout) # "tensor": V was staged as raw fp8 code points; apply the per-tensor scale # on the small side (p) and leave l scale-free: out = sum(p*s*V)/l exactly. - if seg.fmt.KIND == "tensor": + if seg.fmt.KIND == "tensor" and not cfg.FP8_MFMA: p = p * v_scale - p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) + if cfg.FP8_MFMA: + # p goes to e4m3 as-is. A power-of-two pre-scale to lift the small end out + # of e4m3's subnormals was measured and is worth nothing: the output error + # is set by the 3-bit mantissa on the p values near 1, and the subnormal + # tail contributes too little to register (1 / 16 / 256 agree to three + # digits, and 256 is already half of what saturates at p = 1). + p_dot = gl.convert_layout(p.to(gl.float8e4nv), cfg.p_layout) + else: + p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.pv_layout)) acc = acc * alpha_pv[:, None] acc = gl.amd.cdna4.mfma(p_dot, v, acc) @@ -1360,6 +1385,7 @@ def _pa_decode_sparse( IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, + FP8_MFMA: gl.constexpr = False, ): """One program = (query, split, head-block). Two-loop: main (SWA) then extra (top-k). NUM_SPLITS==1 writes the output directly; otherwise stores @@ -1383,6 +1409,18 @@ def _pa_decode_sparse( ) gl.static_assert((not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", "ASM_DEQ is the dsv4 E8M0 dequant") + # FP8_MFMA keeps the tile in fp8 through both dots. That only works where the + # scale is a single positive scalar per cache (it folds outside the loop) and + # the code points are OCP e4m3, which is what the matrix core reads. + gl.static_assert( + (not FP8_MFMA) or (MAIN_FMT == "tensor" and (not HAS_EXTRA or EXTRA_FMT == "tensor")), + "FP8_MFMA requires the per-tensor fp8 format on every segment", + ) + gl.static_assert(not (FP8_MFMA and FP8_FNUZ), "FP8_MFMA is OCP e4m3 only") + gl.static_assert( + not (FP8_MFMA and HAS_EXTRA), + "FP8_MFMA defers the V-side scale to the epilogue, so it needs one segment", + ) # Row bases are block*cs0 + pos*TOK with runtime block/pos, so divisibility # analysis sees 1-byte alignment unless the driver vouches for cs0. if CS0_ALIGN > 1: @@ -1408,6 +1446,7 @@ def _pa_decode_sparse( HAS_INVALID, HEAD_ALIGNED, IDX_BUFFER_LOAD, + FP8_MFMA, ) main_fmt = Fmt( cfg, @@ -1478,6 +1517,13 @@ def _pa_decode_sparse( q = gl.amd.cdna4.buffer_load( ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 ) + if FP8_MFMA: + # One e4m3 scale for this program's whole Q tile (nope + rope), so the + # fold below is a single extra factor on qk_scale. Each program folds its + # own scale, and m/l stay in the true-score domain, so split-K programs + # remain comparable in the reduce. + E4M3_MAX: gl.constexpr = 448.0 + q_amax = gl.max(gl.max(gl.abs(q).to(gl.float32), axis=1), axis=0) q_dot = gl.convert_layout(q, cfg.q_layout) if ROPE_SEPARATE: offs_d_qr = gl.arange(0, ROPE_DIM, layout=gl.SliceLayout(0, cfg.blocked_q)) @@ -1494,6 +1540,26 @@ def _pa_decode_sparse( else: q_rope_dot = q_dot # unused (single-plane QK) -> DCE'd + if FP8_MFMA: + if ROPE_SEPARATE: + q_amax = gl.maximum( + q_amax, gl.max(gl.max(gl.abs(q_rope).to(gl.float32), axis=1), axis=0) + ) + q_amax = gl.maximum(q_amax, 1e-30) + q_rcp = E4M3_MAX / q_amax + q_dot = gl.convert_layout( + (q.to(gl.float32) * q_rcp).to(gl.float8e4nv), cfg.q_layout + ) + if ROPE_SEPARATE: + q_rope_dot = gl.convert_layout( + (q_rope.to(gl.float32) * q_rcp).to(gl.float8e4nv), cfg.q_layout + ) + else: + q_rope_dot = q_dot + q_scale = q_amax / E4M3_MAX + main_qk_scale = main_qk_scale * q_scale + extra_qk_scale = extra_qk_scale * q_scale + # head mask in pv-slice layout (for output / partial masking) offs_m_pv = gl.arange(0, BLOCK_M, layout=gl.SliceLayout(1, cfg.pv_layout)) h_pv = h_off + offs_m_pv @@ -1506,12 +1572,14 @@ def _pa_decode_sparse( l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout)) acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=cfg.pv_layout) + # fp8 planes are half the bytes of the bf16 staging they replace. + SMEM_DT: gl.constexpr = gl.float8e4nv if FP8_MFMA else gl.bfloat16 kv_smem = gl.allocate_shared_memory( - gl.bfloat16, [BLOCK_K, HEAD_SIZE], cfg.kv_shared + SMEM_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared ) if ROPE_SEPARATE: rope_smem = gl.allocate_shared_memory( - gl.bfloat16, [BLOCK_K, ROPE_DIM], cfg.rope_shared + SMEM_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared ) else: rope_smem = kv_smem # never read as plane 1 in this geometry @@ -1614,6 +1682,12 @@ def _pa_decode_sparse( rope_smem, ) + if FP8_MFMA: + # The fp8 PV dot ran on the cache's raw code points, so the V-side + # per-tensor scale comes off here -- once per program instead of once per + # tile. l is untouched, so out = acc*s/l is what the bf16 path computes. + acc = acc * main_v_scale + # Move the row reductions into pv-slice space for output/partials. m_pv = gl.convert_layout(m_i, gl.SliceLayout(1, cfg.pv_layout)) l_pv = gl.convert_layout(l_i, gl.SliceLayout(1, cfg.pv_layout)) diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index 282920706c5..6d7ec0ec812 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -109,6 +109,18 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): raise ValueError(f"unrecognized sparse-MLA kv cache: {tuple(kv.shape)} {kv.dtype}") +def _lds_pad_for(fp8_mfma: bool) -> int: + """kv_smem row padding, in *elements* of the staged type. + + The pad exists to offset which banks a transposed K read lands on, and that + is a byte property: _LDS_PAD is tuned for the bf16 staging, so an fp8 plane + needs twice as many elements to move the row pitch by the same bytes. + Measured at the GLM shape (C=64, topk=2048): carrying bf16's 8 over to fp8 + costs 1.08x, and 16 / 32 / 64 all land at 0.86-0.89x. + """ + return _LDS_PAD * 2 if fp8_mfma else _LDS_PAD + + def sparse_mla_decode_fwd( q: torch.Tensor, kv_buffer: torch.Tensor, @@ -121,6 +133,7 @@ def sparse_mla_decode_fwd( kv_splits: int | None = None, skip_reduce: bool = False, has_invalid: bool = False, + fp8_mfma: bool = False, out: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse (top-k gathered) MLA decode, separated-rope geometry. @@ -140,6 +153,11 @@ def sparse_mla_decode_fwd( instead of launching the combine. has_invalid: index stream carries -1 sentinels (masked out). Default False: vLLM's index converter emits none (it maps -1 to slot 0). + fp8_mfma: keep the tile in fp8 through both dots instead of staging it + as bf16 -- flat per-tensor OCP fp8 cache only. q still arrives bf16 + and is quantized to e4m3 in the kernel, per program, so its scale + folds into that program's qk_scale. Faster (0.75-0.96x) but drops + accuracy to the asm kernel's level, so it is a deliberate trade. out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. Returns: @@ -160,6 +178,11 @@ def sparse_mla_decode_fwd( fmt, cache, alt, scl, block_size, fp8_fnuz = _infer_cache_format( kv_buffer, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale ) + if fp8_mfma and (fmt != "tensor" or fp8_fnuz): + raise ValueError( + "fp8_mfma needs the flat per-tensor OCP fp8 cache (got " + f"fmt={fmt!r}, fp8_fnuz={fp8_fnuz})" + ) kv_indices = _as_int32_contiguous_1d(kv_indices) kv_indptr = _as_int32_contiguous_1d(kv_indptr) assert kv_indptr.numel() == num_queries + 1 @@ -277,7 +300,7 @@ def sparse_mla_decode_fwd( HEAD_ALIGNED=head_aligned, MFMA_K=_MFMA_K, GATHER_TW1=_GATHER_TW1, - LDS_PAD=_LDS_PAD, + LDS_PAD=_lds_pad_for(fp8_mfma), NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", @@ -291,6 +314,7 @@ def sparse_mla_decode_fwd( IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, + FP8_MFMA=fp8_mfma, num_warps=num_warps, waves_per_eu=waves_per_eu, ) From 4815a0412e0d6193a00704bcbda5fbe80e0aed4c Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 20 Aug 2026 04:01:16 +0000 Subject: [PATCH 30/51] [GLUON][GFX950][MLA] Accept caller-quantized fp8 q, the asm calling convention 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) --- .../gfx950/attention/pa_decode_sparse.py | 31 +++++++++++++++---- .../ops/triton/attention/sparse_mla_decode.py | 20 +++++++++++- 2 files changed, 44 insertions(+), 7 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index b05d2c65d95..7747a76fa16 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1386,6 +1386,11 @@ def _pa_decode_sparse( HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, FP8_MFMA: gl.constexpr = False, + # q already quantized to e4m3 by the caller, with the scalar f32 scale it + # was quantized with -- the calling convention aiter's asm mla_decode_fwd + # uses, where vLLM passes layer._q_scale. + q_scl_ptr=None, + Q_FP8: gl.constexpr = False, ): """One program = (query, split, head-block). Two-loop: main (SWA) then extra (top-k). NUM_SPLITS==1 writes the output directly; otherwise stores @@ -1517,11 +1522,11 @@ def _pa_decode_sparse( q = gl.amd.cdna4.buffer_load( ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 ) - if FP8_MFMA: - # One e4m3 scale for this program's whole Q tile (nope + rope), so the - # fold below is a single extra factor on qk_scale. Each program folds its - # own scale, and m/l stay in the true-score domain, so split-K programs - # remain comparable in the reduce. + if FP8_MFMA and not Q_FP8: + # bf16 q: quantize it here, one e4m3 scale for this program's whole Q + # tile (nope + rope), so the fold below is a single extra factor on + # qk_scale. Each program folds its own scale and m/l stay in the + # true-score domain, so split-K programs remain comparable in the reduce. E4M3_MAX: gl.constexpr = 448.0 q_amax = gl.max(gl.max(gl.abs(q).to(gl.float32), axis=1), axis=0) q_dot = gl.convert_layout(q, cfg.q_layout) @@ -1540,7 +1545,21 @@ def _pa_decode_sparse( else: q_rope_dot = q_dot # unused (single-plane QK) -> DCE'd - if FP8_MFMA: + if Q_FP8: + # Nothing to quantize; pick up the caller's scale and fold it exactly + # the way the cache's k_scale is folded. + q_scale = gl.load(q_scl_ptr) + if not FP8_MFMA: + # bf16 dots want bf16 operands, and fp8 -> bf16 is exact (3 mantissa + # bits into 8), so a quantized q costs nothing extra on this path. + q_dot = gl.convert_layout(q.to(gl.bfloat16), cfg.q_layout) + if ROPE_SEPARATE: + q_rope_dot = gl.convert_layout(q_rope.to(gl.bfloat16), cfg.q_layout) + else: + q_rope_dot = q_dot + main_qk_scale = main_qk_scale * q_scale + extra_qk_scale = extra_qk_scale * q_scale + elif FP8_MFMA: if ROPE_SEPARATE: q_amax = gl.maximum( q_amax, gl.max(gl.max(gl.abs(q_rope).to(gl.float32), axis=1), axis=0) diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index 6d7ec0ec812..40ab68fde2d 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -134,6 +134,7 @@ def sparse_mla_decode_fwd( skip_reduce: bool = False, has_invalid: bool = False, fp8_mfma: bool = False, + q_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse (top-k gathered) MLA decode, separated-rope geometry. @@ -158,6 +159,9 @@ def sparse_mla_decode_fwd( and is quantized to e4m3 in the kernel, per program, so its scale folds into that program's qk_scale. Faster (0.75-0.96x) but drops accuracy to the asm kernel's level, so it is a deliberate trade. + q_scale: scalar f32 scale q was quantized with. Required when q is + fp8, unused otherwise. Matches aiter's asm mla_decode_fwd, where + vLLM passes the layer's calibrated ``_q_scale``. out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. Returns: @@ -165,7 +169,19 @@ def sparse_mla_decode_fwd( """ assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert arch_info.get_arch() == "gfx950", "sparse_mla_decode_fwd is gfx950-only" - assert q.dtype == torch.bfloat16 + q_is_fp8 = q.dtype == torch.float8_e4m3fn + assert q.dtype == torch.bfloat16 or q_is_fp8, ( + f"q must be bf16 or float8_e4m3fn, got {q.dtype}" + ) + 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: + raise ValueError("fp8 q needs q_scale (the scale it was quantized with)") + if q_scale.numel() != 1: + raise ValueError(f"q_scale must be a scalar, got {tuple(q_scale.shape)}") + q_scale = q_scale.reshape(1).to(torch.float32).contiguous() num_queries, num_heads, d_qk = q.shape assert d_qk == kv_lora_rank + qk_rope_head_dim, ( f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" @@ -315,6 +331,8 @@ def sparse_mla_decode_fwd( HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, FP8_MFMA=fp8_mfma, + q_scl_ptr=q_scale, + Q_FP8=q_is_fp8, num_warps=num_warps, waves_per_eu=waves_per_eu, ) From 9d188bfb11b0810c3263d1e43cdd7c453176cde9 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 20 Aug 2026 04:35:37 +0000 Subject: [PATCH 31/51] [GLUON][GFX950][MLA] Give the sparse-MLA path its own split-K policy _decode_num_splits_occ is tuned for the DSv4 two-segment SWA + top-k launch and over-splits badly at MLA shapes. Below one workgroup per CU it asks for `tiles` splits, so at C=1 topk=2048 it picks 32 for 32 tiles: every split owns a single tile, 32 programs each pay the full prologue, and the reduce walks 32 partials. That is the whole reason the port lost to the asm kernel at small batch -- 1.19-1.24x at C<=4, where the kernel itself was never the problem. _mla_num_splits replaces it for this path only (DSv4 keeps the shared policy untouched). Two rules, measured on MI355X across C in 1..256 x topk in {512, 2048} x H in {8, 16}, sweeping the split count in every cell: * Below one workgroup per CU: split to fill the machine but never past 8. Past that, splits stop buying parallelism and start paying prologue and reduce traffic. 8 is optimal or within 10% in every cell measured; the 16 that topk=2048 prefers at C=8-16 is worth only 4-10%, not worth a second rule. * At or above one workgroup per CU: the launch is already full, so a split only buys the second occupancy slot and must earn its extra partial round trip -- which takes ~16 tiles. At C=256 that is the difference between one split and two, and topk=512 (8 tiles) wants none: splitting anyway costs 22%. Also capped by the tile count, so a short top-k cannot ask for splits that own no tiles at all. vs aiter's asm mla_decode_fwd, fp8 q + fp8 dots, arms interleaved (before -> after this commit): topk=2048 C=1 1.243 -> 0.981 C=16 0.856 -> 0.850 C=2 1.205 -> 0.964 C=32 0.838 -> 0.797 C=4 1.186 -> 0.958 C=64 0.900 -> 0.879 C=8 0.908 -> 0.905 C=128 1.019 -> 1.013 (H=16) C=256 1.087 -> 1.083 topk=512 C=256 1.025 -> 0.853 C=64 0.740 -> 0.751 (H=16) C=128 0.811 -> 0.812 C=16 0.749 -> 0.710 H=8 is better still throughout (asm repeat-interleaves 8 heads to 16, this path runs BLOCK_M=8 natively): topk=512 is now a win in every measured cell, 0.628-0.853x, and topk=2048 wins C=1..64 at 0.743-0.981x. op_tests 128 passed / 10 skipped; DSv4 identity gate unaffected, still 11/11. Co-Authored-By: Claude Opus 5 (1M context) --- .../ops/triton/attention/sparse_mla_decode.py | 38 +++++++++++++++++-- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index 40ab68fde2d..f6a47ecc38b 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -16,6 +16,8 @@ kernel; <1% accuracy effect). """ +import math + import torch from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( @@ -26,7 +28,6 @@ ) from aiter.ops.triton.attention.pa_decode_sparse import ( _as_int32_contiguous_1d, - _decode_num_splits_occ, ) from aiter.ops.triton.utils._triton import arch_info from aiter.ops.triton.utils.common_utils import max_addressable_bytes @@ -109,6 +110,37 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): raise ValueError(f"unrecognized sparse-MLA kv cache: {tuple(kv.shape)} {kv.dtype}") +def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int: + """Split-K count for the sparse-MLA decode. + + Separate from the DSv4 policy (_decode_num_splits_occ): that one is tuned for + a two-segment SWA + top-k launch and, measured here, over-splits badly at MLA + shapes -- at C=1, topk=2048 it asks for one split per tile, so 32 programs + each pay the full prologue and the reduce walks 32 partials, which costs + 1.23-1.28x. + + Two rules, both measured on MI355X over C in 1..256 x topk in {512, 2048} x + H in {8, 16}: + + * Below one workgroup per CU, split to fill the machine but never past 8. + Splits beyond that stop buying parallelism and start paying prologue and + reduce traffic: 8 is optimal or within 10% in every cell measured, and + the 16 that topk=2048 prefers at C=8-16 is worth only 4-10%. + * At or above one workgroup per CU the launch is already full, so a split + only buys the second occupancy slot and has to earn its extra partial + round trip -- it takes ~16 tiles to do that. At C=256 this is the + difference between one split and two: topk=512 (8 tiles) wants none, and + splitting anyway costs 22%. + """ + num_sms = get_num_sms() + base_wg = max(1, num_queries * heads_blocks) + cta_cap = max(1, (2 * num_sms) // base_wg) + tiles = max(1, math.ceil(avg_topk / _BLOCK_K)) + if base_wg >= num_sms: + return max(1, min(cta_cap, tiles // 16)) + return max(1, min(cta_cap, tiles, 8)) + + def _lds_pad_for(fp8_mfma: bool) -> int: """kv_smem row padding, in *elements* of the staged type. @@ -230,9 +262,7 @@ def sparse_mla_decode_fwd( if kv_splits is not None: num_splits = max(1, int(kv_splits)) else: - num_splits = _decode_num_splits_occ( - num_queries, heads_blocks, avg_topk, 0.0, _BLOCK_K - ) + num_splits = _mla_num_splits(num_queries, heads_blocks, avg_topk) if num_splits > 1: part_m = torch.empty( From 29872800c1b5f9a864138f40f08e6e88d8cd8f77 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 20 Aug 2026 04:58:44 +0000 Subject: [PATCH 32/51] [GLUON][GFX950][MLA] Select the dot precision, and adapt q to it Replaces the fp8_mfma bool with dot_precision="bf16" (default) or "fp8", and makes the wrapper reconcile it with what the caller actually passed instead of requiring them to line it up by hand. q is adapted to the choice: bf16 q + "bf16" nothing to do fp8 q + "bf16" widened in-kernel, exact (3 mantissa bits into 8), so a caller who already quantized for the asm path can use the accurate dots without re-materializing q bf16 q + "fp8" quantized here, the way production does it: one scaled_fp8_quant over [C, H*d_qk]. A supplied q_scale is used as the calibrated static scale, otherwise a per-tensor amax is taken fp8 q + "fp8" passed straight through Rejections, all of them the same constraint, that the fp8 path needs one positive scalar scale per cache because that is what folds outside the tile loop: "fp8" + fp8_ds_mla its scale varies per 128 elements along the head dim, the QK contraction axis, so it cannot fold out; and scaled MFMA cannot express the PV side either, where the scale varies per token against MX's one per 32 "fp8" + bf16 cache nothing to feed the matrix core "fp8" + fnuz the matrix core reads OCP e4m3 fp8 q, no q_scale the fold has no scale to apply q_scale, bf16 q + "bf16" dots nothing would consume it Also trims the comments added across this fp8 work to the house style. All six accepted combinations verified against the bf16 path, all six rejections verified to raise. op_tests 128 passed / 10 skipped, DSv4 identity gate 11/11 byte-identical, C=64 topk=2048 still 0.839x (H=8) / 0.892x (H=16) of the asm kernel. Co-Authored-By: Claude Opus 5 (1M context) --- .../gfx950/attention/pa_decode_sparse.py | 60 ++++--- .../ops/triton/attention/sparse_mla_decode.py | 150 +++++++++++++----- 2 files changed, 135 insertions(+), 75 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 7747a76fa16..06fdd0de976 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -252,8 +252,8 @@ class Cfg: HAS_INVALID: gl.constexpr HEAD_ALIGNED: gl.constexpr IDX_BUFFER_LOAD: gl.constexpr - FP8_MFMA: gl.constexpr # "tensor" only: keep the tile fp8 all the way into - # the matrix core instead of dequantizing to bf16 + FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's + # own fp8 instead of dequantizing to bf16 # operator layouts qk_layout: gl.constexpr pv_layout: gl.constexpr @@ -302,9 +302,9 @@ def __init__( self.HEAD_ALIGNED = gl.constexpr(HEAD_ALIGNED) self.IDX_BUFFER_LOAD = gl.constexpr(IDX_BUFFER_LOAD) self.FP8_MFMA = gl.constexpr(FP8_MFMA) - # fp8 operands need v_mfma_f32_16x16x32_fp8_fp8: K=32 is the only shape - # the backend has for plain fp8, and it divides KV_DIM, ROPE_DIM and - # BLOCK_K alike, so one layout still covers both dots. + # K=32 is the only shape the backend offers for plain fp8 operands, and + # it divides KV_DIM, ROPE_DIM and BLOCK_K alike, so one layout still + # covers both dots. if FP8_MFMA: MFMA_K = 32 self.MFMA_K = gl.constexpr(MFMA_K) @@ -822,9 +822,9 @@ def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): the gather stays pow-2 wide.""" fmt = seg.fmt if cfg.FP8_MFMA: - # No dequant at all: the matrix core reads the cache's own code points. - # The per-tensor scale is folded outside the loop (qk_scale / the output), - # so what lands in LDS is exactly what the gather returned. + # No dequant: the scale is folded outside the loop (qk_scale on the K + # side, the accumulator on the V side), so what lands in LDS is exactly + # what the gather returned. kv_smem.store(x_u8.to(gl.float8e4nv, bitcast=True)) if cfg.ROPE_SEPARATE: rope_smem.store(k_rope.to(gl.float8e4nv, bitcast=True)) @@ -912,11 +912,9 @@ def _qkpv( if seg.fmt.KIND == "tensor" and not cfg.FP8_MFMA: p = p * v_scale if cfg.FP8_MFMA: - # p goes to e4m3 as-is. A power-of-two pre-scale to lift the small end out - # of e4m3's subnormals was measured and is worth nothing: the output error - # is set by the 3-bit mantissa on the p values near 1, and the subnormal - # tail contributes too little to register (1 / 16 / 256 agree to three - # digits, and 256 is already half of what saturates at p = 1). + # p goes to e4m3 as it is. Pre-scaling it by a power of two to lift the + # small end out of e4m3's subnormals is worth nothing: the error is set + # by the p values near 1, and 1 / 16 / 256 agree to three digits. p_dot = gl.convert_layout(p.to(gl.float8e4nv), cfg.p_layout) else: p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) @@ -1386,9 +1384,9 @@ def _pa_decode_sparse( HAS_INVALID: gl.constexpr, FP8_FNUZ: gl.constexpr, FP8_MFMA: gl.constexpr = False, - # q already quantized to e4m3 by the caller, with the scalar f32 scale it - # was quantized with -- the calling convention aiter's asm mla_decode_fwd - # uses, where vLLM passes layer._q_scale. + # q already quantized to e4m3 by the caller, plus the scalar f32 scale it + # was quantized with. This is the calling convention aiter's asm + # mla_decode_fwd uses, where vLLM passes layer._q_scale. q_scl_ptr=None, Q_FP8: gl.constexpr = False, ): @@ -1414,9 +1412,9 @@ def _pa_decode_sparse( ) gl.static_assert((not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", "ASM_DEQ is the dsv4 E8M0 dequant") - # FP8_MFMA keeps the tile in fp8 through both dots. That only works where the - # scale is a single positive scalar per cache (it folds outside the loop) and - # the code points are OCP e4m3, which is what the matrix core reads. + # The fp8 path needs one positive scalar scale per cache, since that is what + # folds outside the loop, and OCP e4m3 code points, which is what the matrix + # core reads. gl.static_assert( (not FP8_MFMA) or (MAIN_FMT == "tensor" and (not HAS_EXTRA or EXTRA_FMT == "tensor")), "FP8_MFMA requires the per-tensor fp8 format on every segment", @@ -1523,10 +1521,10 @@ def _pa_decode_sparse( ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 ) if FP8_MFMA and not Q_FP8: - # bf16 q: quantize it here, one e4m3 scale for this program's whole Q - # tile (nope + rope), so the fold below is a single extra factor on - # qk_scale. Each program folds its own scale and m/l stay in the - # true-score domain, so split-K programs remain comparable in the reduce. + # bf16 q: quantize here, one e4m3 scale for this program's whole Q tile + # (nope and rope), so the fold below is one extra factor on qk_scale. + # m/l stay in the true-score domain, so split-K programs stay comparable + # in the reduce. E4M3_MAX: gl.constexpr = 448.0 q_amax = gl.max(gl.max(gl.abs(q).to(gl.float32), axis=1), axis=0) q_dot = gl.convert_layout(q, cfg.q_layout) @@ -1546,12 +1544,12 @@ def _pa_decode_sparse( q_rope_dot = q_dot # unused (single-plane QK) -> DCE'd if Q_FP8: - # Nothing to quantize; pick up the caller's scale and fold it exactly - # the way the cache's k_scale is folded. + # Nothing to quantize; fold the caller's scale the way k_scale is + # folded. q_scale = gl.load(q_scl_ptr) if not FP8_MFMA: - # bf16 dots want bf16 operands, and fp8 -> bf16 is exact (3 mantissa - # bits into 8), so a quantized q costs nothing extra on this path. + # fp8 -> bf16 is exact (3 mantissa bits into 8), so a quantized q + # costs nothing extra on the bf16 dots. q_dot = gl.convert_layout(q.to(gl.bfloat16), cfg.q_layout) if ROPE_SEPARATE: q_rope_dot = gl.convert_layout(q_rope.to(gl.bfloat16), cfg.q_layout) @@ -1591,7 +1589,7 @@ def _pa_decode_sparse( l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout)) acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=cfg.pv_layout) - # fp8 planes are half the bytes of the bf16 staging they replace. + # An fp8 plane is half the bytes of the bf16 staging it replaces. SMEM_DT: gl.constexpr = gl.float8e4nv if FP8_MFMA else gl.bfloat16 kv_smem = gl.allocate_shared_memory( SMEM_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared @@ -1702,9 +1700,9 @@ def _pa_decode_sparse( ) if FP8_MFMA: - # The fp8 PV dot ran on the cache's raw code points, so the V-side - # per-tensor scale comes off here -- once per program instead of once per - # tile. l is untouched, so out = acc*s/l is what the bf16 path computes. + # The fp8 PV dot ran on raw code points, so the V-side scale comes off + # here, once per program instead of once per tile. l is untouched, so + # out = acc*s/l is what the bf16 path computes. acc = acc * main_v_scale # Move the row reductions into pv-slice space for output/partials. diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla_decode.py index f6a47ecc38b..68561ad386e 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla_decode.py @@ -68,7 +68,7 @@ def _check_out(out, q, kv_lora_rank): def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): """-> (fmt, cache, alt_ptr, scl_ptr, block_size, fp8_fnuz). Pointer roles follow the kernel's Seg contract (see its docstring).""" - if kv.ndim == 4: # [slots, 1, 1, R] -- the asm mla_decode_fwd view + if kv.ndim == 4: # [slots, 1, 1, R], the asm mla_decode_fwd view assert kv.shape[1] == 1 and kv.shape[2] == 1 kv = kv.reshape(kv.shape[0], kv.shape[3]) elif kv.ndim == 3 and kv.shape[2] == d_qk: @@ -113,24 +113,17 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int: """Split-K count for the sparse-MLA decode. - Separate from the DSv4 policy (_decode_num_splits_occ): that one is tuned for - a two-segment SWA + top-k launch and, measured here, over-splits badly at MLA - shapes -- at C=1, topk=2048 it asks for one split per tile, so 32 programs - each pay the full prologue and the reduce walks 32 partials, which costs - 1.23-1.28x. - - Two rules, both measured on MI355X over C in 1..256 x topk in {512, 2048} x - H in {8, 16}: - - * Below one workgroup per CU, split to fill the machine but never past 8. - Splits beyond that stop buying parallelism and start paying prologue and - reduce traffic: 8 is optimal or within 10% in every cell measured, and - the 16 that topk=2048 prefers at C=8-16 is worth only 4-10%. - * At or above one workgroup per CU the launch is already full, so a split - only buys the second occupancy slot and has to earn its extra partial - round trip -- it takes ~16 tiles to do that. At C=256 this is the - difference between one split and two: topk=512 (8 tiles) wants none, and - splitting anyway costs 22%. + Separate from the DSv4 policy in _decode_num_splits_occ, which is tuned for a + two-segment SWA + top-k launch and over-splits here: at C=1 topk=2048 it asks + for one split per tile, so 32 programs each pay the full prologue and the + reduce walks 32 partials, costing 1.23-1.28x. + + Below one workgroup per CU, split to fill the machine but never past 8. Past + that, splits stop buying parallelism and start paying prologue and reduce + traffic; 8 is optimal or within 10% everywhere measured. At or above one + workgroup per CU the launch is already full, so a split only buys the second + occupancy slot and needs ~16 tiles to earn its extra partial round trip. At + C=256 that decides one split versus two, and topk=512 wants none. """ num_sms = get_num_sms() base_wg = max(1, num_queries * heads_blocks) @@ -141,16 +134,54 @@ def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int return max(1, min(cta_cap, tiles, 8)) -def _lds_pad_for(fp8_mfma: bool) -> int: - """kv_smem row padding, in *elements* of the staged type. +_E4M3_MAX = 448.0 - The pad exists to offset which banks a transposed K read lands on, and that - is a byte property: _LDS_PAD is tuned for the bf16 staging, so an fp8 plane - needs twice as many elements to move the row pitch by the same bytes. - Measured at the GLM shape (C=64, topk=2048): carrying bf16's 8 over to fp8 - costs 1.08x, and 16 / 32 / 64 all land at 0.86-0.89x. + +def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool: + """True for the fp8 matrix-core path. + + Every rejection below is the same constraint: the fp8 path needs the cache + scale to be one positive scalar, because that is what folds outside the tile + loop (K-side into qk_scale, V-side onto the accumulator). A scale that varies + along a dot's contraction axis cannot be folded out at all. """ - return _LDS_PAD * 2 if fp8_mfma else _LDS_PAD + if dot_precision not in ("bf16", "fp8"): + raise ValueError( + f"dot_precision must be 'bf16' or 'fp8', got {dot_precision!r}" + ) + if dot_precision == "bf16": + return False + if fmt == "dsmla": + raise ValueError( + "dot_precision='fp8' does not support the fp8_ds_mla cache. Its " + "scale varies per 128 elements along the head dim, the QK " + "contraction axis, so it cannot fold outside the tile loop, and " + "scaled MFMA cannot express the PV side either (one scale per 32 " + "contraction elements, where this needs one per token). " + "Use dot_precision='bf16'." + ) + if fmt == "bf16": + raise ValueError( + "dot_precision='fp8' needs an fp8 cache; this one is bf16, and " + "quantizing the pool per call would cost more than the dots save." + ) + if fp8_fnuz: + raise ValueError( + "dot_precision='fp8' is OCP e4m3 only; this cache is fnuz, a " + "different encoding from what the matrix core reads." + ) + return True + + +def _lds_pad_for(fp8_dots: bool) -> int: + """kv_smem row padding, in elements of the staged type. + + The pad shifts which banks a transposed K read lands on, which is a byte + property, so an fp8 plane needs twice the elements to move the row pitch by + the same bytes. Carrying bf16's value over to fp8 costs 1.08x at C=64 + topk=2048; 16, 32 and 64 all land at 0.86-0.89x. + """ + return _LDS_PAD * 2 if fp8_dots else _LDS_PAD def sparse_mla_decode_fwd( @@ -165,7 +196,7 @@ def sparse_mla_decode_fwd( kv_splits: int | None = None, skip_reduce: bool = False, has_invalid: bool = False, - fp8_mfma: bool = False, + dot_precision: str = "bf16", q_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: @@ -186,14 +217,25 @@ def sparse_mla_decode_fwd( instead of launching the combine. has_invalid: index stream carries -1 sentinels (masked out). Default False: vLLM's index converter emits none (it maps -1 to slot 0). - fp8_mfma: keep the tile in fp8 through both dots instead of staging it - as bf16 -- flat per-tensor OCP fp8 cache only. q still arrives bf16 - and is quantized to e4m3 in the kernel, per program, so its scale - folds into that program's qk_scale. Faster (0.75-0.96x) but drops - accuracy to the asm kernel's level, so it is a deliberate trade. - q_scale: scalar f32 scale q was quantized with. Required when q is - fp8, unused otherwise. Matches aiter's asm mla_decode_fwd, where - vLLM passes the layer's calibrated ``_q_scale``. + dot_precision: what the QK and PV matrix-core ops run in. + + "bf16" (default): the KV tile is dequantized to bf16 on its way + into LDS and both dots are bf16. Works with every cache format. + "fp8": the cache's own code points go to the fp8 matrix core with no + dequant, and the per-tensor scale folds outside the tile loop. + 0.75-0.96x of "bf16", but q and p both round-trip through e4m3, + so accuracy lands at roughly the asm kernel's level. Needs the + flat per-tensor OCP fp8 cache; see _resolve_dot_precision. + + q is adapted to the choice. bf16 q is quantized here when "fp8" is + asked for, fp8 q is passed straight through, and fp8 q under "bf16" + dots is widened in-kernel, which is exact. + q_scale: scalar f32. Required when q arrives already fp8 (the scale it + was quantized with; the aiter asm convention, where vLLM passes + the layer's calibrated ``_q_scale``). Optional when q is bf16 and + dot_precision="fp8": given, it is used as the static quantization + scale, matching production's scaled_fp8_quant(q, layer._q_scale); + omitted, a per-tensor amax is taken here. out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. Returns: @@ -202,9 +244,12 @@ def sparse_mla_decode_fwd( assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert arch_info.get_arch() == "gfx950", "sparse_mla_decode_fwd is gfx950-only" q_is_fp8 = q.dtype == torch.float8_e4m3fn - assert q.dtype == torch.bfloat16 or q_is_fp8, ( - f"q must be bf16 or float8_e4m3fn, got {q.dtype}" - ) + if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError( + f"q must be bf16 or float8_e4m3fn, got {q.dtype}" + + (" (fnuz is a different encoding from what the matrix core reads)" + if "fnuz" in str(q.dtype) else "") + ) 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 @@ -226,10 +271,27 @@ def sparse_mla_decode_fwd( fmt, cache, alt, scl, block_size, fp8_fnuz = _infer_cache_format( kv_buffer, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale ) - if fp8_mfma and (fmt != "tensor" or fp8_fnuz): + fp8_dots = _resolve_dot_precision(dot_precision, fmt, fp8_fnuz) + if fp8_dots and not q_is_fp8: + # Quantize the way production does, one scaled_fp8_quant over + # [C, H*d_qk]. A caller-supplied q_scale is the layer's calibrated + # static scale; without one, take a per-tensor amax. + if q_scale is None: + q_scale = ( + (q.detach().float().abs().amax() / _E4M3_MAX) + .clamp_min(torch.finfo(torch.float32).tiny) + .reshape(1) + ) + q = ( + (q.float() / q_scale).clamp(-_E4M3_MAX, _E4M3_MAX) + .to(torch.float8_e4m3fn) + .contiguous() + ) + q_is_fp8 = True + elif q_scale is not None and not q_is_fp8: raise ValueError( - "fp8_mfma needs the flat per-tensor OCP fp8 cache (got " - f"fmt={fmt!r}, fp8_fnuz={fp8_fnuz})" + "q_scale was given but q is bf16 and dot_precision='bf16', so " + "nothing would use it" ) kv_indices = _as_int32_contiguous_1d(kv_indices) kv_indptr = _as_int32_contiguous_1d(kv_indptr) @@ -346,7 +408,7 @@ def sparse_mla_decode_fwd( HEAD_ALIGNED=head_aligned, MFMA_K=_MFMA_K, GATHER_TW1=_GATHER_TW1, - LDS_PAD=_lds_pad_for(fp8_mfma), + LDS_PAD=_lds_pad_for(fp8_dots), NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", @@ -360,7 +422,7 @@ def sparse_mla_decode_fwd( IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, - FP8_MFMA=fp8_mfma, + FP8_MFMA=fp8_dots, q_scl_ptr=q_scale, Q_FP8=q_is_fp8, num_warps=num_warps, From bf3c53c90eaf3656c8b309c5867e083a0eac78a8 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 20 Aug 2026 19:22:24 +0000 Subject: [PATCH 33/51] restructure asm constexprs --- .../gfx950/attention/pa_decode_sparse.py | 28 +++++-------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 06fdd0de976..c233aaf2a9e 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -23,10 +23,7 @@ "dsmla" vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope (656 B rows); requires ROPE_SEPARATE. -Compile-time state is bundled into aggregates instead of long parameter lists: -Cfg (geometry, layouts, behavior knobs), Fmt (per-segment cache format), Seg -(runtime pointers/bounds). Two-loop (SWA + top-k) or a single segment; 2D and -3D (split-K + reduce) share one kernel. Launchers: +Launchers: aiter/ops/triton/attention/pa_decode_sparse.py (DSv4 / uniform pool) and aiter/ops/triton/attention/sparse_mla_decode.py (separated-rope MLA). """ @@ -187,17 +184,6 @@ def _split_ax(x, AXIS: gl.constexpr): return a, b -# fp8 (e4m3 OCP) x E8M0 -> bf16, one instruction per 2 elements. Emitted as two -# single-output asm calls, not one two-output blob: a single instruction reads -# all its sources before writing dst, so no output needs early clobber. The -# blob form does, and that forced =&v liveness spills a kernel with no headroom. -_DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") -_DEQ_HI: gl.constexpr = gl.constexpr( - "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" -) -_DEQ_CONS: gl.constexpr = gl.constexpr("=v,v,v") - - @gluon.jit def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): """[BLOCK_K, W8/2] int16 (4 packed fp8) + raw E8M0 byte -> [BLOCK_K, W8] bf16. @@ -206,6 +192,12 @@ def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): of its high half, so a 16-bit e << 7 lands the exponent where an f32 e << 23 would and costs one register instead of two. The scale is a power of two, so this stays bit-identical to the unfused convert + multiply.""" + + _DEQ_LO: gl.constexpr = gl.constexpr("v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2") + _DEQ_HI: gl.constexpr = gl.constexpr( + "v_cvt_scalef32_pk_bf16_fp8 $0, $1, $2 op_sel:[1,0,0]" + ) + _DEQ_CONS: gl.constexpr = gl.constexpr("=v,v,v") sc16 = e_u8.to(gl.uint16) << 7 lo = gl.inline_asm_elementwise( _DEQ_LO, _DEQ_CONS, [x16, sc16], dtype=gl.bfloat16, is_pure=True, pack=2 @@ -223,12 +215,6 @@ def _deq_asm(x16, e_u8, W8: gl.constexpr, out_l: gl.constexpr): return gl.convert_layout(both, out_l, assert_trivial=True) -# --------------------------------------------------------------------------- -# Aggregates: Cfg (shared compile-time), Fmt (per-segment compile-time), -# Seg (per-segment runtime). Built once at the kernel entry. -# --------------------------------------------------------------------------- - - @aggregate @strip_annotate class Cfg: From f78942dd15837236bafac3ef38092b71440128ba Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Fri, 21 Aug 2026 18:10:51 +0000 Subject: [PATCH 34/51] Enable async direct-to-LDS gather for gfx950 sparse MLA Adds the fp8 direct-to-LDS tile pipeline and turns it on by default, fixes the KV gather's cache policy, and renames sparse_mla_decode_fwd to sparse_mla_fwd since it serves prefill and decode alike. Co-Authored-By: Claude Opus 5 (1M context) --- .../gfx950/attention/pa_decode_sparse.py | 323 ++++++++++++++++-- .../ops/triton/attention/pa_decode_sparse.py | 2 +- .../{sparse_mla_decode.py => sparse_mla.py} | 123 ++++--- ...parse_mla_decode.py => test_sparse_mla.py} | 53 +-- 4 files changed, 400 insertions(+), 101 deletions(-) rename aiter/ops/triton/attention/{sparse_mla_decode.py => sparse_mla.py} (81%) rename op_tests/triton_tests/attention/{test_sparse_mla_decode.py => test_sparse_mla.py} (75%) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index c233aaf2a9e..b6d03641c16 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -25,7 +25,7 @@ Launchers: aiter/ops/triton/attention/pa_decode_sparse.py (DSv4 / uniform pool) and -aiter/ops/triton/attention/sparse_mla_decode.py (separated-rope MLA). +aiter/ops/triton/attention/sparse_mla.py (separated-rope MLA). """ from triton.experimental import gluon @@ -53,25 +53,31 @@ def _rmax(x, axis): @gluon.jit -def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None): +def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None, + CACHE: gl.constexpr = ".cg"): """Gather rows[i] + col[j]. row is the per-token offset in ptr's element units; col a small compile-time arange. Keeping them apart resolves one pointer per token on the 64-bit path (the column offset folds into the - load's immediate) instead of a 64-bit add per element.""" + load's immediate) instead of a 64-bit add per element. + + CACHE is the gather's cache policy. ".cg" lowers to `sc0 nt` (evict-first), which + is wrong for a KV gather: a 576/584 B row spans five cache lines and shares its end + line with the rope read of the same row, so evict-first turns one line fill into + two. The MLA launcher passes an empty modifier (a plain cached load) instead.""" if USE_BUFFER_LOAD: return gl.amd.cdna4.buffer_load( ptr=ptr, offsets=row.to(gl.int32)[:, None] + col.to(gl.int32)[None, :], mask=mask, other=other, - cache=".cg", + cache=CACHE, ) row_ptr = ptr + row.to(gl.int64) return gl.load( row_ptr[:, None] + col[None, :], mask=mask, other=other, - cache_modifier=".cg", + cache_modifier=CACHE, ) @@ -240,6 +246,18 @@ class Cfg: IDX_BUFFER_LOAD: gl.constexpr FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's # own fp8 instead of dequantizing to bf16 + # Cache policy per load site. The KV gather wants to be cached (its lines are + # shared between neighbouring rows and between a row's latent and rope + # halves); the index list is read exactly once per program, so it is the + # stream that actually wants ".cg" (sc0 nt). + GATHER_CACHE: gl.constexpr + IDX_CACHE: gl.constexpr + # Gather straight into LDS (no register staging). FP8_MFMA only: the bf16 path + # stages dequantized values, which pass through registers by definition. + # Single-buffered; the launcher decides when it is on. + ASYNC_LDS: gl.constexpr + RELAXED_LOAD: gl.constexpr # read LDS with the syncedViaAsyncWait hint + ROPE_VEC: gl.constexpr # bytes per lane in the rope plane's copy # operator layouts qk_layout: gl.constexpr pv_layout: gl.constexpr @@ -249,6 +267,10 @@ class Cfg: v_layout: gl.constexpr # memory layouts gather_l: gl.constexpr + async_l: gl.constexpr + async_rope_l: gl.constexpr + slot_a_l: gl.constexpr + slot_rope_a_l: gl.constexpr gather_rope_l: gl.constexpr gather16_l: gl.constexpr slot_l: gl.constexpr @@ -273,6 +295,12 @@ def __init__( HEAD_ALIGNED, IDX_BUFFER_LOAD, FP8_MFMA=False, + GATHER_CACHE=".cg", + IDX_CACHE="", + ASYNC_LDS=False, + RELAXED_LOAD=True, + ROPE_VEC=16, + PAD_INTERVAL=1024, ): self.BLOCK_M = gl.constexpr(BLOCK_M) self.BLOCK_K = gl.constexpr(BLOCK_K) @@ -288,6 +316,11 @@ def __init__( self.HEAD_ALIGNED = gl.constexpr(HEAD_ALIGNED) self.IDX_BUFFER_LOAD = gl.constexpr(IDX_BUFFER_LOAD) self.FP8_MFMA = gl.constexpr(FP8_MFMA) + self.GATHER_CACHE = gl.constexpr(GATHER_CACHE) + self.IDX_CACHE = gl.constexpr(IDX_CACHE) + self.ASYNC_LDS = gl.constexpr(ASYNC_LDS) + self.RELAXED_LOAD = gl.constexpr(RELAXED_LOAD) + self.ROPE_VEC = gl.constexpr(ROPE_VEC) # K=32 is the only shape the backend offers for plain fp8 operands, and # it divides KV_DIM, ROPE_DIM and BLOCK_K alike, so one layout still # covers both dots. @@ -374,6 +407,47 @@ def __init__( ) ) + # LDS-DMA takes one local address per warp (lane i writes warp_base + + # i*VEC), so the source layout must map lane i to LDS offset i*VEC: that + # pins threads_per_warp and makes a warp cover 64*VEC contiguous bytes. + # Padding therefore cannot sit inside a row; the interval must be a multiple + # of the warp's run so a pad only shifts whole runs. + AVEC = 16 # 128-bit LDS-DMA + ATW = KV_DIM // AVEC + self.async_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, AVEC], + threads_per_warp=[64 // ATW, ATW], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + RTW = ROPE_DIM // ROPE_VEC + self.async_rope_l = gl.constexpr( + gl.BlockedLayout( + size_per_thread=[1, ROPE_VEC], + threads_per_warp=[64 // RTW, RTW], + warps_per_cta=[NUM_WARPS, 1], + order=[1, 0], + ) + ) + self.slot_a_l = gl.constexpr(gl.SliceLayout(1, self.async_l.value)) + self.slot_rope_a_l = gl.constexpr( + gl.SliceLayout(1, self.async_rope_l.value) + ) + if ASYNC_LDS: + self.kv_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[PAD_INTERVAL, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] + ) + ) + # 64 * ROPE_VEC is this plane's warp run: the smallest legal interval. + self.rope_shared = gl.constexpr( + gl.PaddedSharedLayout.with_identity_for( + [[64 * ROPE_VEC, LDS_PAD]], [BLOCK_K, ROPE_DIM], [1, 0] + ) + ) + @aggregate @strip_annotate @@ -620,19 +694,23 @@ def _slots( in_range = k_pos < hi if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load( - ptr=indices_ptr + seg_start, offsets=k_pos, mask=in_range, other=-1 + ptr=indices_ptr + seg_start, offsets=k_pos, mask=in_range, other=-1, + cache=cfg.IDX_CACHE, ) else: - slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1) + slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1, + cache_modifier=cfg.IDX_CACHE) valid = in_range & (slot >= 0) & (slot < num_rows) slot = gl.where(valid, slot, 0) else: # hi >= 1 whenever UNI_TILE runs (guarded by n_full > 0). off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos if IDX_BUFFER_LOAD: - slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off) + slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off, + cache=cfg.IDX_CACHE) else: - slot = gl.load(indices_ptr + seg_start + off) + slot = gl.load(indices_ptr + seg_start + off, + cache_modifier=cfg.IDX_CACHE) valid = (k_pos < hi) if UNI_TILE else (slot >= 0) if HAS_INVALID: if UNI_TILE: @@ -645,18 +723,123 @@ def _slots( def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): """QK scores for one tile; ROPE_SEPARATE chains a second MFMA over the rope plane (MFMA accumulates natively, so KV_DIM + ROPE_DIM is two dots).""" - k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + # the async_wait already ordered the copy; stops the backend re-inserting + # a conservative vmcnt(0) before every LDS read + k = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem.permute([1, 0]), cfg.k_layout + ) + else: + k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] + if cfg.ASYNC_LDS: + k = k.to(gl.float8e4nv, bitcast=True) # raw cache bytes; layout-preserving S = gl.amd.cdna4.mfma( q_dot, k, gl.zeros([cfg.BLOCK_M, cfg.BLOCK_K], gl.float32, layout=cfg.qk_layout), ) if cfg.ROPE_SEPARATE: - k_rope = rope_smem.permute([1, 0]).load(cfg.k_layout) # [ROPE_DIM, BLOCK_K] + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + k_rope = gl.amd.cdna4.async_copy.load_shared_relaxed( + rope_smem.permute([1, 0]), cfg.k_layout + ) + else: + k_rope = rope_smem.permute([1, 0]).load(cfg.k_layout) + if cfg.ASYNC_LDS: + k_rope = k_rope.to(gl.float8e4nv, bitcast=True) S = gl.amd.cdna4.mfma(q_rope_dot, k_rope, S) return S +@gluon.jit +def _tile_rows(cfg, seg, k_start, seg_hi, k_rng, ROPE: gl.constexpr): + """Per-token row offsets for one tile, in k_rng's layout. Clamped like UNI_TILE + (the score mask kills the duplicate), so an over-fetched tile is addressable.""" + fmt = seg.fmt + cs0 = seg.cs0 + if not fmt.USE_BUFFER_LOAD: + cs0 = cs0.to(gl.int64) # >2 GB cache: 64-bit gather offsets + bg, pg, _ = _slots(cfg, seg, k_start + k_rng, seg_hi, 0, False, cfg.UNI_TILE) + row = bg * cs0 + pg * fmt.TOK_EL + if ROPE: + row = row + cfg.KV_DIM + return row + + +@gluon.jit +def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, + CACHE: gl.constexpr): + """One plane of one tile, global -> LDS, no register staging.""" + if USE_BUFFER_LOAD: + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dst, + seg.cache_ptr, + row.to(gl.int32)[:, None] + col.to(gl.int32)[None, :], + cache_modifier=CACHE, + ) + else: + row_ptr = seg.cache_ptr + row.to(gl.int64) + gl.amd.cdna4.async_copy.global_load_to_shared( + dst, row_ptr[:, None] + col[None, :], cache_modifier=CACHE + ) + + +@gluon.jit +def _copy_tile(cfg, seg, kv_smem, rope_smem, row_l, row_r, offs_l, offs_r): + """One commit group = one tile (both planes), so wait_group counts tiles.""" + _copy_plane(seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, + cfg.GATHER_CACHE) + if cfg.ROPE_SEPARATE: + _copy_plane(seg, rope_smem, row_r, offs_r, seg.fmt.USE_BUFFER_LOAD, + cfg.GATHER_CACHE) + gl.amd.cdna4.async_copy.commit_group() + + +@gluon.jit +def _async_segment( + cfg, + seg, + q_dot, + q_rope_dot, + lo, + hi, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, +): + """Direct-to-LDS tile walk, one LDS buffer. + + vmcnt is one in-order FIFO, so the wait goes before any plain load in the body: a + plain load outstanding past the newest copy group would drag `wait_group` into + waiting for that copy too. The body has no masked variant: the slot read is + clamped and the partial tail falls out of the score mask (UNI_TILE). + """ + BK: gl.constexpr = cfg.BLOCK_K + offs_l = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.async_l)) + offs_r = gl.arange(0, cfg.ROPE_DIM, layout=gl.SliceLayout(0, cfg.async_rope_l)) + rng_l = gl.arange(0, BK, layout=cfg.slot_a_l) + rng_r = gl.arange(0, BK, layout=cfg.slot_rope_a_l) + n_full = (hi - lo + BK - 1) // BK + for i in range(n_full): + _copy_tile( + cfg, seg, kv_smem, rope_smem, + _tile_rows(cfg, seg, lo + i * BK, hi, rng_l, False), + _tile_rows(cfg, seg, lo + i * BK, hi, rng_r, True), + offs_l, offs_r, + ) + gl.amd.cdna4.async_copy.wait_group(0) + m_i, l_i, acc = _qkpv_lds( + cfg, seg, None, q_dot, q_rope_dot, m_i, l_i, acc, head_mask, + qk_scale, v_scale, kv_smem, rope_smem, lo + i * BK, hi, + ) + return m_i, l_i, acc + + # --------------------------------------------------------------------------- # Prefetched fp8 pipeline: _gather_full issues tile N+1's loads while _qkpv # stages and dots tile N. @@ -705,7 +888,8 @@ def _gather_full( # Per-tensor scale is folded outside the loop (qk_scale / p), so this # is a bare gather; the K-only rope tail follows when separated. x_u8 = _cache_load( - seg.cache_ptr, bg * cs0 + pg * fmt.TOK_EL, offs_full, fmt.USE_BUFFER_LOAD + seg.cache_ptr, bg * cs0 + pg * fmt.TOK_EL, offs_full, + fmt.USE_BUFFER_LOAD, CACHE=cfg.GATHER_CACHE, ) sc = x_u8 # no scale vector -> DCE'd if cfg.ROPE_SEPARATE: @@ -723,6 +907,7 @@ def _gather_full( bgr * cs0 + pgr * fmt.TOK_EL + cfg.KV_DIM, offs_rope, fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) else: k_rope = x_u8 # rope lives inside plane 0 -> DCE'd @@ -848,8 +1033,33 @@ def _qkpv( seg_hi=0, ): """Stage a prefetched fp8 tile into LDS, then QK -> softmax -> PV.""" - neg_inf = float("-inf") _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem) + return _qkpv_lds( + cfg, seg, valid, q_dot, q_rope_dot, m_i, l_i, acc, head_mask, + qk_scale, v_scale, kv_smem, rope_smem, k_start, seg_hi, + ) + + +@gluon.jit +def _qkpv_lds( + cfg, + seg, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start=0, + seg_hi=0, +): + """QK -> softmax -> PV over a tile already resident in LDS.""" + neg_inf = float("-inf") S = _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem) # UNI_TILE folds the tile's range test into the score mask, which is what # makes the partial last tile correct without a peeled masked body. @@ -892,15 +1102,17 @@ def _qkpv( p = gl.exp2(S * qk_scale - m_new[:, None]) alpha = gl.exp2(m_i - m_new) l_new = l_i * alpha + gl.sum(p, axis=1) - v = kv_smem.load(cfg.v_layout) + if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: + v = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem, cfg.v_layout) + else: + v = kv_smem.load(cfg.v_layout) + if cfg.ASYNC_LDS: + v = v.to(gl.float8e4nv, bitcast=True) # "tensor": V was staged as raw fp8 code points; apply the per-tensor scale # on the small side (p) and leave l scale-free: out = sum(p*s*V)/l exactly. if seg.fmt.KIND == "tensor" and not cfg.FP8_MFMA: p = p * v_scale if cfg.FP8_MFMA: - # p goes to e4m3 as it is. Pre-scaling it by a power of two to lift the - # small end out of e4m3's subnormals is worth nothing: the error is set - # by the p values near 1, and 1 / 16 / 256 agree to three digits. p_dot = gl.convert_layout(p.to(gl.float8e4nv), cfg.p_layout) else: p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) @@ -1375,6 +1587,13 @@ def _pa_decode_sparse( # mla_decode_fwd uses, where vLLM passes layer._q_scale. q_scl_ptr=None, Q_FP8: gl.constexpr = False, + # Defaults keep every existing launch byte-identical; the MLA launcher opts in. + GATHER_CACHE: gl.constexpr = ".cg", + IDX_CACHE: gl.constexpr = "", + ASYNC_LDS: gl.constexpr = False, + RELAXED_LOAD: gl.constexpr = True, + ROPE_VEC: gl.constexpr = 16, + PAD_INTERVAL: gl.constexpr = 1024, ): """One program = (query, split, head-block). Two-loop: main (SWA) then extra (top-k). NUM_SPLITS==1 writes the output directly; otherwise stores @@ -1406,6 +1625,13 @@ def _pa_decode_sparse( "FP8_MFMA requires the per-tensor fp8 format on every segment", ) gl.static_assert(not (FP8_MFMA and FP8_FNUZ), "FP8_MFMA is OCP e4m3 only") + # The direct-to-LDS path stages raw code points, needs the clamped (branch-free) + # tail, and carries no per-tile validity vector. + gl.static_assert( + (not ASYNC_LDS) or (FP8_MFMA and UNI_TILE and not HAS_INVALID + and not HAS_EXTRA), + "ASYNC_LDS requires FP8_MFMA + UNI_TILE, no -1 sentinels, one segment", + ) gl.static_assert( not (FP8_MFMA and HAS_EXTRA), "FP8_MFMA defers the V-side scale to the epilogue, so it needs one segment", @@ -1436,6 +1662,12 @@ def _pa_decode_sparse( HEAD_ALIGNED, IDX_BUFFER_LOAD, FP8_MFMA, + GATHER_CACHE, + IDX_CACHE, + ASYNC_LDS, + RELAXED_LOAD, + ROPE_VEC, + PAD_INTERVAL, ) main_fmt = Fmt( cfg, @@ -1577,12 +1809,15 @@ def _pa_decode_sparse( # An fp8 plane is half the bytes of the bf16 staging it replaces. SMEM_DT: gl.constexpr = gl.float8e4nv if FP8_MFMA else gl.bfloat16 + # The LDS-DMA converts nothing, so an async plane's element type has to be the + # cache's own u8; the dot operands bitcast on read (free, layout-preserving). + PLANE_DT: gl.constexpr = gl.uint8 if ASYNC_LDS else SMEM_DT kv_smem = gl.allocate_shared_memory( - SMEM_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared + PLANE_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared ) if ROPE_SEPARATE: rope_smem = gl.allocate_shared_memory( - SMEM_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared + PLANE_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared ) else: rope_smem = kv_smem # never read as plane 1 in this geometry @@ -1637,22 +1872,40 @@ def _pa_decode_sparse( main_chunk = (main_len + main_splits - 1) // main_splits main_lo = gl.minimum(split_id * main_chunk, main_len) main_hi = gl.minimum(main_lo + main_chunk, main_len) - m_i, l_i, acc = _process_segment( - cfg, - main_seg, - q_dot, - q_rope_dot, - main_lo, - main_hi, - m_i, - l_i, - acc, - head_mask_pv, - main_qk_scale, - main_v_scale, - kv_smem, - rope_smem, - ) + if ASYNC_LDS: + m_i, l_i, acc = _async_segment( + cfg, + main_seg, + q_dot, + q_rope_dot, + main_lo, + main_hi, + m_i, + l_i, + acc, + head_mask_pv, + main_qk_scale, + main_v_scale, + kv_smem, + rope_smem, + ) + else: + m_i, l_i, acc = _process_segment( + cfg, + main_seg, + q_dot, + q_rope_dot, + main_lo, + main_hi, + m_i, + l_i, + acc, + head_mask_pv, + main_qk_scale, + main_v_scale, + kv_smem, + rope_smem, + ) if HAS_EXTRA: extra_seg = Seg( diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 06f32f1ea89..f347f3b174a 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -637,7 +637,7 @@ def _pa_decode_sparse_gfx950_gluon( avg_main = indices.numel() / max(1, num_queries) avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 - # Kernel-side cache-format tags (kernel shared with sparse_mla_decode.py). + # Kernel-side cache-format tags (kernel shared with sparse_mla.py). if UNIFORM: main_fmt = "uniform" if main_is_fp8 else "bf16" extra_fmt = main_fmt diff --git a/aiter/ops/triton/attention/sparse_mla_decode.py b/aiter/ops/triton/attention/sparse_mla.py similarity index 81% rename from aiter/ops/triton/attention/sparse_mla_decode.py rename to aiter/ops/triton/attention/sparse_mla.py index 68561ad386e..a50f79cd895 100644 --- a/aiter/ops/triton/attention/sparse_mla_decode.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Sparse MLA decode (gfx950 gluon): ``kv_lora_rank`` latent + appended +"""Sparse MLA attention (gfx950 gluon): ``kv_lora_rank`` latent + appended decoupled rope, token-granular top-k gather. +Prefill and decode are the same MQA operator on this path, one program per query +token over that token's own gathered KV, so this serves both phases; the launcher +picks the tile from what the grid supplies. + Separated-rope sibling of ``pa_decode_sparse`` — both drive the same gluon kernel, but the calling structures differ enough to warrant separate APIs: DSv4 (rope inside the row, K == V, optional SWA + top-k two-loop, attention @@ -113,17 +117,7 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int: """Split-K count for the sparse-MLA decode. - Separate from the DSv4 policy in _decode_num_splits_occ, which is tuned for a - two-segment SWA + top-k launch and over-splits here: at C=1 topk=2048 it asks - for one split per tile, so 32 programs each pay the full prologue and the - reduce walks 32 partials, costing 1.23-1.28x. - - Below one workgroup per CU, split to fill the machine but never past 8. Past - that, splits stop buying parallelism and start paying prologue and reduce - traffic; 8 is optimal or within 10% everywhere measured. At or above one - workgroup per CU the launch is already full, so a split only buys the second - occupancy slot and needs ~16 tiles to earn its extra partial round trip. At - C=256 that decides one split versus two, and topk=512 wants none. + Below one workgroup per CU, split to fill the machine but never past 8. """ num_sms = get_num_sms() base_wg = max(1, num_queries * heads_blocks) @@ -134,6 +128,51 @@ def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int return max(1, min(cta_cap, tiles, 8)) +# Direct-to-LDS tile pipeline; wins wherever the launch gives a CU more than one +# workgroup to hide the copy behind. +_ASYNC_LDS_DEFAULT = True +_ASYNC_ROPE_VEC = 16 # bytes/lane in the rope copy + + +def _async_launch_config( + fp8_dots: bool, + has_invalid: bool, + num_queries: int, + heads_blocks: int, + num_splits: int, + avg_topk: float, + use_buffer_load: bool, + uni_tile: bool = True, + has_extra: bool = False, +) -> tuple[bool, int, int]: + """-> (ASYNC_LDS, BLOCK_K, waves_per_eu) for this launch. + + BLOCK_K follows what the grid supplies, not the token count: given enough + workgroups to fill four waves/SIMD (prefill) the small tile takes that + occupancy; otherwise (decode, where split-K caps the grid near two + workgroups/CU) the large tile wins instead, by halving the cross-warp softmax + exchange rate per token. Keying on tokens-per-program misclassifies + high-concurrency decode, where the split policy backs off to a couple of splits. + + Async is weakest at exactly one workgroup per CU, having then neither a second + workgroup to hide the copy nor the register path's in-wave prefetch; it stays on + because the surrounding shapes are parity. + """ + enabled = _ASYNC_LDS_DEFAULT + # The kernel static_asserts these; fall back rather than fail to trace. fp8 dots + # because the DMA can only deliver the cache's own code points; the rest because + # the async body has no masked variant and no per-tile validity vector. + enabled = enabled and fp8_dots and uni_tile and not has_invalid and not has_extra + workgroups = num_queries * heads_blocks * max(1, num_splits) + num_sms = get_num_sms() + if enabled and workgroups >= 4 * num_sms: + return True, 64, 4 + waves_per_eu = 2 + if use_buffer_load and workgroups <= num_sms: + waves_per_eu = 1 + return enabled, (128 if enabled else _BLOCK_K), waves_per_eu + + _E4M3_MAX = 448.0 @@ -153,38 +192,21 @@ def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool return False if fmt == "dsmla": raise ValueError( - "dot_precision='fp8' does not support the fp8_ds_mla cache. Its " - "scale varies per 128 elements along the head dim, the QK " - "contraction axis, so it cannot fold outside the tile loop, and " - "scaled MFMA cannot express the PV side either (one scale per 32 " - "contraction elements, where this needs one per token). " + "dot_precision='fp8' does not support the fp8_ds_mla cache." "Use dot_precision='bf16'." ) if fmt == "bf16": raise ValueError( - "dot_precision='fp8' needs an fp8 cache; this one is bf16, and " - "quantizing the pool per call would cost more than the dots save." + "dot_precision='fp8' needs an fp8 cache." ) if fp8_fnuz: raise ValueError( - "dot_precision='fp8' is OCP e4m3 only; this cache is fnuz, a " - "different encoding from what the matrix core reads." + "dot_precision='fp8' is OCP e4m3 only." ) return True -def _lds_pad_for(fp8_dots: bool) -> int: - """kv_smem row padding, in elements of the staged type. - - The pad shifts which banks a transposed K read lands on, which is a byte - property, so an fp8 plane needs twice the elements to move the row pitch by - the same bytes. Carrying bf16's value over to fp8 costs 1.08x at C=64 - topk=2048; 16, 32 and 64 all land at 0.86-0.89x. - """ - return _LDS_PAD * 2 if fp8_dots else _LDS_PAD - - -def sparse_mla_decode_fwd( +def sparse_mla_fwd( q: torch.Tensor, kv_buffer: torch.Tensor, kv_indptr: torch.Tensor, @@ -200,10 +222,11 @@ def sparse_mla_decode_fwd( q_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: - """Sparse (top-k gathered) MLA decode, separated-rope geometry. + """Sparse (top-k gathered) MLA attention, separated-rope geometry. Args: - q: ``[C, H, kv_lora_rank + qk_rope_head_dim]`` bf16 decode queries. + q: ``[C, H, kv_lora_rank + qk_rope_head_dim]`` queries, one row per + query token (prefill and decode alike). kv_buffer: the KV pool — ``[nb, block, R]`` / ``[slots, 1, 1, R]`` / ``[slots, R]`` in bf16, the same shapes in fp8 (+ scalar ``kv_scale``), or ``[nb, block, 656]`` uint8 (vLLM ``fp8_ds_mla``). @@ -216,16 +239,13 @@ def sparse_mla_decode_fwd( skip_reduce: with split-K active, return ``(part_acc, part_m, part_l)`` instead of launching the combine. has_invalid: index stream carries -1 sentinels (masked out). Default - False: vLLM's index converter emits none (it maps -1 to slot 0). + False. dot_precision: what the QK and PV matrix-core ops run in. "bf16" (default): the KV tile is dequantized to bf16 on its way into LDS and both dots are bf16. Works with every cache format. "fp8": the cache's own code points go to the fp8 matrix core with no dequant, and the per-tensor scale folds outside the tile loop. - 0.75-0.96x of "bf16", but q and p both round-trip through e4m3, - so accuracy lands at roughly the asm kernel's level. Needs the - flat per-tensor OCP fp8 cache; see _resolve_dot_precision. q is adapted to the choice. bf16 q is quantized here when "fp8" is asked for, fp8 q is passed straight through, and fp8 q under "bf16" @@ -242,7 +262,7 @@ def sparse_mla_decode_fwd( ``[C, H, kv_lora_rank]`` bf16 attention output (the latent V). """ assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" - assert arch_info.get_arch() == "gfx950", "sparse_mla_decode_fwd is gfx950-only" + assert arch_info.get_arch() == "gfx950", "sparse_mla_fwd is gfx950-only" q_is_fp8 = q.dtype == torch.float8_e4m3fn if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): raise ValueError( @@ -264,7 +284,7 @@ def sparse_mla_decode_fwd( f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" ) _LOGGER.info( - f"SPARSE_MLA_DECODE C={num_queries} H={num_heads} d_qk={d_qk} " + f"SPARSE_MLA C={num_queries} H={num_heads} d_qk={d_qk} " f"nnz={kv_indices.shape[0]}" ) @@ -354,9 +374,14 @@ def sparse_mla_decode_fwd( col_reps = kv_lora_rank // (_GATHER_TW1 * 16) chunk_axis = 1 if col_reps >= 4 else 0 nope_chunk = max(1, _BLOCK_K // 4) if chunk_axis == 0 else min(128, kv_lora_rank) - waves_per_eu = 2 - if use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms(): - waves_per_eu = 1 + # UNI_TILE / HAS_EXTRA are literals at the launch below; passed so the policy + # owns every condition the kernel asserts on. + async_lds_on, block_k, waves_per_eu = _async_launch_config( + fp8_dots, has_invalid, num_queries, heads_blocks, num_splits, avg_topk, + use_buffer_load, uni_tile=True, has_extra=False, + ) + # num_warps stays 4: warps tile the dot's N (MFMA N=16), so a larger BLOCK_K just + # gives each warp two N tiles. grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( @@ -403,12 +428,12 @@ def sparse_mla_decode_fwd( HEAD_SIZE=kv_lora_rank, ROPE_SEPARATE=True, BLOCK_M=block_m, - BLOCK_K=_BLOCK_K, + BLOCK_K=block_k, NUM_SPLITS=num_splits, HEAD_ALIGNED=head_aligned, MFMA_K=_MFMA_K, GATHER_TW1=_GATHER_TW1, - LDS_PAD=_lds_pad_for(fp8_dots), + LDS_PAD=(_LDS_PAD * 2 if fp8_dots else _LDS_PAD), NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", @@ -423,6 +448,12 @@ def sparse_mla_decode_fwd( HAS_INVALID=has_invalid, FP8_FNUZ=fp8_fnuz, FP8_MFMA=fp8_dots, + ASYNC_LDS=async_lds_on, + ROPE_VEC=_ASYNC_ROPE_VEC, + # A 576 B row is not 128 B aligned: it spans five lines and shares its end + # line with its own rope read. ".cg" (sc0 nt) marks that evict-first and + # throws the sharing away, so this gather wants a plain cached load. + GATHER_CACHE="", q_scl_ptr=q_scale, Q_FP8=q_is_fp8, num_warps=num_warps, diff --git a/op_tests/triton_tests/attention/test_sparse_mla_decode.py b/op_tests/triton_tests/attention/test_sparse_mla.py similarity index 75% rename from op_tests/triton_tests/attention/test_sparse_mla_decode.py rename to op_tests/triton_tests/attention/test_sparse_mla.py index 89d7d554e3a..2736c423c14 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla_decode.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Tests for ``sparse_mla_decode_fwd`` (gfx950 gluon, separated-rope MLA). +"""Tests for ``sparse_mla_fwd`` (gfx950 gluon, separated-rope MLA). Target-oriented: the production formats (bf16 and flat per-tensor fp8) get the -shape matrix; fp8_ds_mla gets a single smoke case. Index streams are uniform or -ragged with no -1 sentinels, matching what vLLM's converter emits. +shape matrix, in both dot precisions and at decode and prefill launch shapes; +fp8_ds_mla gets a single smoke case. Index streams are uniform or ragged with no +-1 sentinels, matching what vLLM's converter emits. """ import pytest @@ -14,8 +15,8 @@ from aiter.ops.triton.utils._triton import arch_info if arch_info.get_arch() == "gfx950": - import aiter.ops.triton.attention.sparse_mla_decode as smd - from aiter.ops.triton.attention.sparse_mla_decode import sparse_mla_decode_fwd + import aiter.ops.triton.attention.sparse_mla as smd + from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd FP8_MAX = 448.0 KV_LORA, ROPE = 512, 64 @@ -24,7 +25,7 @@ def _skip_unless_gfx950(): if arch_info.get_arch() != "gfx950": - pytest.skip("sparse_mla_decode_fwd is gfx950-only") + pytest.skip("sparse_mla_fwd is gfx950-only") def quantize_flat_fp8(kv): @@ -113,19 +114,31 @@ def _run_and_check(fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, **kwargs): sm = D_QK**-0.5 q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + out = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) e = rel_err(out, ref) assert e < tol, f"{fmt} C={C} H={H} topk={topk} ragged={ragged}: rel-err {e:.3e}" -@pytest.mark.parametrize("fmt", ["bf16", "tensor"]) +@pytest.mark.parametrize( + "fmt,dots,tol", + [("bf16", "bf16", 2e-2), ("tensor", "bf16", 2e-2), ("tensor", "fp8", 7e-2)], + ids=["bf16", "tensor", "tensor_fp8"], +) @pytest.mark.parametrize("H", [8, 16]) @pytest.mark.parametrize( - "topk,ragged", [(2048, False), (500, True)], ids=["topk2048", "ragged500"] + "C,topk,ragged,pool", + [(8, 2048, False, 1 << 16), (8, 500, True, 1 << 16), (2048, 256, True, 1 << 13)], + ids=["topk2048", "ragged500", "prefill"], ) -def test_sparse_mla_decode(fmt, H, topk, ragged): +def test_sparse_mla(fmt, dots, tol, H, C, topk, ragged, pool): + """The prefill row is the same operator vLLM runs for prefill: one program per + query token over its own ragged prefix, so the launch is sized by the grid + rather than by split-K, and the launcher picks a different tile for it. fp8 + dots get a looser tolerance because q and p both round-trip through e4m3. + """ _skip_unless_gfx950() - _run_and_check(fmt, C=8, H=H, topk=topk, ragged=ragged) + _run_and_check(fmt, C=C, H=H, topk=topk, ragged=ragged, pool=pool, tol=tol, + dot_precision=dots) def test_ds_mla_format(): @@ -134,10 +147,12 @@ def test_ds_mla_format(): @pytest.mark.parametrize("kv_splits", [1, None], ids=["splits1", "splitsauto"]) -def test_split_k(kv_splits): +@pytest.mark.parametrize("dots", ["bf16", "fp8"]) +def test_split_k(kv_splits, dots): _skip_unless_gfx950() - _run_and_check("tensor", C=4, H=16, topk=2048, ragged=False, - pool=1 << 15, kv_splits=kv_splits) + _run_and_check("tensor", C=4, H=16, topk=2048, ragged=False, pool=1 << 15, + tol=7e-2 if dots == "fp8" else 2e-2, + kv_splits=kv_splits, dot_precision=dots) def test_per_tensor_scale_folding_bitwise(): @@ -152,8 +167,8 @@ def test_per_tensor_scale_folding_bitwise(): ks = torch.ones(1, dtype=torch.float32, device="cuda") q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 idx, ptr = make_indices(C, pool, topk, "cuda", 11, ragged=True) - o_t = sparse_mla_decode_fwd(q, exact.view(torch.uint8), ptr, idx, sm, kv_scale=ks) - o_b = sparse_mla_decode_fwd(q, exact.to(torch.bfloat16), ptr, idx, sm) + o_t = sparse_mla_fwd(q, exact.view(torch.uint8), ptr, idx, sm, kv_scale=ks) + o_b = sparse_mla_fwd(q, exact.to(torch.bfloat16), ptr, idx, sm) assert torch.equal(o_t.view(torch.int16), o_b.view(torch.int16)) @@ -164,11 +179,11 @@ def test_vllm_cache_shapes(): C, pool, topk = 4, 1 << 14, 512 sm = D_QK**-0.5 q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, ragged=False) - o_flat = sparse_mla_decode_fwd(q, cache, ptr, idx, sm, kv_scale=ks) - o_4d = sparse_mla_decode_fwd( + o_flat = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks) + o_4d = sparse_mla_fwd( q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks ) - o_paged = sparse_mla_decode_fwd( + o_paged = sparse_mla_fwd( q, cache.view(pool // 64, 64, D_QK), ptr, idx, sm, kv_scale=ks ) assert torch.equal(o_flat, o_4d) and torch.equal(o_flat, o_paged) From 54b1eda7bb8b3ccaaae41394468ac3b49ac3c85d Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Fri, 21 Aug 2026 18:45:51 +0000 Subject: [PATCH 35/51] add bench, simplify comments --- .../ops/triton/attention/pa_decode_sparse.py | 68 +---- .../op_benchmarks/triton/bench_sparse_mla.py | 287 ++++++++++++++++++ 2 files changed, 292 insertions(+), 63 deletions(-) create mode 100644 op_tests/op_benchmarks/triton/bench_sparse_mla.py diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index f347f3b174a..ca62beca1fe 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -469,15 +469,7 @@ def _decode_num_splits( best_splits, best_cost = splits, cost # Collapse to the FEWEST splits that keeps both the wave count and the - # per-split BLOCK_K iteration count. The cost model above treats extra - # splits as nearly free while ``waves`` stays 1 (GAMMA is small and it has no - # reduce term), so at small per-query work it over-splits: at C=64/H=16 with - # main=128, extra=8 it picked 3, where 2 does the same 2 iterations in the - # same single wave but launches 1/3 fewer CTAs and reduces over 2 partials - # instead of 3 (measured 16.2 -> 14.4 us on gfx950/MI355X). - # - # Ported from vLLM's _decode_gfx950_num_splits (PR #52212), which fixed the - # same over-splitting in the in-tree triton decode. + # per-split BLOCK_K iteration count. def _iters(s): m = math.ceil(math.ceil(avg_main / s) / block_k) if avg_main > 0 else 0 e = math.ceil(math.ceil(avg_extra / s) / block_k) if avg_extra > 0 else 0 @@ -497,25 +489,6 @@ def _waves(s): def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block_k): """Split-K count for the gfx950 gluon kernel: fill the machine, but never split a segment finer than one BLOCK_K tile. - - Two facts set this. (a) The kernel is 256 unified VGPRs and 68,608 B of LDS, - so a CU holds two workgroups -- the launch wants ~2*CU programs before extra - splits stop paying. (b) Splitting a segment past its tile count does not divide - the work, it *multiplies* it: every split then owns a partial tile, and a - partial tile costs a full masked gather for a fraction of the tokens. So the - split count is capped by the larger segment's tile count, and the main segment - separately stops at its own (see MAIN_SPLITS in the kernel). - - Measured at C=64 H=16 BLOCK_K=64 (attn = kernel + reduce, us), this rule picks - the sweep optimum at extra in {8, 64, 272, 512, 1024, 2048} and lands within - 1.4% at extra=128: - - extra 8 64 128 272 512 1024 2048 - picked S 2 2 2 5 8 8 8 - best S 2 2 4 5 8 8 8 - old model 2 2 2 3 4 4 4 - old us 14.0 14.6 17.9 21.3 22.0 29.8 45.6 - new us 14.0 14.4 17.4 19.5 20.9 26.9 38.1 """ num_sms = get_num_sms() base_wg = max(1, num_queries * heads_blocks) @@ -524,17 +497,8 @@ def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block extra_tiles = max(1, math.ceil(avg_extra / block_k)) if avg_extra > 0 else 0 tiles = max(1, main_tiles, extra_tiles) if base_wg >= num_sms: - # Already at least one workgroup per CU without splitting, so a split buys - # only the second occupancy slot -- worth an extra round trip of the - # [queries, heads, D] f32 partials plus a reduce launch only when each - # split still owns a real run of tiles. At C=256 extra=8 (2 tiles) forcing - # a split costs 31%; at extra=2048 (32 tiles) it saves 19%. + # Already at least one workgroup per CU without splitting return max(1, min(cta_cap, tiles // 4)) - # Deliberately NOT refined toward split counts that divide `tiles` evenly. - # Minimizing the executed tile count S*ceil(tiles/S) is only right for a - # uniform batch -- it is worth ~5% on the 1-loop seq=1152 shape -- but `tiles` - # here is a batch AVERAGE, and wall clock is set by the longest query. On a - # 16x-ragged batch that refinement picks 5 splits where 8 is 22% faster. return max(1, min(cta_cap, tiles)) @@ -675,14 +639,6 @@ def _pick_splits(hb): num_splits = _pick_splits(heads_blocks) - # Short top-k is latency- and coverage-bound rather than throughput-bound: at - # <= 2 KV tiles per query the whole launch is a few dozen workgroups, so what - # helps is more programs, not fatter ones. Halving BLOCK_M doubles heads_blocks, - # which is real parallelism, unlike a finer split-K that hands every split a - # partial tile. It costs MFMA M, though, and doubling heads_blocks can push the - # launch past one workgroup per CU, which makes the split policy answer with - # fewer splits: take the halving only when it costs neither a split nor the - # one-workgroup-per-CU property. tiles = max( math.ceil(avg_main / BLOCK_K) if avg_main > 0 else 1, math.ceil(avg_extra / BLOCK_K) if avg_extra > 0 else 1, @@ -724,19 +680,11 @@ def _pick_splits(hb): part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # Dequant chunking: split the axis that still has per-lane register repeats, - # rows at gather_tw1=32 and 128-wide column pieces otherwise (128 is the - # narrowest free split -- below that a piece falls inside a thread's 16-byte - # run). Either way a quarter of the tile's f32 expansion is live at a time. + # Dequant chunking col_reps = head_dim // (gather_tw1 * 16) chunk_axis = 1 if col_reps >= 4 else 0 nope_chunk = max(1, BLOCK_K // 4) if chunk_axis == 0 else min(128, head_dim) - # waves_per_eu=2 caps the allocator at 256 unified VGPRs, the 2 waves/SIMD - # threshold on gfx950's 512-VGPR file, so a CU can hold a second workgroup. - # num_warps is 4, one wave per SIMD, so that second slot only exists if the - # launch has more than one workgroup per CU; below that the cap buys no overlap - # and only spills. waves_per_eu = 2 one_wg_per_cu = ( use_buffer_load and num_queries * heads_blocks * num_splits <= get_num_sms() @@ -744,9 +692,6 @@ def _pick_splits(hb): if one_wg_per_cu: waves_per_eu = 1 - # Cap the main segment's split count at whole BLOCK_K tiles: past that every - # split's main range is a partial tile, which costs a full masked gather for - # half the tokens. extra keeps all num_splits programs. main_splits = num_splits if has_extra and avg_main > 0: main_splits = max(1, min(num_splits, math.ceil(avg_main / BLOCK_K))) @@ -758,8 +703,7 @@ def _pick_splits(hb): # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm: # packed OCP fp8 only (uniform scales are f32, fnuz is a different encoding), # bit-identical there because the scale is a power of two. It trades a shorter - # dependent chain for a wider live range, so it pays only where the extra - # registers are free -- the same condition that lifts the occupancy cap. + # dependent chain for a wider live range asm_deq = one_wg_per_cu and not UNIFORM and not fp8_fnuz and main_is_fp8 # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear @@ -842,9 +786,7 @@ def _pick_splits(hb): if skip_reduce: return part_acc, part_m, part_l - # One head per reduce workgroup: the combine is pure bandwidth, so size the grid - # for coverage rather than for the attention kernel's tile. num_queries*num_heads - # workgroups give enough concurrent waves to hide the partial loads. + # One head per reduce workgroup rgrid = (num_queries, num_heads) _pa_decode_sparse_reduce_gfx950[rgrid]( part_m, diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py new file mode 100644 index 00000000000..d05fff5e662 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmark for sparse_mla_fwd (gfx950 gluon, separated-rope MLA). + +Both dot precisions are timed side by side, since the fp8 matrix-core path is +the interesting one and only differs by a flag. + +A shape is num_seqs sequences of num_tokens query tokens each. Prefill and decode +are the same operator here, one program per query token over that token's own +gathered KV, so there is no phase to pick: num_tokens=1 is a decode step and +anything above that is prefill. + +The cache is flushed between iterations by default, which matters: a decode shape +runs up to twice as fast when the loop is allowed to re-read its KV out of cache. + +Usage: + python op_tests/op_benchmarks/triton/bench_sparse_mla.py + python op_tests/op_benchmarks/triton/bench_sparse_mla.py --num_seqs 1 --num_tokens 8192 + python op_tests/op_benchmarks/triton/bench_sparse_mla.py --num_seqs 64 --metric bandwidth +""" + +import argparse + +import torch +import triton +from torch.autograd import DeviceType +from torch.profiler import ProfilerActivity, profile + +from aiter.ops.triton.attention.sparse_mla import _mla_num_splits, sparse_mla_fwd +from aiter.ops.triton.utils._triton import arch_info +from op_tests.op_benchmarks.triton.utils.benchmark_utils import ( + get_caller_name_no_ext, +) + +KV_LORA_RANK = 512 +QK_ROPE_HEAD_DIM = 64 +D_QK = KV_LORA_RANK + QK_ROPE_HEAD_DIM +E4M3_MAX = 448.0 + +# Both kernels the launch runs are named after the gluon kernel, the combine as +# _pa_decode_sparse_reduce, so one substring picks up both and nothing else. +KERNEL_MATCH = "_pa_decode_sparse" +FLUSH_BYTES = 512 << 20 + +_flush_buf = None + + +def flush_l2(): + """Push a buffer bigger than the last-level cache through it. + + A benchmark loop re-reads the same KV every iteration, so without this the + gather is served from cache and the numbers describe the cache rather than + the gather. Any reuse inside a single launch is real and survives this. + """ + global _flush_buf + if _flush_buf is None: + _flush_buf = torch.empty(FLUSH_BYTES, dtype=torch.uint8, device="cuda") + _flush_buf.fill_(1) + + +def device_time_ms(func, warmup=25, rep=100, flush=True): + """Mean GPU time per call, counting only this op's kernels. + + do_bench would measure wall time instead, and a decode launch is small + enough that the wrapper's own python cost exceeds it, which is enough to make + the faster kernel look like the slower one. Filtering by kernel name also + keeps the cache flush out of the total. + """ + for _ in range(warmup): + if flush: + flush_l2() + func() + torch.cuda.synchronize() + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(rep): + if flush: + flush_l2() + func() + torch.cuda.synchronize() + total = sum( + e.self_device_time_total + for e in prof.key_averages() + if e.device_type == DeviceType.CUDA + and KERNEL_MATCH in e.key + and e.self_device_time_total > 0 + ) + if total == 0: + raise RuntimeError(f"no kernel matching {KERNEL_MATCH!r} was profiled") + return total / rep / 1e3 + + +def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes, num_splits): + """Bytes the launch actually moves, counting rows as gathered. + + nnz is the summed per-token top-k length, so a row shared by several query + tokens counts once per token that reads it. That is what the kernel loads, + not the unique footprint. Split-K partials get written and read back by the + combine, so they count twice. + """ + kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it + idx = nnz * 4 # int32 index stream, read once + q = num_tokens * num_heads * D_QK * kv_elem_bytes + out = num_tokens * num_heads * KV_LORA_RANK * 2 + partials = 0 + if num_splits > 1: + # [tokens, splits, heads, lora] bf16 accumulator plus the m/l pair + partials = 2 * num_tokens * num_splits * num_heads * (KV_LORA_RANK * 2 + 8) + return kv + idx + q + out + partials + + +def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): + """Build one launch the way vLLM's sparse-MLA path lays it out. + + Every sequence owns its own slice of the pool, so there is no cross-request + KV sharing to inflate the cache hit rate. Per-token KV length is + min(pos + 1, topk), which is what the index converter emits. + + A single query token per sequence is a decode step, so it sits at the end of + the context. More than one means prefill, and they walk positions 0..n-1. + """ + torch.manual_seed(0) + gen = torch.Generator().manual_seed(1) + pool = num_seqs * context + + kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device=device) * 0.125 + q = ( + torch.randn( + num_seqs * num_tokens, num_heads, D_QK, dtype=torch.bfloat16, + device=device, + ) + * 0.125 + ) + + rows, lens = [], [] + for s in range(num_seqs): + base = s * context + for t in range(num_tokens): + pos = context - 1 if num_tokens == 1 else t + if pos + 1 <= topk: + sel = torch.arange(pos + 1, dtype=torch.int64) # everything fits + else: + sel = torch.randperm(pos + 1, generator=gen)[:topk] + rows.append((sel + base).to(torch.int32)) + lens.append(sel.numel()) + + indices = torch.cat(rows).to(device) + indptr = torch.zeros(len(lens) + 1, dtype=torch.int32) + indptr[1:] = torch.tensor(lens, dtype=torch.int32).cumsum(0) + + # fp8 dots want the flat per-tensor format; bf16 dots read kv as it is. + scale = (kv.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) + kv_fp8 = ( + (kv.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + ).view(torch.uint8) + + return q, kv, kv_fp8, scale.to(device), indices, indptr.to(device) + + +def run_benchmark(args): + if args.metric == "time": + ylabel = "Time (ms)" + elif args.metric == "throughput": + ylabel = "TFLOPs" + else: + ylabel = "GB/s" + + # "phase" is a label, not a knob: it is num_tokens read out loud. It also + # keeps the row from being all-numeric, which is what makes perf_report print + # the shape columns as ints instead of floats. + x_vals_list = [ + ["decode" if tokens == 1 else "prefill", seqs, tokens, args.num_heads, + args.context, args.topk] + for seqs in args.num_seqs + for tokens in args.num_tokens + ] + + benchmark = triton.testing.Benchmark( + x_names=["phase", "num_seqs", "num_tokens", "num_heads", "context", + "topk"], + x_vals=x_vals_list, + line_arg="dots", + line_vals=["bf16", "fp8"], + line_names=["bf16 dots", "fp8 dots"], + styles=[("green", "-"), ("blue", "-")], + ylabel=ylabel, + plot_name=get_caller_name_no_ext(), + args={"metric": args.metric}, + ) + + @triton.testing.perf_report([benchmark]) + def bench_sparse_mla(phase, num_seqs, num_tokens, num_heads, context, topk, + dots, metric, **kwargs): + q, kv, kv_fp8, kv_scale, indices, indptr = build_case( + num_seqs, num_tokens, num_heads, context, topk + ) + sm_scale = D_QK**-0.5 + out = torch.empty( + q.shape[0], num_heads, KV_LORA_RANK, dtype=torch.bfloat16, + device=q.device, + ) + if dots == "fp8": + # Hand q over already quantized, the way production does. Letting + # the wrapper do it would put an amax and a cast inside the timed + # region, which is more than a decode launch costs. + cache, scale = kv_fp8, kv_scale + q_scale = (q.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) + q = ( + (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX) + .to(torch.float8_e4m3fn) + ) + else: + cache, scale, q_scale = kv, None, None + + def func(): + return sparse_mla_fwd( + q, cache, indptr, indices, sm_scale, kv_scale=scale, + q_scale=q_scale, dot_precision=dots, out=out, + ) + + time_ms = device_time_ms(func, flush=not args.no_flush_l2) + + nnz = int(indptr[-1].item()) + num_tokens = q.shape[0] + # QK reads the whole row, PV only the latent half + flops = 2.0 * num_heads * nnz * (D_QK + KV_LORA_RANK) + moved = bytes_moved( + num_tokens, num_heads, nnz, 1 if dots == "fp8" else 2, + _mla_num_splits(num_tokens, 1, nnz / num_tokens), + ) + + if metric == "time": + return time_ms + if metric == "throughput": + return flops / (time_ms * 1e-3) / 1e12 + return moved / (time_ms * 1e-3) / 1e9 + + bench_sparse_mla.run(save_path="." if args.o else None, print_data=True) + + +def main(): + parser = argparse.ArgumentParser( + description="Sparse MLA Benchmark", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--num_seqs", type=int, nargs="+", default=[1, 8, 64, 256], + help="sequences in the launch; every value is benchmarked", + ) + parser.add_argument( + "--num_tokens", type=int, nargs="+", default=[1], + help="query tokens per sequence. 1 is a decode step, more is prefill; " + "every value is benchmarked against every --num_seqs", + ) + parser.add_argument("--num_heads", type=int, default=16, help="q heads (TP4=16)") + parser.add_argument("--context", type=int, default=8192, help="context length") + parser.add_argument("--topk", type=int, default=2048, help="indexer top-k") + parser.add_argument( + "--no_flush_l2", action="store_true", + help="skip the cache flush between iterations, so the gather runs warm", + ) + parser.add_argument( + "-o", action="store_true", help="Write performance results to CSV file" + ) + parser.add_argument( + "--metric", + type=str, + choices=["time", "throughput", "bandwidth"], + default="time", + help="metric to plot. bandwidth is the interesting one for a " + "gather-bound kernel, but an fp8 cache moves half the bytes of a bf16 " + "one, so only time compares the two precisions directly", + ) + args = parser.parse_args() + over = [t for t in args.num_tokens if t > args.context] + if over: + raise SystemExit( + f"--num_tokens {over} exceeds --context {args.context}: a sequence " + "cannot prefill more tokens than its context holds" + ) + if arch_info.get_arch() != "gfx950": + raise SystemExit(f"sparse_mla_fwd is gfx950-only (got {arch_info.get_arch()})") + run_benchmark(args) + + +if __name__ == "__main__": + main() From 7946a1e30dcd708db9442470220c952921ab0a59 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Fri, 21 Aug 2026 19:07:26 +0000 Subject: [PATCH 36/51] cleanup --- .../ops/triton/attention/pa_decode_sparse.py | 52 +------------------ aiter/ops/triton/attention/sparse_mla.py | 39 ++------------ .../op_benchmarks/triton/bench_sparse_mla.py | 37 +------------ 3 files changed, 8 insertions(+), 120 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index ca62beca1fe..3e0056294ba 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -442,49 +442,6 @@ def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor: return x.to(torch.int32).contiguous() -def _decode_num_splits( - num_queries, heads_blocks, avg_main=0.0, avg_extra=0.0, block_k=64 -): - """Pick the split-K count by minimizing a cost model of the decode work: - - cost(s) = waves(s) * iters(s) + GAMMA * s + DELTA * fill(s) - s = # splits - Tuned on gfx950 for DSv4 decode (H=16, D=512, BLOCK_K=64); split count is - capped at 16. - """ - cu = max(1, get_num_sms()) - base = max(1, num_queries * heads_blocks) - GAMMA, DELTA, FILL_CU = 0.32, 2.0, 0.75 - thr = FILL_CU * cu - best_splits, best_cost = 1, None - for splits in range(1, 17): - m_it = math.ceil(math.ceil(avg_main / splits) / block_k) if avg_main > 0 else 0 - e_it = ( - math.ceil(math.ceil(avg_extra / splits) / block_k) if avg_extra > 0 else 0 - ) - waves = (base * splits + cu - 1) // cu - fill = max(0.0, 1.0 - base * splits / thr) / splits - cost = waves * (m_it + e_it) + GAMMA * splits + DELTA * fill - if best_cost is None or cost < best_cost - 1e-9: - best_splits, best_cost = splits, cost - - # Collapse to the FEWEST splits that keeps both the wave count and the - # per-split BLOCK_K iteration count. - def _iters(s): - m = math.ceil(math.ceil(avg_main / s) / block_k) if avg_main > 0 else 0 - e = math.ceil(math.ceil(avg_extra / s) / block_k) if avg_extra > 0 else 0 - return m + e - - def _waves(s): - return (base * s + cu - 1) // cu - - if best_splits > 1: - target_waves, target_iters = _waves(best_splits), _iters(best_splits) - for splits in range(1, best_splits): - if _waves(splits) == target_waves and _iters(splits) == target_iters: - return splits - return best_splits - def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block_k): """Split-K count for the gfx950 gluon kernel: fill the machine, but never @@ -662,9 +619,7 @@ def _pick_splits(hb): ) part_l = torch.empty_like(part_m) # bf16 partials halve both the split-K HBM traffic (~31% of the kernel's - # bytes) and the store count, and the mantissa bits that costs are well - # inside the error the fp8 KV format already carries. skip_reduce hands the - # partials back to the caller, so it keeps their dtype. + # bytes) part_acc = torch.empty( (num_queries, num_splits, num_heads, head_dim), dtype=torch.float32 if skip_reduce else torch.bfloat16, @@ -700,10 +655,7 @@ def _pick_splits(hb): # than one split to give back. adaptive_splits = num_splits > 1 - # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm: - # packed OCP fp8 only (uniform scales are f32, fnuz is a different encoding), - # bit-identical there because the scale is a power of two. It trades a shorter - # dependent chain for a wider live range + # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm asm_deq = one_wg_per_cu and not UNIFORM and not fp8_fnuz and main_is_fp8 # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index a50f79cd895..541f961943a 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -7,17 +7,6 @@ Prefill and decode are the same MQA operator on this path, one program per query token over that token's own gathered KV, so this serves both phases; the launcher picks the tile from what the grid supplies. - -Separated-rope sibling of ``pa_decode_sparse`` — both drive the same gluon -kernel, but the calling structures differ enough to warrant separate APIs: -DSv4 (rope inside the row, K == V, optional SWA + top-k two-loop, attention -sink) stays on ``pa_decode_sparse``; MLA-lineage models (single top-k segment, -no sink, QK over ``kv_lora_rank + qk_rope_head_dim``, V = the latent) use this. - -Follows the ``aiter.mla.mla_decode_fwd`` calling convention as vLLM's -``ROCMAiterMLASparseImpl`` uses it, with two deviations: no head padding -(H < 16 runs natively at BLOCK_M = next_pow2(H)) and q stays bf16 (no q-quant -kernel; <1% accuracy effect). """ import math @@ -41,7 +30,7 @@ _LOGGER = AiterTritonLogger() # buffer_load carries a signed 32-bit offset; caches past this span gather -# through 64-bit addresses (a production MLA pool crosses it at ~3.7M tokens). +# through 64-bit addresses MAX_BYTES = 2**31 - 1 # Tuned launch config (gfx950 / MI355). num_warps = BLOCK_K // 16 (warps tile @@ -151,17 +140,9 @@ def _async_launch_config( workgroups to fill four waves/SIMD (prefill) the small tile takes that occupancy; otherwise (decode, where split-K caps the grid near two workgroups/CU) the large tile wins instead, by halving the cross-warp softmax - exchange rate per token. Keying on tokens-per-program misclassifies - high-concurrency decode, where the split policy backs off to a couple of splits. - - Async is weakest at exactly one workgroup per CU, having then neither a second - workgroup to hide the copy nor the register path's in-wave prefetch; it stays on - because the surrounding shapes are parity. + exchange rate per token. """ enabled = _ASYNC_LDS_DEFAULT - # The kernel static_asserts these; fall back rather than fail to trace. fp8 dots - # because the DMA can only deliver the cache's own code points; the rest because - # the async body has no masked variant and no per-tile validity vector. enabled = enabled and fp8_dots and uni_tile and not has_invalid and not has_extra workgroups = num_queries * heads_blocks * max(1, num_splits) num_sms = get_num_sms() @@ -177,13 +158,6 @@ def _async_launch_config( def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool: - """True for the fp8 matrix-core path. - - Every rejection below is the same constraint: the fp8 path needs the cache - scale to be one positive scalar, because that is what folds outside the tile - loop (K-side into qk_scale, V-side onto the accumulator). A scale that varies - along a dot's contraction axis cannot be folded out at all. - """ if dot_precision not in ("bf16", "fp8"): raise ValueError( f"dot_precision must be 'bf16' or 'fp8', got {dot_precision!r}" @@ -368,14 +342,10 @@ def sparse_mla_fwd( part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # Dequant chunking: split the axis that keeps per-lane register repeats - # (rows at GATHER_TW1=32); waves_per_eu=2 caps the allocator at 256 VGPRs - # for the second occupancy slot, except when the launch cannot fill it. + # Dequant chunking col_reps = kv_lora_rank // (_GATHER_TW1 * 16) chunk_axis = 1 if col_reps >= 4 else 0 nope_chunk = max(1, _BLOCK_K // 4) if chunk_axis == 0 else min(128, kv_lora_rank) - # UNI_TILE / HAS_EXTRA are literals at the launch below; passed so the policy - # owns every condition the kernel asserts on. async_lds_on, block_k, waves_per_eu = _async_launch_config( fp8_dots, has_invalid, num_queries, heads_blocks, num_splits, avg_topk, use_buffer_load, uni_tile=True, has_extra=False, @@ -465,8 +435,7 @@ def sparse_mla_fwd( if skip_reduce: return part_acc, part_m, part_l - # One head per reduce workgroup: the combine is pure bandwidth, so size the - # grid for coverage, not for the attention kernel's tile. + # One head per reduce workgroup rgrid = (num_queries, num_heads) _pa_decode_sparse_reduce_gfx950[rgrid]( part_m, diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py index d05fff5e662..99193cbb632 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_mla.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -3,14 +3,6 @@ """Benchmark for sparse_mla_fwd (gfx950 gluon, separated-rope MLA). -Both dot precisions are timed side by side, since the fp8 matrix-core path is -the interesting one and only differs by a flag. - -A shape is num_seqs sequences of num_tokens query tokens each. Prefill and decode -are the same operator here, one program per query token over that token's own -gathered KV, so there is no phase to pick: num_tokens=1 is a decode step and -anything above that is prefill. - The cache is flushed between iterations by default, which matters: a decode shape runs up to twice as fast when the loop is allowed to re-read its KV out of cache. @@ -47,12 +39,6 @@ def flush_l2(): - """Push a buffer bigger than the last-level cache through it. - - A benchmark loop re-reads the same KV every iteration, so without this the - gather is served from cache and the numbers describe the cache rather than - the gather. Any reuse inside a single launch is real and survives this. - """ global _flush_buf if _flush_buf is None: _flush_buf = torch.empty(FLUSH_BYTES, dtype=torch.uint8, device="cuda") @@ -60,13 +46,6 @@ def flush_l2(): def device_time_ms(func, warmup=25, rep=100, flush=True): - """Mean GPU time per call, counting only this op's kernels. - - do_bench would measure wall time instead, and a decode launch is small - enough that the wrapper's own python cost exceeds it, which is enough to make - the faster kernel look like the slower one. Filtering by kernel name also - keeps the cache flush out of the total. - """ for _ in range(warmup): if flush: flush_l2() @@ -92,11 +71,6 @@ def device_time_ms(func, warmup=25, rep=100, flush=True): def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes, num_splits): """Bytes the launch actually moves, counting rows as gathered. - - nnz is the summed per-token top-k length, so a row shared by several query - tokens counts once per token that reads it. That is what the kernel loads, - not the unique footprint. Split-K partials get written and read back by the - combine, so they count twice. """ kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it idx = nnz * 4 # int32 index stream, read once @@ -165,9 +139,6 @@ def run_benchmark(args): else: ylabel = "GB/s" - # "phase" is a label, not a knob: it is num_tokens read out loud. It also - # keeps the row from being all-numeric, which is what makes perf_report print - # the shape columns as ints instead of floats. x_vals_list = [ ["decode" if tokens == 1 else "prefill", seqs, tokens, args.num_heads, args.context, args.topk] @@ -200,9 +171,7 @@ def bench_sparse_mla(phase, num_seqs, num_tokens, num_heads, context, topk, device=q.device, ) if dots == "fp8": - # Hand q over already quantized, the way production does. Letting - # the wrapper do it would put an amax and a cast inside the timed - # region, which is more than a decode launch costs. + # Hand q over already quantized, the way production does. cache, scale = kv_fp8, kv_scale q_scale = (q.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) q = ( @@ -267,9 +236,7 @@ def main(): type=str, choices=["time", "throughput", "bandwidth"], default="time", - help="metric to plot. bandwidth is the interesting one for a " - "gather-bound kernel, but an fp8 cache moves half the bytes of a bf16 " - "one, so only time compares the two precisions directly", + help="metric to plot", ) args = parser.parse_args() over = [t for t in args.num_tokens if t > args.context] From 4d86ac59df030f7852444b721e7c5d024b870c4f Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Fri, 21 Aug 2026 19:08:48 +0000 Subject: [PATCH 37/51] more cleanup --- aiter/ops/triton/attention/pa_decode_sparse.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 3e0056294ba..d5d3bf0ad25 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -11,9 +11,7 @@ On gfx950 (CDNA4) DeepSeek-V4 sparse-MLA decode has a dedicated gluon implementation (bottom of this module): ``pa_decode_sparse`` routes all formats -to the merged ``_pa_decode_sparse_gfx950_gluon`` driver -- packed fp8_ds_mla / -bf16 block cache (3D; optional SWA+top-k two-loop via ``extra_*``) and the -uniform fp8 / bf16 pool (2D). +to the merged ``_pa_decode_sparse_gfx950_gluon`` driver. """ import math From 2961459d09421ab1240de734f6cd9e42c125c24d Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Sun, 23 Aug 2026 04:45:39 +0000 Subject: [PATCH 38/51] for reviewdog comments --- .../gfx950/attention/pa_decode_sparse.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index b6d03641c16..5221012d34a 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1005,9 +1005,8 @@ def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) elif fmt.KIND == "dsmla": rope_smem.store(k_rope) - elif fmt.KIND == "tensor": - if cfg.ROPE_SEPARATE: - rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) + elif fmt.KIND == "tensor" and cfg.ROPE_SEPARATE: + rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) # "uniform": the whole head is one fp8 tile; nothing else to store. @@ -1464,8 +1463,7 @@ def _process_segment( False, ) - if (not cfg.UNI_TILE) or (not seg.fmt.IS_FP8): - if hi_full < hi: + if ((not cfg.UNI_TILE) or (not seg.fmt.IS_FP8)) and hi_full < hi: m_i, l_i, acc = _decode_tile( cfg, seg, @@ -2041,7 +2039,7 @@ def _pa_decode_sparse_reduce( pid_h = gl.program_id(1) # Lay the 64 lanes out so a small BLOCK_M spends them on the head dim. - TPW0: gl.constexpr = BLOCK_M if BLOCK_M < 8 else 8 + TPW0: gl.constexpr = min(8, BLOCK_M) TPW1: gl.constexpr = 64 // TPW0 BLK: gl.constexpr = gl.BlockedLayout( size_per_thread=[1, 8], From e62372c77acadad91b911dd449ab5adfa8014fff Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Sun, 23 Aug 2026 04:54:46 +0000 Subject: [PATCH 39/51] black and formatting --- .../gfx950/attention/pa_decode_sparse.py | 356 ++++++++++++------ .../ops/triton/attention/pa_decode_sparse.py | 1 - aiter/ops/triton/attention/sparse_mla.py | 66 ++-- .../op_benchmarks/triton/bench_sparse_mla.py | 69 ++-- .../triton_tests/attention/test_sparse_mla.py | 30 +- 5 files changed, 339 insertions(+), 183 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 5221012d34a..7e2f31dd560 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -53,8 +53,15 @@ def _rmax(x, axis): @gluon.jit -def _cache_load(ptr, row, col, USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None, - CACHE: gl.constexpr = ".cg"): +def _cache_load( + ptr, + row, + col, + USE_BUFFER_LOAD: gl.constexpr, + mask=None, + other=None, + CACHE: gl.constexpr = ".cg", +): """Gather rows[i] + col[j]. row is the per-token offset in ptr's element units; col a small compile-time arange. Keeping them apart resolves one pointer per token on the 64-bit path (the column offset folds into the @@ -229,12 +236,12 @@ class Cfg: # geometry BLOCK_M: gl.constexpr BLOCK_K: gl.constexpr - KV_DIM: gl.constexpr # plane-0 width = LDS tile width = V width (pow-2) - ROPE_DIM: gl.constexpr # plane-1 width when ROPE_SEPARATE, else the bf16 - # tail width inside plane 0 (DSv4 packed) - ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0. True: K-only second - # plane; QK contracts over KV_DIM + ROPE_DIM. - QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) + KV_DIM: gl.constexpr # plane-0 width = LDS tile width = V width (pow-2) + ROPE_DIM: gl.constexpr # plane-1 width when ROPE_SEPARATE, else the bf16 + # tail width inside plane 0 (DSv4 packed) + ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0. True: K-only second + # plane; QK contracts over KV_DIM + ROPE_DIM. + QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) MFMA_K: gl.constexpr NUM_WARPS: gl.constexpr GATHER_TW1: gl.constexpr @@ -244,8 +251,8 @@ class Cfg: HAS_INVALID: gl.constexpr HEAD_ALIGNED: gl.constexpr IDX_BUFFER_LOAD: gl.constexpr - FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's - # own fp8 instead of dequantizing to bf16 + FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's + # own fp8 instead of dequantizing to bf16 # Cache policy per load site. The KV gather wants to be cached (its lines are # shared between neighbouring rows and between a row's latent and rope # halves); the index list is read exactly once per program, so it is the @@ -256,8 +263,8 @@ class Cfg: # stages dequantized values, which pass through registers by definition. # Single-buffered; the launcher decides when it is on. ASYNC_LDS: gl.constexpr - RELAXED_LOAD: gl.constexpr # read LDS with the syncedViaAsyncWait hint - ROPE_VEC: gl.constexpr # bytes per lane in the rope plane's copy + RELAXED_LOAD: gl.constexpr # read LDS with the syncedViaAsyncWait hint + ROPE_VEC: gl.constexpr # bytes per lane in the rope plane's copy # operator layouts qk_layout: gl.constexpr pv_layout: gl.constexpr @@ -412,7 +419,7 @@ def __init__( # pins threads_per_warp and makes a warp cover 64*VEC contiguous bytes. # Padding therefore cannot sit inside a row; the interval must be a multiple # of the warp's run so a pad only shifts whole runs. - AVEC = 16 # 128-bit LDS-DMA + AVEC = 16 # 128-bit LDS-DMA ATW = KV_DIM // AVEC self.async_l = gl.constexpr( gl.BlockedLayout( @@ -432,9 +439,7 @@ def __init__( ) ) self.slot_a_l = gl.constexpr(gl.SliceLayout(1, self.async_l.value)) - self.slot_rope_a_l = gl.constexpr( - gl.SliceLayout(1, self.async_rope_l.value) - ) + self.slot_rope_a_l = gl.constexpr(gl.SliceLayout(1, self.async_rope_l.value)) if ASYNC_LDS: self.kv_shared = gl.constexpr( gl.PaddedSharedLayout.with_identity_for( @@ -454,27 +459,27 @@ def __init__( class Fmt: """Compile-time description of one segment's cache format.""" - KIND: gl.constexpr # "bf16" | "dsv4" | "uniform" | "tensor" | "dsmla" - IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop + KIND: gl.constexpr # "bf16" | "dsv4" | "uniform" | "tensor" | "dsmla" + IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop BLOCK_SIZE: gl.constexpr USE_BUFFER_LOAD: gl.constexpr FP8_FNUZ: gl.constexpr ASM_DEQ: gl.constexpr NOPE_CHUNK: gl.constexpr CHUNK_AXIS: gl.constexpr - NOPE_DIM: gl.constexpr # fp8 payload width (448 dsv4; KV_DIM elsewhere) - GROUP: gl.constexpr # scale group width (64 dsv4/uniform, 128 dsmla) - NG: gl.constexpr # scale groups per row (KV_DIM // GROUP) + NOPE_DIM: gl.constexpr # fp8 payload width (448 dsv4; KV_DIM elsewhere) + GROUP: gl.constexpr # scale group width (64 dsv4/uniform, 128 dsmla) + NG: gl.constexpr # scale groups per row (KV_DIM // GROUP) NARROW_SCALE: gl.constexpr scl_l: gl.constexpr # packed-row addressing constants (element units named in the suffix) - TOK_U8: gl.constexpr # bytes per token row inside a block (576 / 656) - TOK_U16: gl.constexpr # the same row in bf16-view units - ROPE_U16_OFF: gl.constexpr # bf16-view offset of the rope tail in a row + TOK_U8: gl.constexpr # bytes per token row inside a block (576 / 656) + TOK_U16: gl.constexpr # the same row in bf16-view units + ROPE_U16_OFF: gl.constexpr # bf16-view offset of the rope tail in a row SCL_TRAILER_U8: gl.constexpr # dsv4: scale bytes per token in the block trailer - TOK_F32: gl.constexpr # dsmla: f32-view stride per token row - SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the group scales - TOK_EL: gl.constexpr # flat formats: cache elements per token row + TOK_F32: gl.constexpr # dsmla: f32-view stride per token row + SCL_F32_OFF: gl.constexpr # dsmla: f32-view offset of the group scales + TOK_EL: gl.constexpr # flat formats: cache elements per token row @gluon.constexpr_function def __init__( @@ -528,9 +533,9 @@ def __init__( # dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] + 8 B UE8M0 per token after # the block; dsmla row: [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] inline. if KIND == "dsv4": - TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 + TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 elif KIND == "dsmla": - TOK_U8 = KV_DIM + 4 * NG + 2 * ROPE_DIM # 512 + 16 + 128 = 656 + TOK_U8 = KV_DIM + 4 * NG + 2 * ROPE_DIM # 512 + 16 + 128 = 656 else: TOK_U8 = 0 self.TOK_U8 = gl.constexpr(TOK_U8) @@ -694,23 +699,30 @@ def _slots( in_range = k_pos < hi if IDX_BUFFER_LOAD: slot = gl.amd.cdna4.buffer_load( - ptr=indices_ptr + seg_start, offsets=k_pos, mask=in_range, other=-1, + ptr=indices_ptr + seg_start, + offsets=k_pos, + mask=in_range, + other=-1, cache=cfg.IDX_CACHE, ) else: - slot = gl.load(indices_ptr + seg_start + k_pos, mask=in_range, other=-1, - cache_modifier=cfg.IDX_CACHE) + slot = gl.load( + indices_ptr + seg_start + k_pos, + mask=in_range, + other=-1, + cache_modifier=cfg.IDX_CACHE, + ) valid = in_range & (slot >= 0) & (slot < num_rows) slot = gl.where(valid, slot, 0) else: # hi >= 1 whenever UNI_TILE runs (guarded by n_full > 0). off = gl.minimum(k_pos, hi - 1) if UNI_TILE else k_pos if IDX_BUFFER_LOAD: - slot = gl.amd.cdna4.buffer_load(ptr=indices_ptr + seg_start, offsets=off, - cache=cfg.IDX_CACHE) + slot = gl.amd.cdna4.buffer_load( + ptr=indices_ptr + seg_start, offsets=off, cache=cfg.IDX_CACHE + ) else: - slot = gl.load(indices_ptr + seg_start + off, - cache_modifier=cfg.IDX_CACHE) + slot = gl.load(indices_ptr + seg_start + off, cache_modifier=cfg.IDX_CACHE) valid = (k_pos < hi) if UNI_TILE else (slot >= 0) if HAS_INVALID: if UNI_TILE: @@ -732,7 +744,7 @@ def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): else: k = kv_smem.permute([1, 0]).load(cfg.k_layout) # [KV_DIM, BLOCK_K] if cfg.ASYNC_LDS: - k = k.to(gl.float8e4nv, bitcast=True) # raw cache bytes; layout-preserving + k = k.to(gl.float8e4nv, bitcast=True) # raw cache bytes; layout-preserving S = gl.amd.cdna4.mfma( q_dot, k, @@ -767,8 +779,7 @@ def _tile_rows(cfg, seg, k_start, seg_hi, k_rng, ROPE: gl.constexpr): @gluon.jit -def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, - CACHE: gl.constexpr): +def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, CACHE: gl.constexpr): """One plane of one tile, global -> LDS, no register staging.""" if USE_BUFFER_LOAD: gl.amd.cdna4.async_copy.buffer_load_to_shared( @@ -787,11 +798,11 @@ def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, @gluon.jit def _copy_tile(cfg, seg, kv_smem, rope_smem, row_l, row_r, offs_l, offs_r): """One commit group = one tile (both planes), so wait_group counts tiles.""" - _copy_plane(seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, - cfg.GATHER_CACHE) + _copy_plane(seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE) if cfg.ROPE_SEPARATE: - _copy_plane(seg, rope_smem, row_r, offs_r, seg.fmt.USE_BUFFER_LOAD, - cfg.GATHER_CACHE) + _copy_plane( + seg, rope_smem, row_r, offs_r, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE + ) gl.amd.cdna4.async_copy.commit_group() @@ -827,15 +838,32 @@ def _async_segment( n_full = (hi - lo + BK - 1) // BK for i in range(n_full): _copy_tile( - cfg, seg, kv_smem, rope_smem, + cfg, + seg, + kv_smem, + rope_smem, _tile_rows(cfg, seg, lo + i * BK, hi, rng_l, False), _tile_rows(cfg, seg, lo + i * BK, hi, rng_r, True), - offs_l, offs_r, + offs_l, + offs_r, ) gl.amd.cdna4.async_copy.wait_group(0) m_i, l_i, acc = _qkpv_lds( - cfg, seg, None, q_dot, q_rope_dot, m_i, l_i, acc, head_mask, - qk_scale, v_scale, kv_smem, rope_smem, lo + i * BK, hi, + cfg, + seg, + None, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + lo + i * BK, + hi, ) return m_i, l_i, acc @@ -880,16 +908,17 @@ def _gather_full( x_u8 = _cache_load( seg.cache_ptr, bg * cs0 + pg * cfg.KV_DIM, offs_full, fmt.USE_BUFFER_LOAD ) - sc = _cache_load( - seg.alt_ptr, bg * NGRP, offs_full // 64, fmt.USE_BUFFER_LOAD - ) + sc = _cache_load(seg.alt_ptr, bg * NGRP, offs_full // 64, fmt.USE_BUFFER_LOAD) k_rope = x_u8 # no rope side-channel -> DCE'd elif fmt.KIND == "tensor": # Per-tensor scale is folded outside the loop (qk_scale / p), so this # is a bare gather; the K-only rope tail follows when separated. x_u8 = _cache_load( - seg.cache_ptr, bg * cs0 + pg * fmt.TOK_EL, offs_full, - fmt.USE_BUFFER_LOAD, CACHE=cfg.GATHER_CACHE, + seg.cache_ptr, + bg * cs0 + pg * fmt.TOK_EL, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) sc = x_u8 # no scale vector -> DCE'd if cfg.ROPE_SEPARATE: @@ -919,8 +948,16 @@ def _gather_full( # data load as well. if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: sc = _scale_load( - seg.scl_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, - fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 0.0, + seg.scl_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 0.0, ) else: sc = _cache_load( @@ -948,8 +985,16 @@ def _gather_full( # Scales first (vmcnt FIFO; see the dsmla branch). if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: sc = _scale_load( - seg.cache_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, cfg.gather_l, - fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 127, + seg.cache_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 127, ) else: sc = _cache_load( @@ -959,12 +1004,8 @@ def _gather_full( # 2-byte elements: <2 x i16> = 4 packed fp8 per VGPR out of one # dword load. Same byte as nope_row (both even), addressed through # the bf16 view; the layout convert is a rename (shared dim-0 tiling). - row16 = gl.convert_layout( - nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) - ) - x_u8 = _cache_load( - seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD - ) + row16 = gl.convert_layout(nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l)) + x_u8 = _cache_load(seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD) else: x_u8 = _cache_load(seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD) bgr, pgr, _ = _slots( @@ -1034,8 +1075,21 @@ def _qkpv( """Stage a prefetched fp8 tile into LDS, then QK -> softmax -> PV.""" _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem) return _qkpv_lds( - cfg, seg, valid, q_dot, q_rope_dot, m_i, l_i, acc, head_mask, - qk_scale, v_scale, kv_smem, rope_smem, k_start, seg_hi, + cfg, + seg, + valid, + q_dot, + q_rope_dot, + m_i, + l_i, + acc, + head_mask, + qk_scale, + v_scale, + kv_smem, + rope_smem, + k_start, + seg_hi, ) @@ -1076,9 +1130,9 @@ def _qkpv_lds( < seg_hi )[None, :] else: - col_mask = gl.convert_layout( - valid, gl.SliceLayout(0, cfg.qk_layout) - )[None, :] + col_mask = gl.convert_layout(valid, gl.SliceLayout(0, cfg.qk_layout))[ + None, : + ] if not cfg.HEAD_ALIGNED: col_mask = ( gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ @@ -1178,8 +1232,12 @@ def _decode_tile( scl_col = offs_full // 64 if MASKED: x_u8 = _cache_load( - seg.cache_ptr, kv_row, offs_full, fmt.USE_BUFFER_LOAD, - mask=valid_g[:, None], other=0, + seg.cache_ptr, + kv_row, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0, ) sc = _cache_load( seg.alt_ptr, @@ -1202,35 +1260,62 @@ def _decode_tile( if MASKED: # scales first: see _gather_full if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: exps = _scale_load( - seg.cache_ptr, scl_row, valid1d, fmt.USE_BUFFER_LOAD, - cfg.gather_l, fmt.scl_l, fmt.NG, cfg.KV_DIM, True, 127, + seg.cache_ptr, + scl_row, + valid1d, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + True, + 127, ) else: exps = _cache_load( - seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD, - mask=valid_g[:, None], other=127, + seg.cache_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=127, ) if fmt.ASM_DEQ: row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) x_u8 = _cache_load( - seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD, - mask=gl.convert_layout( - valid1d, gl.SliceLayout(1, cfg.gather16_l) - )[:, None], + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + mask=gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather16_l))[ + :, None + ], other=0.0, ) else: x_u8 = _cache_load( - seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD, - mask=valid_g[:, None], other=0, + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0, ) else: if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: exps = _scale_load( - seg.cache_ptr, scl_row, scl_row, fmt.USE_BUFFER_LOAD, - cfg.gather_l, fmt.scl_l, fmt.NG, cfg.KV_DIM, False, 127, + seg.cache_ptr, + scl_row, + scl_row, + fmt.USE_BUFFER_LOAD, + cfg.gather_l, + fmt.scl_l, + fmt.NG, + cfg.KV_DIM, + False, + 127, ) else: exps = _cache_load(seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) @@ -1238,9 +1323,7 @@ def _decode_tile( row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) - x_u8 = _cache_load( - seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD - ) + x_u8 = _cache_load(seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD) else: x_u8 = _cache_load( seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD @@ -1271,8 +1354,12 @@ def _decode_tile( kv_row2 = block_idx_g * cs0 + pos_g * fmt.TOK_EL if MASKED: kv = _cache_load( - seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD, - mask=valid_g[:, None], other=0.0, + seg.alt_ptr, + kv_row2, + offs_full, + fmt.USE_BUFFER_LOAD, + mask=valid_g[:, None], + other=0.0, ) else: kv = _cache_load(seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD) @@ -1289,8 +1376,12 @@ def _decode_tile( rope_row = block_idx_gr * cs0 + pos_gr * fmt.TOK_EL + cfg.KV_DIM if MASKED: k_rope = _cache_load( - seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD, - mask=valid_gr[:, None], other=0.0, + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + mask=valid_gr[:, None], + other=0.0, ) else: k_rope = _cache_load( @@ -1303,9 +1394,9 @@ def _decode_tile( NEED_MASK: gl.constexpr = COL_VALID or (not cfg.HEAD_ALIGNED) if NEED_MASK: if COL_VALID: - col_mask = gl.convert_layout( - valid1d, gl.SliceLayout(0, cfg.qk_layout) - )[None, :] + col_mask = gl.convert_layout(valid1d, gl.SliceLayout(0, cfg.qk_layout))[ + None, : + ] if not cfg.HEAD_ALIGNED: col_mask = ( gl.convert_layout(head_mask, gl.SliceLayout(1, cfg.qk_layout))[ @@ -1499,27 +1590,27 @@ def _process_segment( def _pa_decode_sparse( # Shapes below: C = queries, H = num_heads, S = HEAD_SIZE (the V width), # R = ROPE_DIM, nnz = total gathered tokens in a segment's index list. - q_ptr, # [C, H, S (+R when ROPE_SEPARATE)] bf16 + q_ptr, # [C, H, S (+R when ROPE_SEPARATE)] bf16 # One segment = a cache plus its index list. The cache is paged # [num_blocks, BLOCK_SIZE, row] (BLOCK_SIZE = 1 for a flat pool), and # indices[indptr[t]:indptr[t + 1]] are the rows query t attends to. The two # cache pointers are the same allocation under different element types; # which of them is live depends on the format (see Seg). - main_cache_ptr, # main (SWA) cache, u8 view + main_cache_ptr, # main (SWA) cache, u8 view main_cache_bf16_ptr, # bf16 view of it, or the f32 scale pool ("uniform") - main_indices_ptr, # [nnz_main] int32 row ids - main_indptr_ptr, # [C + 1] int32 - extra_cache_ptr, # top-k segment; aliases main when HAS_EXTRA=False + main_indices_ptr, # [nnz_main] int32 row ids + main_indptr_ptr, # [C + 1] int32 + extra_cache_ptr, # top-k segment; aliases main when HAS_EXTRA=False extra_cache_bf16_ptr, - extra_indices_ptr, # [nnz_extra] int32 - extra_indptr_ptr, # [C + 1] int32 - attn_sink_ptr, # [H] f32, HAS_SINK only - out_ptr, # [C, H, S] bf16, written when NUM_SPLITS == 1 + extra_indices_ptr, # [nnz_extra] int32 + extra_indptr_ptr, # [C + 1] int32 + attn_sink_ptr, # [H] f32, HAS_SINK only + out_ptr, # [C, H, S] bf16, written when NUM_SPLITS == 1 # Split-K partials, written instead of out_ptr when NUM_SPLITS > 1 (unused # placeholders otherwise). - part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain - part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum - part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized + part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain + part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum + part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view # ("dsmla"). None elides the argument, keeping other formats' kernarg # layouts unchanged. @@ -1613,21 +1704,24 @@ def _pa_decode_sparse( MAIN_FMT != "dsmla" or ROPE_SEPARATE, "fp8_ds_mla is a separated-rope (MLA) format", ) - gl.static_assert((not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", - "ASM_DEQ is the dsv4 E8M0 dequant") + gl.static_assert( + (not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", + "ASM_DEQ is the dsv4 E8M0 dequant", + ) # The fp8 path needs one positive scalar scale per cache, since that is what # folds outside the loop, and OCP e4m3 code points, which is what the matrix # core reads. gl.static_assert( - (not FP8_MFMA) or (MAIN_FMT == "tensor" and (not HAS_EXTRA or EXTRA_FMT == "tensor")), + (not FP8_MFMA) + or (MAIN_FMT == "tensor" and (not HAS_EXTRA or EXTRA_FMT == "tensor")), "FP8_MFMA requires the per-tensor fp8 format on every segment", ) gl.static_assert(not (FP8_MFMA and FP8_FNUZ), "FP8_MFMA is OCP e4m3 only") # The direct-to-LDS path stages raw code points, needs the clamped (branch-free) # tail, and carries no per-tile validity vector. gl.static_assert( - (not ASYNC_LDS) or (FP8_MFMA and UNI_TILE and not HAS_INVALID - and not HAS_EXTRA), + (not ASYNC_LDS) + or (FP8_MFMA and UNI_TILE and not HAS_INVALID and not HAS_EXTRA), "ASYNC_LDS requires FP8_MFMA + UNI_TILE, no -1 sentinels, one segment", ) gl.static_assert( @@ -1810,9 +1904,7 @@ def _pa_decode_sparse( # The LDS-DMA converts nothing, so an async plane's element type has to be the # cache's own u8; the dot operands bitcast on read (free, layout-preserving). PLANE_DT: gl.constexpr = gl.uint8 if ASYNC_LDS else SMEM_DT - kv_smem = gl.allocate_shared_memory( - PLANE_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared - ) + kv_smem = gl.allocate_shared_memory(PLANE_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared) if ROPE_SEPARATE: rope_smem = gl.allocate_shared_memory( PLANE_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared @@ -1846,7 +1938,9 @@ def _pa_decode_sparse( mask=head_mask_pv, ) gl.amd.cdna4.buffer_store( - gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.pv_layout)), + gl.zeros( + [BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.pv_layout) + ), ptr=part_l_ptr + pm_base, offsets=h_pv.to(gl.int32), mask=head_mask_pv, @@ -1949,9 +2043,15 @@ def _pa_decode_sparse( if NUM_SPLITS == 1: if HAS_SINK: # m_pv is in the base-2 exponent domain; lift the sink into it. - sink = gl.amd.cdna4.buffer_load( - ptr=attn_sink_ptr, offsets=h_pv, mask=head_mask_pv, other=float("-inf") - ).to(gl.float32) * RCP_LN2 + sink = ( + gl.amd.cdna4.buffer_load( + ptr=attn_sink_ptr, + offsets=h_pv, + mask=head_mask_pv, + other=float("-inf"), + ).to(gl.float32) + * RCP_LN2 + ) m_final = _max2(m_pv, sink) alpha = gl.exp2(m_pv - m_final) l_final = l_pv * alpha + gl.exp2(sink - m_final) @@ -2064,12 +2164,19 @@ def _pa_decode_sparse_reduce( for s in range(NUM_SPLITS): base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( - ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, + ptr=part_m_ptr + base, + offsets=h, + mask=head_mask, + other=neg_inf, ) m_final = _max2(m_final, m_s) # m_s already in base-2 exponent domain if HAS_SINK: sink = gl.amd.cdna4.buffer_load( - ptr=attn_sink_ptr, offsets=h, mask=head_mask, other=neg_inf, cache=".cg", + ptr=attn_sink_ptr, + offsets=h, + mask=head_mask, + other=neg_inf, + cache=".cg", ).to(gl.float32) scaled_sink = sink * RCP_LN2 m_final = _max2(m_final, scaled_sink) # lift sink to base-2 @@ -2080,11 +2187,17 @@ def _pa_decode_sparse_reduce( for s in range(NUM_SPLITS): base = query_idx * pm_stride0 + s * pm_stride_s m_s = gl.amd.cdna4.buffer_load( - ptr=part_m_ptr + base, offsets=h, mask=head_mask, other=neg_inf, + ptr=part_m_ptr + base, + offsets=h, + mask=head_mask, + other=neg_inf, cache=".cg", ) l_s = gl.amd.cdna4.buffer_load( - ptr=part_l_ptr + base, offsets=h, mask=head_mask, other=0.0, + ptr=part_l_ptr + base, + offsets=h, + mask=head_mask, + other=0.0, cache=".cg", ) w = gl.exp2(m_s - m_final) @@ -2098,7 +2211,10 @@ def _pa_decode_sparse_reduce( else: acc_mask = head_mask[:, None] acc_s = gl.amd.cdna4.buffer_load( - ptr=part_acc_ptr, offsets=a_off, mask=acc_mask, other=0.0, + ptr=part_acc_ptr, + offsets=a_off, + mask=acc_mask, + other=0.0, cache=".cg", ) acc = acc + w[:, None] * acc_s.to(gl.float32) diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index d5d3bf0ad25..016b0468a0d 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -440,7 +440,6 @@ def _as_int32_contiguous_1d(x: torch.Tensor) -> torch.Tensor: return x.to(torch.int32).contiguous() - def _decode_num_splits_occ(num_queries, heads_blocks, avg_main, avg_extra, block_k): """Split-K count for the gfx950 gluon kernel: fill the machine, but never split a segment finer than one BLOCK_K tile. diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 541f961943a..6e0cd33324f 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -47,12 +47,10 @@ def _check_out(out, q, kv_lora_rank): kv_lora_rank is accepted and written in its leading columns.""" C, H = q.shape[0], q.shape[1] if out is None: - return torch.empty( - (C, H, kv_lora_rank), dtype=torch.bfloat16, device=q.device - ) - assert out.shape[0] == C and out.shape[1] == H, ( - f"out shape {tuple(out.shape)} != [{C}, {H}, >= {kv_lora_rank}]" - ) + return torch.empty((C, H, kv_lora_rank), dtype=torch.bfloat16, device=q.device) + assert ( + out.shape[0] == C and out.shape[1] == H + ), f"out shape {tuple(out.shape)} != [{C}, {H}, >= {kv_lora_rank}]" assert out.shape[2] >= kv_lora_rank and out.stride(2) == 1 assert out.dtype == torch.bfloat16 and out.device == q.device return out @@ -75,9 +73,9 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): if kv.dtype == torch.bfloat16: return "bf16", kv, kv, None, 1, False assert kv.dtype in fp8_dtypes, f"unsupported flat cache dtype {kv.dtype}" - assert kv_scale is not None, ( - "flat fp8 cache needs the per-tensor kv_scale (layer._k_scale)" - ) + assert ( + kv_scale is not None + ), "flat fp8 cache needs the per-tensor kv_scale (layer._k_scale)" assert kv_scale.dtype == torch.float32 u8 = kv.view(torch.uint8) fnuz = kv.dtype == torch.float8_e4m3fnuz @@ -85,13 +83,14 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): if ( 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 + and kv.shape[2] + == kv_lora_rank + 4 * (kv_lora_rank // 128) + 2 * qk_rope_head_dim ): # vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) - assert u8.stride(2) == 1 and u8.stride(1) == u8.shape[2], ( - "fp8_ds_mla rows must be contiguous 656-byte records" - ) + assert ( + u8.stride(2) == 1 and u8.stride(1) == u8.shape[2] + ), "fp8_ds_mla rows must be contiguous 656-byte records" return ( "dsmla", u8, @@ -120,7 +119,7 @@ def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int # Direct-to-LDS tile pipeline; wins wherever the launch gives a CU more than one # workgroup to hide the copy behind. _ASYNC_LDS_DEFAULT = True -_ASYNC_ROPE_VEC = 16 # bytes/lane in the rope copy +_ASYNC_ROPE_VEC = 16 # bytes/lane in the rope copy def _async_launch_config( @@ -170,13 +169,9 @@ def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool "Use dot_precision='bf16'." ) if fmt == "bf16": - raise ValueError( - "dot_precision='fp8' needs an fp8 cache." - ) + raise ValueError("dot_precision='fp8' needs an fp8 cache.") if fp8_fnuz: - raise ValueError( - "dot_precision='fp8' is OCP e4m3 only." - ) + raise ValueError("dot_precision='fp8' is OCP e4m3 only.") return True @@ -241,8 +236,11 @@ def sparse_mla_fwd( if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): raise ValueError( f"q must be bf16 or float8_e4m3fn, got {q.dtype}" - + (" (fnuz is a different encoding from what the matrix core reads)" - if "fnuz" in str(q.dtype) else "") + + ( + " (fnuz is a different encoding from what the matrix core reads)" + if "fnuz" in str(q.dtype) + else "" + ) ) if q_is_fp8: # Caller-quantized q, the asm calling convention: one scaled_fp8_quant @@ -254,9 +252,9 @@ def sparse_mla_fwd( raise ValueError(f"q_scale must be a scalar, got {tuple(q_scale.shape)}") q_scale = q_scale.reshape(1).to(torch.float32).contiguous() num_queries, num_heads, d_qk = q.shape - assert d_qk == kv_lora_rank + qk_rope_head_dim, ( - f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" - ) + assert ( + d_qk == kv_lora_rank + qk_rope_head_dim + ), f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" _LOGGER.info( f"SPARSE_MLA C={num_queries} H={num_heads} d_qk={d_qk} " f"nnz={kv_indices.shape[0]}" @@ -277,7 +275,8 @@ def sparse_mla_fwd( .reshape(1) ) q = ( - (q.float() / q_scale).clamp(-_E4M3_MAX, _E4M3_MAX) + (q.float() / q_scale) + .clamp(-_E4M3_MAX, _E4M3_MAX) .to(torch.float8_e4m3fn) .contiguous() ) @@ -347,8 +346,15 @@ def sparse_mla_fwd( chunk_axis = 1 if col_reps >= 4 else 0 nope_chunk = max(1, _BLOCK_K // 4) if chunk_axis == 0 else min(128, kv_lora_rank) async_lds_on, block_k, waves_per_eu = _async_launch_config( - fp8_dots, has_invalid, num_queries, heads_blocks, num_splits, avg_topk, - use_buffer_load, uni_tile=True, has_extra=False, + fp8_dots, + has_invalid, + num_queries, + heads_blocks, + num_splits, + avg_topk, + use_buffer_load, + uni_tile=True, + has_extra=False, ) # num_warps stays 4: warps tile the dot's N (MFMA N=16), so a larger BLOCK_K just # gives each warp two N tiles. @@ -360,7 +366,7 @@ def sparse_mla_fwd( alt, kv_indices, kv_indptr, - cache, # extra_* segment: unread placeholders (HAS_EXTRA=False) + cache, # extra_* segment: unread placeholders (HAS_EXTRA=False) alt, kv_indices, kv_indptr, @@ -369,7 +375,7 @@ def sparse_mla_fwd( part_m, part_l, part_acc, - scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") + scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") scl, float(softmax_scale), q.stride(0), diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py index 99193cbb632..0b84c8d1a6d 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_mla.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -70,10 +70,9 @@ def device_time_ms(func, warmup=25, rep=100, flush=True): def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes, num_splits): - """Bytes the launch actually moves, counting rows as gathered. - """ - kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it - idx = nnz * 4 # int32 index stream, read once + """Bytes the launch actually moves, counting rows as gathered.""" + kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it + idx = nnz * 4 # int32 index stream, read once q = num_tokens * num_heads * D_QK * kv_elem_bytes out = num_tokens * num_heads * KV_LORA_RANK * 2 partials = 0 @@ -100,7 +99,10 @@ def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device=device) * 0.125 q = ( torch.randn( - num_seqs * num_tokens, num_heads, D_QK, dtype=torch.bfloat16, + num_seqs * num_tokens, + num_heads, + D_QK, + dtype=torch.bfloat16, device=device, ) * 0.125 @@ -140,15 +142,20 @@ def run_benchmark(args): ylabel = "GB/s" x_vals_list = [ - ["decode" if tokens == 1 else "prefill", seqs, tokens, args.num_heads, - args.context, args.topk] + [ + "decode" if tokens == 1 else "prefill", + seqs, + tokens, + args.num_heads, + args.context, + args.topk, + ] for seqs in args.num_seqs for tokens in args.num_tokens ] benchmark = triton.testing.Benchmark( - x_names=["phase", "num_seqs", "num_tokens", "num_heads", "context", - "topk"], + x_names=["phase", "num_seqs", "num_tokens", "num_heads", "context", "topk"], x_vals=x_vals_list, line_arg="dots", line_vals=["bf16", "fp8"], @@ -160,31 +167,39 @@ def run_benchmark(args): ) @triton.testing.perf_report([benchmark]) - def bench_sparse_mla(phase, num_seqs, num_tokens, num_heads, context, topk, - dots, metric, **kwargs): + def bench_sparse_mla( + phase, num_seqs, num_tokens, num_heads, context, topk, dots, metric, **kwargs + ): q, kv, kv_fp8, kv_scale, indices, indptr = build_case( num_seqs, num_tokens, num_heads, context, topk ) sm_scale = D_QK**-0.5 out = torch.empty( - q.shape[0], num_heads, KV_LORA_RANK, dtype=torch.bfloat16, + q.shape[0], + num_heads, + KV_LORA_RANK, + dtype=torch.bfloat16, device=q.device, ) if dots == "fp8": # Hand q over already quantized, the way production does. cache, scale = kv_fp8, kv_scale q_scale = (q.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) - q = ( - (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX) - .to(torch.float8_e4m3fn) - ) + q = (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) else: cache, scale, q_scale = kv, None, None def func(): return sparse_mla_fwd( - q, cache, indptr, indices, sm_scale, kv_scale=scale, - q_scale=q_scale, dot_precision=dots, out=out, + q, + cache, + indptr, + indices, + sm_scale, + kv_scale=scale, + q_scale=q_scale, + dot_precision=dots, + out=out, ) time_ms = device_time_ms(func, flush=not args.no_flush_l2) @@ -194,7 +209,10 @@ def func(): # QK reads the whole row, PV only the latent half flops = 2.0 * num_heads * nnz * (D_QK + KV_LORA_RANK) moved = bytes_moved( - num_tokens, num_heads, nnz, 1 if dots == "fp8" else 2, + num_tokens, + num_heads, + nnz, + 1 if dots == "fp8" else 2, _mla_num_splits(num_tokens, 1, nnz / num_tokens), ) @@ -213,11 +231,17 @@ def main(): formatter_class=argparse.ArgumentDefaultsHelpFormatter, ) parser.add_argument( - "--num_seqs", type=int, nargs="+", default=[1, 8, 64, 256], + "--num_seqs", + type=int, + nargs="+", + default=[1, 8, 64, 256], help="sequences in the launch; every value is benchmarked", ) parser.add_argument( - "--num_tokens", type=int, nargs="+", default=[1], + "--num_tokens", + type=int, + nargs="+", + default=[1], help="query tokens per sequence. 1 is a decode step, more is prefill; " "every value is benchmarked against every --num_seqs", ) @@ -225,7 +249,8 @@ def main(): parser.add_argument("--context", type=int, default=8192, help="context length") parser.add_argument("--topk", type=int, default=2048, help="indexer top-k") parser.add_argument( - "--no_flush_l2", action="store_true", + "--no_flush_l2", + action="store_true", help="skip the cache flush between iterations, so the gather runs warm", ) parser.add_argument( diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index 2736c423c14..abbec1e7861 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -137,8 +137,9 @@ def test_sparse_mla(fmt, dots, tol, H, C, topk, ragged, pool): dots get a looser tolerance because q and p both round-trip through e4m3. """ _skip_unless_gfx950() - _run_and_check(fmt, C=C, H=H, topk=topk, ragged=ragged, pool=pool, tol=tol, - dot_precision=dots) + _run_and_check( + fmt, C=C, H=H, topk=topk, ragged=ragged, pool=pool, tol=tol, dot_precision=dots + ) def test_ds_mla_format(): @@ -150,9 +151,17 @@ def test_ds_mla_format(): @pytest.mark.parametrize("dots", ["bf16", "fp8"]) def test_split_k(kv_splits, dots): _skip_unless_gfx950() - _run_and_check("tensor", C=4, H=16, topk=2048, ragged=False, pool=1 << 15, - tol=7e-2 if dots == "fp8" else 2e-2, - kv_splits=kv_splits, dot_precision=dots) + _run_and_check( + "tensor", + C=4, + H=16, + topk=2048, + ragged=False, + pool=1 << 15, + tol=7e-2 if dots == "fp8" else 2e-2, + kv_splits=kv_splits, + dot_precision=dots, + ) def test_per_tensor_scale_folding_bitwise(): @@ -162,8 +171,11 @@ def test_per_tensor_scale_folding_bitwise(): torch.manual_seed(3) C, H, topk, pool = 4, 16, 512, 1 << 15 sm = D_QK**-0.5 - exact = (torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25 - ).float().to(torch.float8_e4m3fn) + exact = ( + (torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25) + .float() + .to(torch.float8_e4m3fn) + ) ks = torch.ones(1, dtype=torch.float32, device="cuda") q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 idx, ptr = make_indices(C, pool, topk, "cuda", 11, ragged=True) @@ -180,9 +192,7 @@ def test_vllm_cache_shapes(): sm = D_QK**-0.5 q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, ragged=False) o_flat = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks) - o_4d = sparse_mla_fwd( - q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks - ) + o_4d = sparse_mla_fwd(q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks) o_paged = sparse_mla_fwd( q, cache.view(pool // 64, 64, D_QK), ptr, idx, sm, kv_scale=ks ) From d921e5ba5cd0beaaf3a956b4a04cb860c0307c55 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Mon, 24 Aug 2026 15:22:48 +0000 Subject: [PATCH 40/51] add lse return and trim tests/comments --- .../gfx950/attention/pa_decode_sparse.py | 31 ++++- aiter/ops/triton/attention/sparse_mla.py | 33 ++++-- .../op_benchmarks/triton/bench_sparse_mla.py | 3 +- .../triton_tests/attention/test_sparse_mla.py | 109 +++++++----------- 4 files changed, 99 insertions(+), 77 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 7e2f31dd560..504d7b7d685 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -60,7 +60,7 @@ def _cache_load( USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None, - CACHE: gl.constexpr = ".cg", + CACHE: gl.constexpr = gl.constexpr(".cg"), ): """Gather rows[i] + col[j]. row is the per-token offset in ptr's element units; col a small compile-time arange. Keeping them apart resolves one @@ -1611,6 +1611,7 @@ def _pa_decode_sparse( part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized + lse_ptr, # [C, H] f32 natural-log LSE, HAS_LSE only # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view # ("dsmla"). None elides the argument, keeping other formats' kernarg # layouts unchanged. @@ -1633,6 +1634,7 @@ def _pa_decode_sparse( num_heads: gl.constexpr, HAS_EXTRA: gl.constexpr, HAS_SINK: gl.constexpr, + HAS_LSE: gl.constexpr, MAIN_FMT: gl.constexpr, EXTRA_FMT: gl.constexpr, MAIN_BLOCK_SIZE: gl.constexpr, @@ -1677,8 +1679,8 @@ def _pa_decode_sparse( q_scl_ptr=None, Q_FP8: gl.constexpr = False, # Defaults keep every existing launch byte-identical; the MLA launcher opts in. - GATHER_CACHE: gl.constexpr = ".cg", - IDX_CACHE: gl.constexpr = "", + GATHER_CACHE: gl.constexpr = gl.constexpr(".cg"), + IDX_CACHE: gl.constexpr = gl.constexpr(""), ASYNC_LDS: gl.constexpr = False, RELAXED_LOAD: gl.constexpr = True, ROPE_VEC: gl.constexpr = 16, @@ -1804,6 +1806,7 @@ def _pa_decode_sparse( # k_scale joins the per-segment fold (max/exp2 commute with a positive # scale) and the V-side scale hits p in the loop. Both are exact. RCP_LN2: gl.constexpr = 1.4426950408889634 + LN2: gl.constexpr = 0.6931471805599453 qk_scale = scale * RCP_LN2 main_qk_scale = qk_scale main_v_scale = 1.0 @@ -2057,6 +2060,7 @@ def _pa_decode_sparse( l_final = l_pv * alpha + gl.exp2(sink - m_final) acc = acc * alpha[:, None] else: + m_final = m_pv l_final = l_pv one_over_l = 1.0 / l_final out = acc * one_over_l[:, None] @@ -2070,6 +2074,15 @@ def _pa_decode_sparse( offsets=o_off, mask=head_mask_pv[:, None], ) + if HAS_LSE: + # m is base-2, so sum_j exp(s_j) = 2^m * l and ln of it is + # (m + log2 l) * ln2. A fully masked row keeps -inf, not NaN. + gl.amd.cdna4.buffer_store( + (m_final + gl.log2(l_final)) * LN2, + ptr=lse_ptr + query_idx * num_heads, + offsets=h_pv.to(gl.int32), + mask=head_mask_pv, + ) else: # Un-normalized partials for the reduce kernel; m stays in the base-2 # exponent domain (the triton reduce's convention too). @@ -2114,6 +2127,7 @@ def _pa_decode_sparse_reduce( part_acc_ptr, attn_sink_ptr, out_ptr, + lse_ptr, out_stride0: gl.constexpr, out_stride1: gl.constexpr, pm_stride0: gl.constexpr, @@ -2123,6 +2137,7 @@ def _pa_decode_sparse_reduce( pa_stride_h: gl.constexpr, num_heads: gl.constexpr, HAS_SINK: gl.constexpr, + HAS_LSE: gl.constexpr, HEAD_SIZE: gl.constexpr, BLOCK_M: gl.constexpr, NUM_SPLITS: gl.constexpr, @@ -2135,6 +2150,7 @@ def _pa_decode_sparse_reduce( BLOCK_M is sized for workgroup count, not for the attention kernel's tile.""" NUM_WARPS: gl.constexpr = gl.num_warps() RCP_LN2: gl.constexpr = 1.4426950408889634 + LN2: gl.constexpr = 0.6931471805599453 query_idx = gl.program_id(0) pid_h = gl.program_id(1) @@ -2222,6 +2238,15 @@ def _pa_decode_sparse_reduce( if HAS_SINK: l_final = l_final + gl.exp2(scaled_sink - m_final) + if HAS_LSE: + # Same base-2 -> natural conversion as the no-split epilogue. + gl.amd.cdna4.buffer_store( + (m_final + gl.log2(l_final)) * LN2, + ptr=lse_ptr + query_idx * num_heads, + offsets=h.to(gl.int32), + mask=head_mask, + ) + # One reciprocal per row instead of a per-element f32 divide. one_over_l = 1.0 / l_final out = acc * one_over_l[:, None] diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 6e0cd33324f..5bd3216d411 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -190,15 +190,16 @@ def sparse_mla_fwd( dot_precision: str = "bf16", q_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, -) -> torch.Tensor: + return_lse: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: """Sparse (top-k gathered) MLA attention, separated-rope geometry. Args: - q: ``[C, H, kv_lora_rank + qk_rope_head_dim]`` queries, one row per + q: [C, H, kv_lora_rank + qk_rope_head_dim] queries, one row per query token (prefill and decode alike). - kv_buffer: the KV pool — ``[nb, block, R]`` / ``[slots, 1, 1, R]`` / - ``[slots, R]`` in bf16, the same shapes in fp8 (+ scalar - ``kv_scale``), or ``[nb, block, 656]`` uint8 (vLLM ``fp8_ds_mla``). + kv_buffer: the KV pool [nb, block, R] / [slots, 1, 1, R] / + [slots, R] in bf16, the same shapes in fp8 (+ scalar + kv_scale), or [nb, block, 656] uint8 (vLLM ``fp8_ds_mla``). kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts. kv_indices: flat int32 GLOBAL slot ids into the pool. softmax_scale: the layer's softmax scale. @@ -226,9 +227,13 @@ def sparse_mla_fwd( scale, matching production's scaled_fp8_quant(q, layer._q_scale); omitted, a per-tensor amax is taken here. out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. + return_lse: also return the natural-log log-sum-exp, [C, H] f32, for + merging partials across context-parallel ranks. A fully masked row + reports -inf. Returns: - ``[C, H, kv_lora_rank]`` bf16 attention output (the latent V). + (out, lse), following aiter.mla.mla_decode_fwd: out is + [C, H, kv_lora_rank] bf16 (the latent V), lse is None unless return_lse. """ assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" assert arch_info.get_arch() == "gfx950", "sparse_mla_fwd is gfx950-only" @@ -313,6 +318,14 @@ def sparse_mla_fwd( head_aligned = num_heads % block_m == 0 heads_blocks = (num_heads + block_m - 1) // block_m out = _check_out(out, q, kv_lora_rank) + if return_lse and skip_reduce: + raise ValueError("return_lse needs the combine, so skip_reduce cannot be set") + # A live pointer either way, like attn_sink below. + lse = ( + torch.empty((num_queries, num_heads), dtype=torch.float32, device=q.device) + if return_lse + else torch.empty(1, dtype=torch.float32, device=q.device) + ) if kv_splits is not None: num_splits = max(1, int(kv_splits)) @@ -375,6 +388,7 @@ def sparse_mla_fwd( part_m, part_l, part_acc, + lse, scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") scl, float(softmax_scale), @@ -394,6 +408,7 @@ def sparse_mla_fwd( num_heads, HAS_EXTRA=False, HAS_SINK=False, + HAS_LSE=return_lse, MAIN_FMT=fmt, EXTRA_FMT=fmt, MAIN_BLOCK_SIZE=block_size, @@ -437,7 +452,7 @@ def sparse_mla_fwd( ) if num_splits == 1: - return out + return out, (lse if return_lse else None) if skip_reduce: return part_acc, part_m, part_l @@ -449,6 +464,7 @@ def sparse_mla_fwd( part_acc, attn_sink, out, + lse, out.stride(0), out.stride(1), pm_stride0, @@ -458,6 +474,7 @@ def sparse_mla_fwd( pa_stride_h, num_heads, HAS_SINK=False, + HAS_LSE=return_lse, HEAD_SIZE=kv_lora_rank, BLOCK_M=1, NUM_SPLITS=num_splits, @@ -465,4 +482,4 @@ def sparse_mla_fwd( ADAPTIVE_SPLITS=num_splits > 1, num_warps=1, ) - return out + return out, (lse if return_lse else None) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py index 0b84c8d1a6d..736af74e19c 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_mla.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -190,7 +190,8 @@ def bench_sparse_mla( cache, scale, q_scale = kv, None, None def func(): - return sparse_mla_fwd( + # out= is supplied, so the returned pair is not needed + sparse_mla_fwd( q, cache, indptr, diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index abbec1e7861..16db5ef15e7 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -1,13 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Tests for ``sparse_mla_fwd`` (gfx950 gluon, separated-rope MLA). - -Target-oriented: the production formats (bf16 and flat per-tensor fp8) get the -shape matrix, in both dot precisions and at decode and prefill launch shapes; -fp8_ds_mla gets a single smoke case. Index streams are uniform or ragged with no --1 sentinels, matching what vLLM's converter emits. -""" +"""Tests for sparse_mla_fwd (gfx950 gluon, separated-rope MLA).""" import pytest import torch @@ -79,7 +73,7 @@ def make_indices(C, pool, topk, device, seed, ragged): def reference(q, kv_truth, indices, indptr, sm_scale): - """f32 attention over the dequantized cache (tests the kernel, not fp8).""" + """f32 attention over the dequantized cache.""" outs = [] for c in range(q.shape[0]): kvs = kv_truth[indices[indptr[c] : indptr[c + 1]].long()].float() @@ -114,7 +108,7 @@ def _run_and_check(fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, **kwargs): sm = D_QK**-0.5 q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + out, _ = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) e = rel_err(out, ref) assert e < tol, f"{fmt} C={C} H={H} topk={topk} ragged={ragged}: rel-err {e:.3e}" @@ -131,11 +125,6 @@ def _run_and_check(fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, **kwargs): ids=["topk2048", "ragged500", "prefill"], ) def test_sparse_mla(fmt, dots, tol, H, C, topk, ragged, pool): - """The prefill row is the same operator vLLM runs for prefill: one program per - query token over its own ragged prefix, so the launch is sized by the grid - rather than by split-K, and the launcher picks a different tile for it. fp8 - dots get a looser tolerance because q and p both round-trip through e4m3. - """ _skip_unless_gfx950() _run_and_check( fmt, C=C, H=H, topk=topk, ragged=ragged, pool=pool, tol=tol, dot_precision=dots @@ -147,61 +136,51 @@ def test_ds_mla_format(): _run_and_check("dsmla", C=8, H=16, topk=2048, ragged=True) -@pytest.mark.parametrize("kv_splits", [1, None], ids=["splits1", "splitsauto"]) -@pytest.mark.parametrize("dots", ["bf16", "fp8"]) -def test_split_k(kv_splits, dots): - _skip_unless_gfx950() - _run_and_check( - "tensor", - C=4, - H=16, - topk=2048, - ragged=False, - pool=1 << 15, - tol=7e-2 if dots == "fp8" else 2e-2, - kv_splits=kv_splits, - dot_precision=dots, - ) +def reference_lse(q, kv_truth, indices, indptr, sm_scale): + """Natural-log LSE per (query, head), matching mla_decode_fwd's convention.""" + rows = [] + for c in range(q.shape[0]): + kvs = kv_truth[indices[indptr[c] : indptr[c + 1]].long()].float() + s = torch.einsum("hd,td->ht", q[c].float(), kvs) * sm_scale + rows.append(torch.logsumexp(s, dim=-1)) + return torch.stack(rows) -def test_per_tensor_scale_folding_bitwise(): - """kv_scale=1.0 on fp8-exact values must equal the bf16 path bit-for-bit - (K-scale folds into qk_scale, V-scale into p; fp8 -> bf16 is exact).""" +@pytest.mark.parametrize("fmt", ["bf16", "tensor"]) +@pytest.mark.parametrize( + "C,topk,ragged,splits", + [(8, 2048, False, None), (8, 500, True, None), (2048, 256, True, 1)], + ids=["split", "split_ragged", "nosplit"], +) +def test_return_lse(fmt, C, topk, ragged, splits): _skip_unless_gfx950() - torch.manual_seed(3) - C, H, topk, pool = 4, 16, 512, 1 << 15 + H, pool = 16, 1 << 16 sm = D_QK**-0.5 - exact = ( - (torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.25) - .float() - .to(torch.float8_e4m3fn) + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) + kwargs = {} if splits is None else {"kv_splits": splits} + out, lse = sparse_mla_fwd( + q, cache, ptr, idx, sm, kv_scale=ks, return_lse=True, **kwargs ) - ks = torch.ones(1, dtype=torch.float32, device="cuda") - q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 - idx, ptr = make_indices(C, pool, topk, "cuda", 11, ragged=True) - o_t = sparse_mla_fwd(q, exact.view(torch.uint8), ptr, idx, sm, kv_scale=ks) - o_b = sparse_mla_fwd(q, exact.to(torch.bfloat16), ptr, idx, sm) - assert torch.equal(o_t.view(torch.int16), o_b.view(torch.int16)) - - -def test_vllm_cache_shapes(): - """[slots,1,1,R] (asm view) and [nb,block,R] (vLLM paged) must match the - flat [slots,R] form exactly.""" + assert lse is not None + assert lse.shape == (C, H) and lse.dtype == torch.float32 + ref = reference_lse(q, truth.to(torch.bfloat16), idx, ptr, sm) + e = (lse - ref).abs().max().item() / ref.abs().max().item() + assert e < 2e-2, f"{fmt} C={C} splits={splits}: lse rel-err {e:.3e}" + # the output must not change just because the LSE was asked for + out_no, lse_no = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + assert lse_no is None + assert torch.equal(out.view(torch.int16), out_no.view(torch.int16)) + + +def test_global_load_path(): + """A pool whose addressable span passes buffer_load's 2 GB offset limit.""" _skip_unless_gfx950() - C, pool, topk = 4, 1 << 14, 512 + live = 1 << 16 sm = D_QK**-0.5 - q, cache, ks, idx, ptr, _ = _build("tensor", C, 16, topk, pool, ragged=False) - o_flat = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks) - o_4d = sparse_mla_fwd(q, cache.view(pool, 1, 1, D_QK), ptr, idx, sm, kv_scale=ks) - o_paged = sparse_mla_fwd( - q, cache.view(pool // 64, 64, D_QK), ptr, idx, sm, kv_scale=ks - ) - assert torch.equal(o_flat, o_4d) and torch.equal(o_flat, o_paged) - - -def test_global_load_path(monkeypatch): - """Force the 64-bit gather path (production pools cross buffer_load's 2 GB - offset limit at ~3.7M tokens).""" - _skip_unless_gfx950() - monkeypatch.setattr(smd, "MAX_BYTES", 0) - _run_and_check("tensor", C=8, H=16, topk=2048, ragged=False) + q, cache, ks, idx, ptr, truth = _build("tensor", 8, 16, 2048, live, ragged=False) + big = torch.empty(3_800_000, D_QK, dtype=cache.dtype, device=cache.device) + big[:live] = cache + assert smd.max_addressable_bytes(big) >= smd.MAX_BYTES + out, _ = sparse_mla_fwd(q, big, ptr, idx, sm, kv_scale=ks) + e = rel_err(out, reference(q, truth.to(torch.bfloat16), idx, ptr, sm)) + assert e < 2e-2, f"global load path: rel-err {e:.3e}" From 773a5d2a7d060f9082d7dc504a1c1db5b89b1883 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Mon, 24 Aug 2026 15:50:03 +0000 Subject: [PATCH 41/51] cache constexpr restructure to fix 3.6 issues --- .../_gluon_kernels/gfx950/attention/pa_decode_sparse.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 504d7b7d685..bdc98732956 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -41,6 +41,9 @@ # row is guarded), so propagate instead. _MAX_PROP_NAN: gl.constexpr = gl.constexpr(PropagateNan.ALL) +_CG: gl.constexpr = gl.constexpr(".cg") +_NO_CACHE: gl.constexpr = gl.constexpr("") + @gluon.jit def _max2(a, b): @@ -60,7 +63,7 @@ def _cache_load( USE_BUFFER_LOAD: gl.constexpr, mask=None, other=None, - CACHE: gl.constexpr = gl.constexpr(".cg"), + CACHE: gl.constexpr = _CG, ): """Gather rows[i] + col[j]. row is the per-token offset in ptr's element units; col a small compile-time arange. Keeping them apart resolves one @@ -1679,8 +1682,8 @@ def _pa_decode_sparse( q_scl_ptr=None, Q_FP8: gl.constexpr = False, # Defaults keep every existing launch byte-identical; the MLA launcher opts in. - GATHER_CACHE: gl.constexpr = gl.constexpr(".cg"), - IDX_CACHE: gl.constexpr = gl.constexpr(""), + GATHER_CACHE: gl.constexpr = _CG, + IDX_CACHE: gl.constexpr = _NO_CACHE, ASYNC_LDS: gl.constexpr = False, RELAXED_LOAD: gl.constexpr = True, ROPE_VEC: gl.constexpr = 16, From a6be2c17fc42b4b9241bdeacd82c76867373cdeb Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Mon, 24 Aug 2026 17:20:36 +0000 Subject: [PATCH 42/51] change LSE args --- .../_gluon_kernels/gfx950/attention/pa_decode_sparse.py | 9 +++++---- aiter/ops/triton/attention/sparse_mla.py | 8 ++++---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index bdc98732956..6201f285cc4 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1614,7 +1614,6 @@ def _pa_decode_sparse( part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized - lse_ptr, # [C, H] f32 natural-log LSE, HAS_LSE only # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view # ("dsmla"). None elides the argument, keeping other formats' kernarg # layouts unchanged. @@ -1637,7 +1636,6 @@ def _pa_decode_sparse( num_heads: gl.constexpr, HAS_EXTRA: gl.constexpr, HAS_SINK: gl.constexpr, - HAS_LSE: gl.constexpr, MAIN_FMT: gl.constexpr, EXTRA_FMT: gl.constexpr, MAIN_BLOCK_SIZE: gl.constexpr, @@ -1681,6 +1679,9 @@ def _pa_decode_sparse( # mla_decode_fwd uses, where vLLM passes layer._q_scale. q_scl_ptr=None, Q_FP8: gl.constexpr = False, + # [C, H] f32 natural-log LSE, written only when HAS_LSE. + lse_ptr=None, + HAS_LSE: gl.constexpr = False, # Defaults keep every existing launch byte-identical; the MLA launcher opts in. GATHER_CACHE: gl.constexpr = _CG, IDX_CACHE: gl.constexpr = _NO_CACHE, @@ -2130,7 +2131,6 @@ def _pa_decode_sparse_reduce( part_acc_ptr, attn_sink_ptr, out_ptr, - lse_ptr, out_stride0: gl.constexpr, out_stride1: gl.constexpr, pm_stride0: gl.constexpr, @@ -2140,12 +2140,13 @@ def _pa_decode_sparse_reduce( pa_stride_h: gl.constexpr, num_heads: gl.constexpr, HAS_SINK: gl.constexpr, - HAS_LSE: gl.constexpr, HEAD_SIZE: gl.constexpr, BLOCK_M: gl.constexpr, NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, ADAPTIVE_SPLITS: gl.constexpr, + lse_ptr=None, + HAS_LSE: gl.constexpr = False, ): """Split-KV combine: merge per-split partials, fold the sink, write the output. Identical for both geometries (partials are HEAD_SIZE = V wide). diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 5bd3216d411..6f2c88402a6 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -388,7 +388,6 @@ def sparse_mla_fwd( part_m, part_l, part_acc, - lse, scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") scl, float(softmax_scale), @@ -408,7 +407,6 @@ def sparse_mla_fwd( num_heads, HAS_EXTRA=False, HAS_SINK=False, - HAS_LSE=return_lse, MAIN_FMT=fmt, EXTRA_FMT=fmt, MAIN_BLOCK_SIZE=block_size, @@ -447,6 +445,8 @@ def sparse_mla_fwd( GATHER_CACHE="", q_scl_ptr=q_scale, Q_FP8=q_is_fp8, + lse_ptr=lse, + HAS_LSE=return_lse, num_warps=num_warps, waves_per_eu=waves_per_eu, ) @@ -464,7 +464,6 @@ def sparse_mla_fwd( part_acc, attn_sink, out, - lse, out.stride(0), out.stride(1), pm_stride0, @@ -474,12 +473,13 @@ def sparse_mla_fwd( pa_stride_h, num_heads, HAS_SINK=False, - HAS_LSE=return_lse, HEAD_SIZE=kv_lora_rank, BLOCK_M=1, NUM_SPLITS=num_splits, HEAD_ALIGNED=True, ADAPTIVE_SPLITS=num_splits > 1, + lse_ptr=lse, + HAS_LSE=return_lse, num_warps=1, ) return out, (lse if return_lse else None) From d4db44c298e5b043ddbedd72ba3eef293388523b Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 27 Aug 2026 05:58:35 +0000 Subject: [PATCH 43/51] add glm 5.3 flash support, comment trimming --- .../gfx950/attention/pa_decode_sparse.py | 260 ++++++++++++------ aiter/ops/triton/attention/sparse_mla.py | 11 +- .../triton_tests/attention/test_sparse_mla.py | 51 +++- 3 files changed, 225 insertions(+), 97 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 6201f285cc4..751fd80ee9b 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1,16 +1,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Gluon (gfx950) sparse decode for DeepSeek-V4 and separated-rope MLA (GLM-5). +"""Gluon (gfx950) sparse decode for DeepSeek-V4 and MLA (GLM-5). One kernel serves two geometries, selected by the ROPE_SEPARATE constexpr: - False (DSv4): rope lives inside the pow-2 row. One LDS plane - kv_smem [BLOCK_K, HEAD_SIZE], one QK MFMA chain, V = the same plane. + False: rope gets no LDS buffer of its own, either because it + lives inside the pow-2 row (DSv4) or because there is none at all + (GLM-5.3-Flash, ROPE_DIM=0). One LDS buffer, kv_smem + [BLOCK_K, HEAD_SIZE], one QK MFMA chain, V = that same buffer. True (MLA): rope is appended (kv_lora_rank latent + rope = QK width; - V = the latent only). A second pow-2 LDS plane rope_smem holds the + V = the latent only). A second pow-2 LDS buffer, rope_smem, holds the K-only rope and the QK contraction chains a second MFMA into the first. - Plane 0 is the entire V in both geometries. + The KV buffer is the entire V in both geometries. Cache formats (Fmt.KIND), per segment: @@ -68,12 +70,7 @@ def _cache_load( """Gather rows[i] + col[j]. row is the per-token offset in ptr's element units; col a small compile-time arange. Keeping them apart resolves one pointer per token on the 64-bit path (the column offset folds into the - load's immediate) instead of a 64-bit add per element. - - CACHE is the gather's cache policy. ".cg" lowers to `sc0 nt` (evict-first), which - is wrong for a KV gather: a 576/584 B row spans five cache lines and shares its end - line with the rope read of the same row, so evict-first turns one line fill into - two. The MLA launcher passes an empty modifier (a plain cached load) instead.""" + load's immediate) instead of a 64-bit add per element.""" if USE_BUFFER_LOAD: return gl.amd.cdna4.buffer_load( ptr=ptr, @@ -239,11 +236,13 @@ class Cfg: # geometry BLOCK_M: gl.constexpr BLOCK_K: gl.constexpr - KV_DIM: gl.constexpr # plane-0 width = LDS tile width = V width (pow-2) - ROPE_DIM: gl.constexpr # plane-1 width when ROPE_SEPARATE, else the bf16 - # tail width inside plane 0 (DSv4 packed) - ROPE_SEPARATE: gl.constexpr # False: rope inside plane 0. True: K-only second - # plane; QK contracts over KV_DIM + ROPE_DIM. + KV_DIM: gl.constexpr # KV buffer width = LDS tile width = V width (pow-2) + ROPE_DIM: gl.constexpr # rope buffer width when ROPE_SEPARATE, else the + # bf16 tail width inside the KV buffer (DSv4 packed) + ROPE_L: gl.constexpr # rope buffer layout width. ROPE_DIM, or a stand-in + # when there is no rope at all. + ROPE_SEPARATE: gl.constexpr # False: rope inside the KV buffer. True: a + # K-only second buffer, QK contracts over KV_DIM + ROPE_DIM. QK_DIM: gl.constexpr # q row width: KV_DIM (+ ROPE_DIM if separate) MFMA_K: gl.constexpr NUM_WARPS: gl.constexpr @@ -256,18 +255,15 @@ class Cfg: IDX_BUFFER_LOAD: gl.constexpr FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's # own fp8 instead of dequantizing to bf16 - # Cache policy per load site. The KV gather wants to be cached (its lines are - # shared between neighbouring rows and between a row's latent and rope - # halves); the index list is read exactly once per program, so it is the - # stream that actually wants ".cg" (sc0 nt). + # Cache policy per load site GATHER_CACHE: gl.constexpr IDX_CACHE: gl.constexpr # Gather straight into LDS (no register staging). FP8_MFMA only: the bf16 path # stages dequantized values, which pass through registers by definition. - # Single-buffered; the launcher decides when it is on. + # Single-buffered, the launcher decides when it is on. ASYNC_LDS: gl.constexpr RELAXED_LOAD: gl.constexpr # read LDS with the syncedViaAsyncWait hint - ROPE_VEC: gl.constexpr # bytes per lane in the rope plane's copy + ROPE_VEC: gl.constexpr # bytes per lane in the rope buffer's copy # operator layouts qk_layout: gl.constexpr pv_layout: gl.constexpr @@ -360,9 +356,7 @@ def __init__( self.p_layout = gl.constexpr(gl.DotOperandLayout(0, self.pv_layout, KW)) self.v_layout = gl.constexpr(gl.DotOperandLayout(1, self.pv_layout, KW)) - # 16 uint8 per thread = 128-bit gather loads. Warps tile dim 0 so the - # per-lane slot vector stays short and the column direction keeps - # register repeats (what the chunked dequant splits for free). + # 16 uint8 per thread = 128-bit gather loads # GATHER_TW1 = threads spent on the head dim: 32 requests a whole # 512 B token row per instruction, at the cost of a longer slot vector. GSPT = 16 @@ -410,18 +404,16 @@ def __init__( [[KV_DIM, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] ) ) - # Plane 1 (K-only rope) when ROPE_SEPARATE; dead constexpr otherwise. + # The rope buffer (K-only) exists when ROPE_SEPARATE, dead otherwise. + # compiler removes the dead code, so no side effect beyond eliminating errors + ROPE_L = ROPE_DIM if ROPE_DIM > 0 else 64 + self.ROPE_L = gl.constexpr(ROPE_L) self.rope_shared = gl.constexpr( gl.PaddedSharedLayout.with_identity_for( - [[ROPE_DIM, LDS_PAD]], [BLOCK_K, ROPE_DIM], [1, 0] + [[ROPE_L, LDS_PAD]], [BLOCK_K, ROPE_L], [1, 0] ) ) - # LDS-DMA takes one local address per warp (lane i writes warp_base + - # i*VEC), so the source layout must map lane i to LDS offset i*VEC: that - # pins threads_per_warp and makes a warp cover 64*VEC contiguous bytes. - # Padding therefore cannot sit inside a row; the interval must be a multiple - # of the warp's run so a pad only shifts whole runs. AVEC = 16 # 128-bit LDS-DMA ATW = KV_DIM // AVEC self.async_l = gl.constexpr( @@ -432,7 +424,7 @@ def __init__( order=[1, 0], ) ) - RTW = ROPE_DIM // ROPE_VEC + RTW = ROPE_L // ROPE_VEC self.async_rope_l = gl.constexpr( gl.BlockedLayout( size_per_thread=[1, ROPE_VEC], @@ -449,10 +441,10 @@ def __init__( [[PAD_INTERVAL, LDS_PAD]], [BLOCK_K, KV_DIM], [1, 0] ) ) - # 64 * ROPE_VEC is this plane's warp run: the smallest legal interval. + # 64 * ROPE_VEC is this buffer's warp run: the smallest legal interval. self.rope_shared = gl.constexpr( gl.PaddedSharedLayout.with_identity_for( - [[64 * ROPE_VEC, LDS_PAD]], [BLOCK_K, ROPE_DIM], [1, 0] + [[64 * ROPE_VEC, LDS_PAD]], [BLOCK_K, ROPE_L], [1, 0] ) ) @@ -593,11 +585,6 @@ def __init__( self.num_rows = num_rows -# --------------------------------------------------------------------------- -# Dequant-and-store into the LDS plane(s) -# --------------------------------------------------------------------------- - - @gluon.jit def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 @@ -737,7 +724,7 @@ def _slots( @gluon.jit def _qk_scores(cfg, q_dot, q_rope_dot, kv_smem, rope_smem): """QK scores for one tile; ROPE_SEPARATE chains a second MFMA over the rope - plane (MFMA accumulates natively, so KV_DIM + ROPE_DIM is two dots).""" + buffer (MFMA accumulates natively, so KV_DIM + ROPE_DIM is two dots).""" if cfg.ASYNC_LDS and cfg.RELAXED_LOAD: # the async_wait already ordered the copy; stops the backend re-inserting # a conservative vmcnt(0) before every LDS read @@ -782,8 +769,10 @@ def _tile_rows(cfg, seg, k_start, seg_hi, k_rng, ROPE: gl.constexpr): @gluon.jit -def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, CACHE: gl.constexpr): - """One plane of one tile, global -> LDS, no register staging.""" +def _copy_lds_buffer( + seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, CACHE: gl.constexpr +): + """One LDS buffer of one tile, global -> LDS, no register staging.""" if USE_BUFFER_LOAD: gl.amd.cdna4.async_copy.buffer_load_to_shared( dst, @@ -800,10 +789,12 @@ def _copy_plane(seg, dst, row, col, USE_BUFFER_LOAD: gl.constexpr, CACHE: gl.con @gluon.jit def _copy_tile(cfg, seg, kv_smem, rope_smem, row_l, row_r, offs_l, offs_r): - """One commit group = one tile (both planes), so wait_group counts tiles.""" - _copy_plane(seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE) + """One commit group = one tile (both buffers), so wait_group counts tiles.""" + _copy_lds_buffer( + seg, kv_smem, row_l, offs_l, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE + ) if cfg.ROPE_SEPARATE: - _copy_plane( + _copy_lds_buffer( seg, rope_smem, row_r, offs_r, seg.fmt.USE_BUFFER_LOAD, cfg.GATHER_CACHE ) gl.amd.cdna4.async_copy.commit_group() @@ -835,7 +826,7 @@ def _async_segment( """ BK: gl.constexpr = cfg.BLOCK_K offs_l = gl.arange(0, cfg.KV_DIM, layout=gl.SliceLayout(0, cfg.async_l)) - offs_r = gl.arange(0, cfg.ROPE_DIM, layout=gl.SliceLayout(0, cfg.async_rope_l)) + offs_r = gl.arange(0, cfg.ROPE_L, layout=gl.SliceLayout(0, cfg.async_rope_l)) rng_l = gl.arange(0, BK, layout=cfg.slot_a_l) rng_r = gl.arange(0, BK, layout=cfg.slot_rope_a_l) n_full = (hi - lo + BK - 1) // BK @@ -909,9 +900,19 @@ def _gather_full( if fmt.KIND == "uniform": NGRP: gl.constexpr = cfg.KV_DIM // 64 x_u8 = _cache_load( - seg.cache_ptr, bg * cs0 + pg * cfg.KV_DIM, offs_full, fmt.USE_BUFFER_LOAD + seg.cache_ptr, + bg * cs0 + pg * cfg.KV_DIM, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + sc = _cache_load( + seg.alt_ptr, + bg * NGRP, + offs_full // 64, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) - sc = _cache_load(seg.alt_ptr, bg * NGRP, offs_full // 64, fmt.USE_BUFFER_LOAD) k_rope = x_u8 # no rope side-channel -> DCE'd elif fmt.KIND == "tensor": # Per-tensor scale is folded outside the loop (qk_scale / p), so this @@ -942,7 +943,7 @@ def _gather_full( CACHE=cfg.GATHER_CACHE, ) else: - k_rope = x_u8 # rope lives inside plane 0 -> DCE'd + k_rope = x_u8 # rope lives inside the KV buffer -> DCE'd elif fmt.KIND == "dsmla": nope_row = bg * cs0 + pg * fmt.TOK_U8 scl_row = bg * (cs0 // 4) + pg * fmt.TOK_F32 + fmt.SCL_F32_OFF @@ -964,9 +965,19 @@ def _gather_full( ) else: sc = _cache_load( - seg.scl_ptr, scl_row, offs_full // fmt.GROUP, fmt.USE_BUFFER_LOAD + seg.scl_ptr, + scl_row, + offs_full // fmt.GROUP, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) - x_u8 = _cache_load(seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD) + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) bgr, pgr, _ = _slots( cfg, seg, @@ -981,6 +992,7 @@ def _gather_full( bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, offs_rope, fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) else: # "dsv4" nope_row = bg * cs0 + pg * fmt.TOK_U8 @@ -1001,16 +1013,32 @@ def _gather_full( ) else: sc = _cache_load( - seg.cache_ptr, scl_row, offs_full // 64, fmt.USE_BUFFER_LOAD + seg.cache_ptr, + scl_row, + offs_full // 64, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) if fmt.ASM_DEQ: # 2-byte elements: <2 x i16> = 4 packed fp8 per VGPR out of one # dword load. Same byte as nope_row (both even), addressed through # the bf16 view; the layout convert is a rename (shared dim-0 tiling). row16 = gl.convert_layout(nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l)) - x_u8 = _cache_load(seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD) + x_u8 = _cache_load( + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) else: - x_u8 = _cache_load(seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD) + x_u8 = _cache_load( + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) bgr, pgr, _ = _slots( cfg, seg, @@ -1025,13 +1053,14 @@ def _gather_full( bgr * (cs0 // 2) + pgr * fmt.TOK_U16 + fmt.ROPE_U16_OFF, offs_rope, fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) return x_u8, sc, k_rope, valid @gluon.jit def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): - """Write one prefetched tile into the LDS plane(s). Plane 0 is always the + """Write one prefetched tile into the LDS buffer(s). The KV buffer is always the full KV_DIM-wide dequant: "dsv4" gathers KV_DIM bytes too, the last 64 being bf16 rope read as garbage fp8 and overwritten by the slice-store below, so the gather stays pow-2 wide.""" @@ -1241,6 +1270,7 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0, + CACHE=cfg.GATHER_CACHE, ) sc = _cache_load( seg.alt_ptr, @@ -1249,10 +1279,23 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, + CACHE=cfg.GATHER_CACHE, ) else: - x_u8 = _cache_load(seg.cache_ptr, kv_row, offs_full, fmt.USE_BUFFER_LOAD) - sc = _cache_load(seg.alt_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) + x_u8 = _cache_load( + seg.cache_ptr, + kv_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) + sc = _cache_load( + seg.alt_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) elif fmt.KIND == "dsv4": nope_row = block_idx_g * cs0 + pos_g * fmt.TOK_U8 @@ -1282,6 +1325,7 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=127, + CACHE=cfg.GATHER_CACHE, ) if fmt.ASM_DEQ: row16 = gl.convert_layout( @@ -1296,6 +1340,7 @@ def _decode_tile( :, None ], other=0.0, + CACHE=cfg.GATHER_CACHE, ) else: x_u8 = _cache_load( @@ -1305,6 +1350,7 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0, + CACHE=cfg.GATHER_CACHE, ) else: if fmt.NARROW_SCALE and not fmt.USE_BUFFER_LOAD: @@ -1321,15 +1367,31 @@ def _decode_tile( 127, ) else: - exps = _cache_load(seg.cache_ptr, scl_row, scl_col, fmt.USE_BUFFER_LOAD) + exps = _cache_load( + seg.cache_ptr, + scl_row, + scl_col, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) if fmt.ASM_DEQ: row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) - x_u8 = _cache_load(seg.alt_ptr, row16, offs_full16, fmt.USE_BUFFER_LOAD) + x_u8 = _cache_load( + seg.alt_ptr, + row16, + offs_full16, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) else: x_u8 = _cache_load( - seg.cache_ptr, nope_row, offs_full, fmt.USE_BUFFER_LOAD + seg.cache_ptr, + nope_row, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) _deq_store_tile(x_u8, exps, kv_smem, cfg, fmt) block_idx_gr, pos_gr, valid_gr = _slots( @@ -1349,9 +1411,16 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_gr[:, None], other=0.0, + CACHE=cfg.GATHER_CACHE, ) else: - k_rope = _cache_load(seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD) + k_rope = _cache_load( + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) else: # "bf16" (tensor/dsmla require UNI_TILE, so they never come here) kv_row2 = block_idx_g * cs0 + pos_g * fmt.TOK_EL @@ -1363,9 +1432,16 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_g[:, None], other=0.0, + CACHE=cfg.GATHER_CACHE, ) else: - kv = _cache_load(seg.alt_ptr, kv_row2, offs_full, fmt.USE_BUFFER_LOAD) + kv = _cache_load( + seg.alt_ptr, + kv_row2, + offs_full, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, + ) kv_smem.store(kv) if cfg.ROPE_SEPARATE: block_idx_gr, pos_gr, valid_gr = _slots( @@ -1385,10 +1461,15 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, mask=valid_gr[:, None], other=0.0, + CACHE=cfg.GATHER_CACHE, ) else: k_rope = _cache_load( - seg.alt_ptr, rope_row, offs_rope, fmt.USE_BUFFER_LOAD + seg.alt_ptr, + rope_row, + offs_rope, + fmt.USE_BUFFER_LOAD, + CACHE=cfg.GATHER_CACHE, ) rope_smem.store(k_rope) @@ -1453,7 +1534,7 @@ def _process_segment( offs_full16 = gl.arange( 0, cfg.KV_DIM // 2, layout=gl.SliceLayout(0, cfg.gather16_l) ) - offs_rope = gl.arange(0, cfg.ROPE_DIM, layout=gl.SliceLayout(0, cfg.gather_rope_l)) + offs_rope = gl.arange(0, cfg.ROPE_L, layout=gl.SliceLayout(0, cfg.gather_rope_l)) k_rng_slot = gl.arange(0, cfg.BLOCK_K, layout=cfg.slot_l) k_rng_rope = gl.arange(0, cfg.BLOCK_K, layout=gl.SliceLayout(1, cfg.gather_rope_l)) @@ -1704,12 +1785,26 @@ def _pa_decode_sparse( ) gl.static_assert( (not ROPE_SEPARATE) or (MAIN_FMT != "dsv4" and MAIN_FMT != "uniform"), - "dsv4/uniform formats carry rope inside plane 0 (ROPE_SEPARATE=False)", + "dsv4/uniform formats carry rope inside the KV buffer (ROPE_SEPARATE=False)", ) gl.static_assert( MAIN_FMT != "dsmla" or ROPE_SEPARATE, "fp8_ds_mla is a separated-rope (MLA) format", ) + # No rope means no rope buffer, and the packed formats all define a rope tail. + gl.static_assert( + ROPE_DIM > 0 or not ROPE_SEPARATE, + "ROPE_DIM=0 is the rope-free geometry; it needs ROPE_SEPARATE=False", + ) + gl.static_assert( + ROPE_DIM > 0 + or ( + MAIN_FMT != "dsv4" + and MAIN_FMT != "dsmla" + and (not HAS_EXTRA or (EXTRA_FMT != "dsv4" and EXTRA_FMT != "dsmla")) + ), + "dsv4/fp8_ds_mla rows carry a rope tail; ROPE_DIM=0 is inconsistent", + ) gl.static_assert( (not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", "ASM_DEQ is the dsv4 E8M0 dequant", @@ -1805,10 +1900,7 @@ def _pa_decode_sparse( extra_len = 0 main_len = main_end - main_start - # exp2 softmax: fold scale*log2(e) into the loop exponent; keep raw scale - # for the sink (a scaled-score-space logit). For "tensor" the cache's - # k_scale joins the per-segment fold (max/exp2 commute with a positive - # scale) and the V-side scale hits p in the loop. Both are exact. + # exp2 softmax RCP_LN2: gl.constexpr = 1.4426950408889634 LN2: gl.constexpr = 0.6931471805599453 qk_scale = scale * RCP_LN2 @@ -1840,8 +1932,6 @@ def _pa_decode_sparse( if FP8_MFMA and not Q_FP8: # bf16 q: quantize here, one e4m3 scale for this program's whole Q tile # (nope and rope), so the fold below is one extra factor on qk_scale. - # m/l stay in the true-score domain, so split-K programs stay comparable - # in the reduce. E4M3_MAX: gl.constexpr = 448.0 q_amax = gl.max(gl.max(gl.abs(q).to(gl.float32), axis=1), axis=0) q_dot = gl.convert_layout(q, cfg.q_layout) @@ -1858,15 +1948,13 @@ def _pa_decode_sparse( ) q_rope_dot = gl.convert_layout(q_rope, cfg.q_layout) else: - q_rope_dot = q_dot # unused (single-plane QK) -> DCE'd + q_rope_dot = q_dot # unused (single-buffer QK) -> DCE'd if Q_FP8: - # Nothing to quantize; fold the caller's scale the way k_scale is - # folded. + # Nothing to quantize q_scale = gl.load(q_scl_ptr) if not FP8_MFMA: - # fp8 -> bf16 is exact (3 mantissa bits into 8), so a quantized q - # costs nothing extra on the bf16 dots. + # fp8 -> bf16 is exact (3 mantissa bits into 8) q_dot = gl.convert_layout(q.to(gl.bfloat16), cfg.q_layout) if ROPE_SEPARATE: q_rope_dot = gl.convert_layout(q_rope.to(gl.bfloat16), cfg.q_layout) @@ -1906,18 +1994,18 @@ def _pa_decode_sparse( l_i = gl.zeros([BLOCK_M], gl.float32, layout=gl.SliceLayout(1, cfg.qk_layout)) acc = gl.zeros([BLOCK_M, HEAD_SIZE], gl.float32, layout=cfg.pv_layout) - # An fp8 plane is half the bytes of the bf16 staging it replaces. + # An fp8 buffer is half the bytes of the bf16 staging it replaces. SMEM_DT: gl.constexpr = gl.float8e4nv if FP8_MFMA else gl.bfloat16 - # The LDS-DMA converts nothing, so an async plane's element type has to be the + # The LDS-DMA converts nothing, so an async buffer's element type has to be the # cache's own u8; the dot operands bitcast on read (free, layout-preserving). - PLANE_DT: gl.constexpr = gl.uint8 if ASYNC_LDS else SMEM_DT - kv_smem = gl.allocate_shared_memory(PLANE_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared) + BUF_DT: gl.constexpr = gl.uint8 if ASYNC_LDS else SMEM_DT + kv_smem = gl.allocate_shared_memory(BUF_DT, [BLOCK_K, HEAD_SIZE], cfg.kv_shared) if ROPE_SEPARATE: rope_smem = gl.allocate_shared_memory( - PLANE_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared + BUF_DT, [BLOCK_K, ROPE_DIM], cfg.rope_shared ) else: - rope_smem = kv_smem # never read as plane 1 in this geometry + rope_smem = kv_smem # never read as the rope buffer in this geometry # ADAPTIVE_SPLITS: the host split count is sized for batch averages; in a # ragged batch the surplus programs would each gather a mostly-masked tile @@ -2177,8 +2265,7 @@ def _pa_decode_sparse_reduce( neg_inf = float("-inf") # Deliberately NOT bounded by the per-query split count: that would make - # these dynamic loops and lose the static unroll. A split that bowed out - # wrote m = -inf, which is what actually has to be handled (below). + # these dynamic loops and lose the static unroll m_final = gl.full([BLOCK_M], neg_inf, gl.float32, layout=row_l) # pass 1: global max over splits for s in range(NUM_SPLITS): @@ -2224,8 +2311,7 @@ def _pa_decode_sparse_reduce( l_final = l_final + w * l_s a_base = query_idx * pa_stride0 + s * pa_stride_s a_off = (a_base + h[:, None] * pa_stride_h + offs_d[None, :]).to(gl.int32) - # A bowed-out split's part_acc is uninitialized, so mask the load - # rather than relying on w == 0 (0 * NaN is NaN). + # split's part_acc can be uninitialized, so mask the load if ADAPTIVE_SPLITS: acc_mask = head_mask[:, None] & (m_s > neg_inf)[:, None] else: diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 6f2c88402a6..c17b489936c 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -192,11 +192,17 @@ def sparse_mla_fwd( out: torch.Tensor | None = None, return_lse: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """Sparse (top-k gathered) MLA attention, separated-rope geometry. + """Sparse (top-k gathered) MLA attention. + + ``qk_rope_head_dim`` selects the geometry: separated rope (DeepSeek-V3.2, + GLM-5.1/5.2), where the query is the latent plus an appended rope, or + rope-free (GLM-5.3-Flash), where the query is the latent alone. Args: q: [C, H, kv_lora_rank + qk_rope_head_dim] queries, one row per query token (prefill and decode alike). + qk_rope_head_dim: width of the appended rope. 0 means rope-free, and + the QK contraction becomes a single dot over kv_lora_rank. kv_buffer: the KV pool [nb, block, R] / [slots, 1, 1, R] / [slots, R] in bf16, the same shapes in fp8 (+ scalar kv_scale), or [nb, block, 656] uint8 (vLLM ``fp8_ds_mla``). @@ -260,6 +266,7 @@ def sparse_mla_fwd( assert ( d_qk == kv_lora_rank + qk_rope_head_dim ), f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" + assert qk_rope_head_dim >= 0, "qk_rope_head_dim cannot be negative" _LOGGER.info( f"SPARSE_MLA C={num_queries} H={num_heads} d_qk={d_qk} " f"nnz={kv_indices.shape[0]}" @@ -415,7 +422,7 @@ def sparse_mla_fwd( NOPE_DIM=kv_lora_rank, ROPE_DIM=qk_rope_head_dim, HEAD_SIZE=kv_lora_rank, - ROPE_SEPARATE=True, + ROPE_SEPARATE=qk_rope_head_dim > 0, BLOCK_M=block_m, BLOCK_K=block_k, NUM_SPLITS=num_splits, diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index 16db5ef15e7..731de77b18b 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Tests for sparse_mla_fwd (gfx950 gluon, separated-rope MLA).""" +"""Tests for sparse_mla_fwd (gfx950 gluon MLA). + +Covers both geometries: separated rope (DSV3.2, GLM-5.1, GLM-5.2) and +rope-free (GLM-5.3-Flash), where the query is the latent alone. +""" import pytest import torch @@ -87,10 +91,11 @@ def rel_err(out, ref): return ((out.float() - ref).abs().max() / ref.abs().max()).item() -def _build(fmt, C, H, topk, pool, ragged, seed=0): +def _build(fmt, C, H, topk, pool, ragged, seed=0, rope=ROPE): torch.manual_seed(seed) - q = torch.randn(C, H, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 - kv = torch.randn(pool, D_QK, dtype=torch.bfloat16, device="cuda") * 0.125 + d_qk = KV_LORA + rope + q = torch.randn(C, H, d_qk, dtype=torch.bfloat16, device="cuda") * 0.125 + kv = torch.randn(pool, d_qk, dtype=torch.bfloat16, device="cuda") * 0.125 idx, ptr = make_indices(C, pool, topk, "cuda", seed + 7, ragged) ks = None if fmt == "bf16": @@ -104,11 +109,15 @@ def _build(fmt, C, H, topk, pool, ragged, seed=0): return q, cache, ks, idx, ptr, truth -def _run_and_check(fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, **kwargs): - sm = D_QK**-0.5 - q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged) +def _run_and_check( + fmt, C, H, topk, ragged, tol=2e-2, pool=1 << 16, rope=ROPE, **kwargs +): + sm = (KV_LORA + rope) ** -0.5 + q, cache, ks, idx, ptr, truth = _build(fmt, C, H, topk, pool, ragged, rope=rope) ref = reference(q, truth.to(torch.bfloat16), idx, ptr, sm) - out, _ = sparse_mla_fwd(q, cache, ptr, idx, sm, kv_scale=ks, **kwargs) + out, _ = sparse_mla_fwd( + q, cache, ptr, idx, sm, kv_scale=ks, qk_rope_head_dim=rope, **kwargs + ) e = rel_err(out, ref) assert e < tol, f"{fmt} C={C} H={H} topk={topk} ragged={ragged}: rel-err {e:.3e}" @@ -184,3 +193,29 @@ def test_global_load_path(): out, _ = sparse_mla_fwd(q, big, ptr, idx, sm, kv_scale=ks) e = rel_err(out, reference(q, truth.to(torch.bfloat16), idx, ptr, sm)) assert e < 2e-2, f"global load path: rel-err {e:.3e}" + + +@pytest.mark.parametrize( + "fmt,dots,tol", + [("bf16", "bf16", 2e-2), ("tensor", "bf16", 2e-2), ("tensor", "fp8", 7e-2)], + ids=["bf16", "tensor", "tensor_fp8"], +) +@pytest.mark.parametrize("H", [8, 16]) +@pytest.mark.parametrize( + "C,topk,ragged,pool", + [(8, 2048, False, 1 << 16), (8, 500, True, 1 << 16), (2048, 256, True, 1 << 13)], + ids=["topk2048", "ragged500", "prefill"], +) +def test_sparse_mla_rope_free(fmt, dots, tol, H, C, topk, ragged, pool): + _skip_unless_gfx950() + _run_and_check( + fmt, + C=C, + H=H, + topk=topk, + ragged=ragged, + pool=pool, + tol=tol, + rope=0, + dot_precision=dots, + ) From 9b1b986abb8983c04d8bd14ade86cfaf66cd4da8 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 27 Aug 2026 06:14:37 +0000 Subject: [PATCH 44/51] cache mod for q --- .../gfx950/attention/pa_decode_sparse.py | 5 +++-- aiter/ops/triton/attention/pa_decode_sparse.py | 4 ++++ aiter/ops/triton/attention/sparse_mla.py | 12 ++++++------ 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 751fd80ee9b..3f9f3c84a52 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -1740,6 +1740,7 @@ def _pa_decode_sparse( PART_STORE_CACHE: gl.constexpr, UNI_TILE: gl.constexpr, GRID_ORDER: gl.constexpr, + Q_CACHE: gl.constexpr, # MAIN_SPLITS <= NUM_SPLITS: splitting the SWA window past its tile count # only manufactures masked partial tiles, so main stops early and extra # keeps all programs (surplus ones get an empty main range). @@ -1927,7 +1928,7 @@ def _pa_decode_sparse( gl.int32 ) q = gl.amd.cdna4.buffer_load( - ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0 + ptr=q_ptr, offsets=q_off, mask=h_mask_q[:, None], other=0.0, cache=Q_CACHE ) if FP8_MFMA and not Q_FP8: # bf16 q: quantize here, one e4m3 scale for this program's whole Q tile @@ -1944,7 +1945,7 @@ def _pa_decode_sparse( + offs_d_qr[None, :] ).to(gl.int32) q_rope = gl.amd.cdna4.buffer_load( - ptr=q_ptr, offsets=qr_off, mask=h_mask_q[:, None], other=0.0 + ptr=q_ptr, offsets=qr_off, mask=h_mask_q[:, None], other=0.0, cache=Q_CACHE ) q_rope_dot = gl.convert_layout(q_rope, cfg.q_layout) else: diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 016b0468a0d..c30f7843314 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -610,6 +610,9 @@ def _pick_splits(hb): heads_blocks = hb8 num_splits = ns8 + # Q is read once per query without split-K, and re-read by every split + q_cache = ".cg" if num_splits == 1 else "" + if num_splits > 1: part_m = torch.empty( (num_queries, num_splits, num_heads), dtype=torch.float32, device=q.device @@ -713,6 +716,7 @@ def _pick_splits(hb): NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", + Q_CACHE=q_cache, GRID_ORDER="qsh", # The partial last tile rides the full-tile body. Gluon inlines, so a peeled # masked copy would be a second gather+dequant+MFMA body, and its register diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index c17b489936c..87cb6618db5 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -209,7 +209,7 @@ def sparse_mla_fwd( kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts. kv_indices: flat int32 GLOBAL slot ids into the pool. softmax_scale: the layer's softmax scale. - kv_scale: ``[1]`` f32 per-tensor cache scale (``layer._k_scale``); + kv_scale: ``[1]`` f32 per-tensor cache scale; required for the flat fp8 format. kv_splits: split-K override; default follows the occupancy policy. skip_reduce: with split-K active, return ``(part_acc, part_m, part_l)`` @@ -227,11 +227,8 @@ def sparse_mla_fwd( asked for, fp8 q is passed straight through, and fp8 q under "bf16" dots is widened in-kernel, which is exact. q_scale: scalar f32. Required when q arrives already fp8 (the scale it - was quantized with; the aiter asm convention, where vLLM passes - the layer's calibrated ``_q_scale``). Optional when q is bf16 and - dot_precision="fp8": given, it is used as the static quantization - scale, matching production's scaled_fp8_quant(q, layer._q_scale); - omitted, a per-tensor amax is taken here. + was quantized with; the aiter asm convention. Optional when q is bf16 and + dot_precision="fp8". out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. return_lse: also return the natural-log log-sum-exp, [C, H] f32, for merging partials across context-parallel ranks. A fully masked row @@ -379,6 +376,8 @@ def sparse_mla_fwd( # num_warps stays 4: warps tile the dot's N (MFMA N=16), so a larger BLOCK_K just # gives each warp two N tiles. + # Q is read once per query without split-K, and re-read by every split + q_cache = ".cg" if num_splits == 1 else "" grid = (num_queries, num_splits, heads_blocks) _pa_decode_sparse_gfx950[grid]( q, @@ -435,6 +434,7 @@ def sparse_mla_fwd( PART_STORE_CACHE="", UNI_TILE=True, GRID_ORDER="qsh", + Q_CACHE=q_cache, MAIN_SPLITS=num_splits, ADAPTIVE_SPLITS=num_splits > 1, ASM_DEQ=False, From 83075b6701d554cf096e35fae71777ffe3ef3aac Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 27 Aug 2026 07:16:40 +0000 Subject: [PATCH 45/51] fix gen issue --- op_tests/triton_tests/attention/test_sparse_mla.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index 731de77b18b..b6198c1b6b5 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -61,14 +61,18 @@ def dequant_ds_mla(cache): def make_indices(C, pool, topk, device, seed, ragged): + # Built on CPU g = torch.Generator(device="cpu").manual_seed(seed) lens = ( - torch.randint(1, topk + 1, (C,), generator=g) + torch.randint(1, topk + 1, (C,), generator=g, device="cpu") if ragged - else torch.full((C,), topk) + else torch.full((C,), topk, device="cpu") ) - rows = [torch.randperm(pool, generator=g)[: int(lens[c])] for c in range(C)] - indptr = torch.zeros(C + 1, dtype=torch.int64) + rows = [ + torch.randperm(pool, generator=g, device="cpu")[: int(lens[c])] + for c in range(C) + ] + indptr = torch.zeros(C + 1, dtype=torch.int64, device="cpu") indptr[1:] = lens.cumsum(0) return ( torch.cat(rows).to(torch.int32).to(device), From 88126fc60bd9c519c9bd2c0ca7764f8cd24dc0fe Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Thu, 10 Sep 2026 21:54:10 +0000 Subject: [PATCH 46/51] better name kv cache types --- .../gfx950/attention/pa_decode_sparse.py | 138 +++++++++--------- .../ops/triton/attention/pa_decode_sparse.py | 38 +++-- aiter/ops/triton/attention/sparse_mla.py | 19 ++- 3 files changed, 101 insertions(+), 94 deletions(-) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py index 3f9f3c84a52..8df21a23319 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py @@ -14,20 +14,22 @@ K-only rope and the QK contraction chains a second MFMA into the first. The KV buffer is the entire V in both geometries. -Cache formats (Fmt.KIND), per segment: - - "bf16" full-row bf16 - "dsv4" 448 fp8 | 64 bf16 rope (576 B rows) + 8 B UE8M0-per-64 trailer - "uniform" whole-row fp8 + separate f32 per-64 kv_scales - "tensor" whole-row fp8 + a single per-tensor f32 scale (k_scale). The - scale never touches the tile loop: K-side folds into qk_scale, - V-side into p. - "dsmla" vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope - (656 B rows); requires ROPE_SEPARATE. - -Launchers: -aiter/ops/triton/attention/pa_decode_sparse.py (DSv4 / uniform pool) and -aiter/ops/triton/attention/sparse_mla.py (separated-rope MLA). +Cache formats (Fmt.KIND), per segment. The two packed kinds are both vLLM's +``--kv-cache-dtype fp8_ds_mla``, which is one name for two byte layouts +distinguished by model generation, so they carry the generation instead: + + "bf16" full-row bf16 + "fp8_scalar" whole-row fp8 + a single per-tensor f32 scale (k_scale). + The scale never touches the tile loop: K-side folds into + qk_scale, V-side into p. + "fp8_g64" whole-row fp8 + separate f32 per-64 kv_scales + "fp8_dsv4_mla" DeepSeek-V4, 584 B per token: 448 fp8 | 64 bf16 rope + (576 B data rows) + an 8 B UE8M0-per-64 trailer after the + block's rows (7 scales + 1 B pad). Rope lives inside the + 512-wide head dim, so V is the whole row. + "fp8_dsv32_mla" DeepSeek-V3.2 (also Kimi-K3), 656 B per token, token-major + records: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope. + Rope is appended, so this requires ROPE_SEPARATE. """ from triton.experimental import gluon @@ -253,7 +255,7 @@ class Cfg: HAS_INVALID: gl.constexpr HEAD_ALIGNED: gl.constexpr IDX_BUFFER_LOAD: gl.constexpr - FP8_MFMA: gl.constexpr # "tensor" only: feed the matrix core the cache's + FP8_MFMA: gl.constexpr # "fp8_scalar" only: feed the matrix core the cache's # own fp8 instead of dequantizing to bf16 # Cache policy per load site GATHER_CACHE: gl.constexpr @@ -454,7 +456,9 @@ def __init__( class Fmt: """Compile-time description of one segment's cache format.""" - KIND: gl.constexpr # "bf16" | "dsv4" | "uniform" | "tensor" | "dsmla" + # "bf16" | "fp8_scalar" | "fp8_g64" + # | "fp8_dsv4_mla" (584 B/token) | "fp8_dsv32_mla" (656 B/token) + KIND: gl.constexpr IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop BLOCK_SIZE: gl.constexpr USE_BUFFER_LOAD: gl.constexpr @@ -501,7 +505,7 @@ def __init__( KV_DIM = cfg.KV_DIM.value ROPE_DIM = cfg.ROPE_DIM.value - GROUP = 128 if KIND == "dsmla" else 64 + GROUP = 128 if KIND == "fp8_dsv32_mla" else 64 NG = KV_DIM // GROUP self.GROUP = gl.constexpr(GROUP) self.NG = gl.constexpr(NG) @@ -527,16 +531,16 @@ def __init__( # dsv4 row: [NOPE_DIM fp8 | ROPE_DIM bf16] + 8 B UE8M0 per token after # the block; dsmla row: [KV_DIM fp8 | NG f32 | ROPE_DIM bf16] inline. - if KIND == "dsv4": + if KIND == "fp8_dsv4_mla": TOK_U8 = NOPE_DIM + 2 * ROPE_DIM # 448 + 128 = 576 - elif KIND == "dsmla": + elif KIND == "fp8_dsv32_mla": TOK_U8 = KV_DIM + 4 * NG + 2 * ROPE_DIM # 512 + 16 + 128 = 656 else: TOK_U8 = 0 self.TOK_U8 = gl.constexpr(TOK_U8) self.TOK_U16 = gl.constexpr(TOK_U8 // 2) self.ROPE_U16_OFF = gl.constexpr( - (NOPE_DIM // 2) if KIND == "dsv4" else ((KV_DIM + 4 * NG) // 2) + (NOPE_DIM // 2) if KIND == "fp8_dsv4_mla" else ((KV_DIM + 4 * NG) // 2) ) self.SCL_TRAILER_U8 = gl.constexpr(8) self.TOK_F32 = gl.constexpr(TOK_U8 // 4) @@ -544,7 +548,7 @@ def __init__( # Flat formats gather in cache elements: KV_DIM wide, plus the appended # rope when the geometry separates it. self.TOK_EL = gl.constexpr( - KV_DIM + (ROPE_DIM if cfg.ROPE_SEPARATE.value and KIND != "uniform" else 0) + KV_DIM + (ROPE_DIM if cfg.ROPE_SEPARATE.value and KIND != "fp8_g64" else 0) ) @@ -554,12 +558,12 @@ class Seg: """Runtime context of one segment. Pointer roles by format ("--" = unused duplicate): - KIND cache_ptr alt_ptr scl_ptr - bf16 -- bf16 cache -- - dsv4 u8 cache bf16 view of the cache -- - uniform u8 cache f32 per-64 kv_scales -- - tensor u8 cache -- f32 scalar k_scale - dsmla u8 cache bf16 view (rope tail) f32 view (scales) + KIND cache_ptr alt_ptr scl_ptr + bf16 -- bf16 cache -- + fp8_scalar u8 cache -- f32 scalar k_scale + fp8_g64 u8 cache f32 per-64 kv_scales -- + fp8_dsv4_mla u8 cache bf16 view of the cache -- + fp8_dsv32_mla u8 cache bf16 view (rope tail) f32 view (scales) """ fmt: Fmt @@ -589,7 +593,7 @@ def __init__( def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 (gfx950 has no bf16 multiply). sc: raw UE8M0 byte (dsv4), f32 scale - (uniform/dsmla), or unused ("tensor": bare fp8 -> bf16 convert).""" + (uniform/dsmla), or unused ("fp8_scalar": bare fp8 -> bf16 convert).""" if fmt.ASM_DEQ: # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. W8: gl.constexpr = x_u8.shape[1] * 2 @@ -603,10 +607,10 @@ def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): else: kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) else: - if fmt.KIND == "tensor": + if fmt.KIND == "fp8_scalar": val = _fp8_to_bf16(x_u8, fmt.FP8_FNUZ) else: - if fmt.KIND == "uniform" or fmt.KIND == "dsmla": + if fmt.KIND == "fp8_g64" or fmt.KIND == "fp8_dsv32_mla": scale = sc else: scale = gl.exp2(sc.to(gl.float32) - 127.0) @@ -897,7 +901,7 @@ def _gather_full( False, cfg.UNI_TILE, ) - if fmt.KIND == "uniform": + if fmt.KIND == "fp8_g64": NGRP: gl.constexpr = cfg.KV_DIM // 64 x_u8 = _cache_load( seg.cache_ptr, @@ -914,7 +918,7 @@ def _gather_full( CACHE=cfg.GATHER_CACHE, ) k_rope = x_u8 # no rope side-channel -> DCE'd - elif fmt.KIND == "tensor": + elif fmt.KIND == "fp8_scalar": # Per-tensor scale is folded outside the loop (qk_scale / p), so this # is a bare gather; the K-only rope tail follows when separated. x_u8 = _cache_load( @@ -944,7 +948,7 @@ def _gather_full( ) else: k_rope = x_u8 # rope lives inside the KV buffer -> DCE'd - elif fmt.KIND == "dsmla": + elif fmt.KIND == "fp8_dsv32_mla": nope_row = bg * cs0 + pg * fmt.TOK_U8 scl_row = bg * (cs0 // 4) + pg * fmt.TOK_F32 + fmt.SCL_F32_OFF # Scales issue before the bulk fp8: vmcnt is one in-order FIFO, so @@ -994,7 +998,7 @@ def _gather_full( fmt.USE_BUFFER_LOAD, CACHE=cfg.GATHER_CACHE, ) - else: # "dsv4" + else: # "fp8_dsv4_mla" nope_row = bg * cs0 + pg * fmt.TOK_U8 scl_row = bg * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pg * fmt.SCL_TRAILER_U8 # Scales first (vmcnt FIFO; see the dsmla branch). @@ -1061,7 +1065,7 @@ def _gather_full( @gluon.jit def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): """Write one prefetched tile into the LDS buffer(s). The KV buffer is always the - full KV_DIM-wide dequant: "dsv4" gathers KV_DIM bytes too, the last 64 being + full KV_DIM-wide dequant: "fp8_dsv4_mla" gathers KV_DIM bytes too, the last 64 being bf16 rope read as garbage fp8 and overwritten by the slice-store below, so the gather stays pow-2 wide.""" fmt = seg.fmt @@ -1074,13 +1078,13 @@ def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): rope_smem.store(k_rope.to(gl.float8e4nv, bitcast=True)) else: _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) - if fmt.KIND == "dsv4": + if fmt.KIND == "fp8_dsv4_mla": kv_smem.slice(fmt.NOPE_DIM, cfg.ROPE_DIM, dim=1).store(k_rope) - elif fmt.KIND == "dsmla": + elif fmt.KIND == "fp8_dsv32_mla": rope_smem.store(k_rope) - elif fmt.KIND == "tensor" and cfg.ROPE_SEPARATE: + elif fmt.KIND == "fp8_scalar" and cfg.ROPE_SEPARATE: rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) - # "uniform": the whole head is one fp8 tile; nothing else to store. + # "fp8_g64": the whole head is one fp8 tile; nothing else to store. @gluon.jit @@ -1178,7 +1182,7 @@ def _qkpv_lds( ] S = gl.where(col_mask, S, neg_inf) # Online softmax in the base-2 exponent domain; m_i carries qk_scale - # (= scale * log2e [* k_scale for "tensor"]). max commutes with a positive + # (= scale * log2e [* k_scale for "fp8_scalar"]). max commutes with a positive # scale, so scale the row max instead of every element of S; what is left, # S * qk_scale - m_new, lowers to one FMA, and -inf columns stay -inf. m_block = _rmax(S, 1) * qk_scale @@ -1193,9 +1197,9 @@ def _qkpv_lds( v = kv_smem.load(cfg.v_layout) if cfg.ASYNC_LDS: v = v.to(gl.float8e4nv, bitcast=True) - # "tensor": V was staged as raw fp8 code points; apply the per-tensor scale + # "fp8_scalar": V was staged as raw fp8 code points; apply the per-tensor scale # on the small side (p) and leave l scale-free: out = sum(p*s*V)/l exactly. - if seg.fmt.KIND == "tensor" and not cfg.FP8_MFMA: + if seg.fmt.KIND == "fp8_scalar" and not cfg.FP8_MFMA: p = p * v_scale if cfg.FP8_MFMA: p_dot = gl.convert_layout(p.to(gl.float8e4nv), cfg.p_layout) @@ -1257,7 +1261,7 @@ def _decode_tile( if MASKED: valid_g = gl.convert_layout(valid1d, gl.SliceLayout(1, cfg.gather_l)) - if fmt.KIND == "uniform": + if fmt.KIND == "fp8_g64": NGRP: gl.constexpr = cfg.KV_DIM // 64 kv_row = block_idx_g * cs0 + pos_g * cfg.KV_DIM scl_row = block_idx_g * NGRP @@ -1297,7 +1301,7 @@ def _decode_tile( CACHE=cfg.GATHER_CACHE, ) _deq_store_tile(x_u8, sc, kv_smem, cfg, fmt) - elif fmt.KIND == "dsv4": + elif fmt.KIND == "fp8_dsv4_mla": nope_row = block_idx_g * cs0 + pos_g * fmt.TOK_U8 scl_row = ( block_idx_g * cs0 + fmt.BLOCK_SIZE * fmt.TOK_U8 + pos_g * fmt.SCL_TRAILER_U8 @@ -1504,7 +1508,7 @@ def _decode_tile( l_new = l_i * alpha + gl.sum(p, axis=1) v = kv_smem.load(cfg.v_layout) # [BLOCK_K, KV_DIM] - if seg.fmt.KIND == "tensor": + if seg.fmt.KIND == "fp8_scalar": p = p * v_scale # per-tensor V scale on the small side (see _qkpv) p_dot = gl.convert_layout(p.to(gl.bfloat16), cfg.p_layout) alpha_pv = gl.convert_layout(alpha, gl.SliceLayout(1, cfg.pv_layout)) @@ -1681,7 +1685,7 @@ def _pa_decode_sparse( # cache pointers are the same allocation under different element types; # which of them is live depends on the format (see Seg). main_cache_ptr, # main (SWA) cache, u8 view - main_cache_bf16_ptr, # bf16 view of it, or the f32 scale pool ("uniform") + main_cache_bf16_ptr, # bf16 view of it, or the f32 scale pool ("fp8_g64") main_indices_ptr, # [nnz_main] int32 row ids main_indptr_ptr, # [C + 1] int32 extra_cache_ptr, # top-k segment; aliases main when HAS_EXTRA=False @@ -1695,8 +1699,8 @@ def _pa_decode_sparse( part_m_ptr, # [C, NUM_SPLITS, H] f32 row max, base-2 domain part_l_ptr, # [C, NUM_SPLITS, H] f32 row sum part_acc_ptr, # [C, NUM_SPLITS, H, S] bf16 or f32, un-normalized - # f32 side-channel per segment: scalar k_scale ("tensor") or f32 cache view - # ("dsmla"). None elides the argument, keeping other formats' kernarg + # f32 side-channel per segment: scalar k_scale ("fp8_scalar") or f32 cache view + # ("fp8_dsv32_mla"). None elides the argument, keeping other formats' kernarg # layouts unchanged. main_scl_ptr, extra_scl_ptr, @@ -1777,20 +1781,20 @@ def _pa_decode_sparse( un-normalized partials for the reduce kernel.""" NUM_WARPS: gl.constexpr = gl.num_warps() gl.static_assert( - UNI_TILE or (MAIN_FMT != "tensor" and MAIN_FMT != "dsmla"), + UNI_TILE or (MAIN_FMT != "fp8_scalar" and MAIN_FMT != "fp8_dsv32_mla"), "tensor/dsmla formats require UNI_TILE=1", ) gl.static_assert( - UNI_TILE or (EXTRA_FMT != "tensor" and EXTRA_FMT != "dsmla"), + UNI_TILE or (EXTRA_FMT != "fp8_scalar" and EXTRA_FMT != "fp8_dsv32_mla"), "tensor/dsmla formats require UNI_TILE=1", ) gl.static_assert( - (not ROPE_SEPARATE) or (MAIN_FMT != "dsv4" and MAIN_FMT != "uniform"), + (not ROPE_SEPARATE) or (MAIN_FMT != "fp8_dsv4_mla" and MAIN_FMT != "fp8_g64"), "dsv4/uniform formats carry rope inside the KV buffer (ROPE_SEPARATE=False)", ) gl.static_assert( - MAIN_FMT != "dsmla" or ROPE_SEPARATE, - "fp8_ds_mla is a separated-rope (MLA) format", + MAIN_FMT != "fp8_dsv32_mla" or ROPE_SEPARATE, + "fp8_dsv32_mla is a separated-rope (MLA) format", ) # No rope means no rope buffer, and the packed formats all define a rope tail. gl.static_assert( @@ -1800,14 +1804,18 @@ def _pa_decode_sparse( gl.static_assert( ROPE_DIM > 0 or ( - MAIN_FMT != "dsv4" - and MAIN_FMT != "dsmla" - and (not HAS_EXTRA or (EXTRA_FMT != "dsv4" and EXTRA_FMT != "dsmla")) + MAIN_FMT != "fp8_dsv4_mla" + and MAIN_FMT != "fp8_dsv32_mla" + and ( + not HAS_EXTRA + or (EXTRA_FMT != "fp8_dsv4_mla" and EXTRA_FMT != "fp8_dsv32_mla") + ) ), - "dsv4/fp8_ds_mla rows carry a rope tail; ROPE_DIM=0 is inconsistent", + "fp8_dsv4_mla/fp8_dsv32_mla rows carry a rope tail; ROPE_DIM=0 is " + "inconsistent", ) gl.static_assert( - (not ASM_DEQ) or MAIN_FMT == "dsv4" or EXTRA_FMT == "dsv4", + (not ASM_DEQ) or MAIN_FMT == "fp8_dsv4_mla" or EXTRA_FMT == "fp8_dsv4_mla", "ASM_DEQ is the dsv4 E8M0 dequant", ) # The fp8 path needs one positive scalar scale per cache, since that is what @@ -1815,7 +1823,7 @@ def _pa_decode_sparse( # core reads. gl.static_assert( (not FP8_MFMA) - or (MAIN_FMT == "tensor" and (not HAS_EXTRA or EXTRA_FMT == "tensor")), + or (MAIN_FMT == "fp8_scalar" and (not HAS_EXTRA or EXTRA_FMT == "fp8_scalar")), "FP8_MFMA requires the per-tensor fp8 format on every segment", ) gl.static_assert(not (FP8_MFMA and FP8_FNUZ), "FP8_MFMA is OCP e4m3 only") @@ -1869,7 +1877,7 @@ def _pa_decode_sparse( MAIN_BLOCK_SIZE, MAIN_USE_BUFFER_LOAD, FP8_FNUZ, - ASM_DEQ and MAIN_FMT == "dsv4", + ASM_DEQ and MAIN_FMT == "fp8_dsv4_mla", NOPE_DIM, NOPE_CHUNK, CHUNK_AXIS, @@ -1880,7 +1888,7 @@ def _pa_decode_sparse( EXTRA_BLOCK_SIZE, EXTRA_USE_BUFFER_LOAD, FP8_FNUZ, - ASM_DEQ and EXTRA_FMT == "dsv4", + ASM_DEQ and EXTRA_FMT == "fp8_dsv4_mla", NOPE_DIM, NOPE_CHUNK, CHUNK_AXIS, @@ -1907,13 +1915,13 @@ def _pa_decode_sparse( qk_scale = scale * RCP_LN2 main_qk_scale = qk_scale main_v_scale = 1.0 - if MAIN_FMT == "tensor": + if MAIN_FMT == "fp8_scalar": main_k_scale = gl.load(main_scl_ptr) main_qk_scale = qk_scale * main_k_scale main_v_scale = main_k_scale extra_qk_scale = qk_scale extra_v_scale = 1.0 - if HAS_EXTRA and EXTRA_FMT == "tensor": + if HAS_EXTRA and EXTRA_FMT == "fp8_scalar": extra_k_scale = gl.load(extra_scl_ptr) extra_qk_scale = qk_scale * extra_k_scale extra_v_scale = extra_k_scale @@ -2051,7 +2059,7 @@ def _pa_decode_sparse( main_fmt, main_cache_ptr, main_cache_bf16_ptr, - main_scl_ptr if (MAIN_FMT == "dsmla") else main_cache_ptr, + main_scl_ptr if (MAIN_FMT == "fp8_dsv32_mla") else main_cache_ptr, main_indices_ptr, main_start, main_cs0, @@ -2100,7 +2108,7 @@ def _pa_decode_sparse( extra_fmt, extra_cache_ptr, extra_cache_bf16_ptr, - extra_scl_ptr if (EXTRA_FMT == "dsmla") else extra_cache_ptr, + extra_scl_ptr if (EXTRA_FMT == "fp8_dsv32_mla") else extra_cache_ptr, extra_indices_ptr, extra_start, extra_cs0, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index c30f7843314..964564ee1e9 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -1,17 +1,10 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Sparse paged-decode attention over a unified KV pool with per-token paged -indices. See ``_triton_kernels/attention/pa_decode_sparse.py`` for the -kernels' caller contract. +"""Sparse paged-decode attention over a unified or split KV pool with per-token paged +indices. -This module exposes ``pa_decode_sparse`` — a 3D split-K + widened-BLOCK_H -+ pipelined-K-loop variant suitable for sparse decode (e.g. V4 top-k gather) -where each token's K range is an unordered subset of a unified KV pool. - -On gfx950 (CDNA4) DeepSeek-V4 sparse-MLA decode has a dedicated gluon -implementation (bottom of this module): ``pa_decode_sparse`` routes all formats -to the merged ``_pa_decode_sparse_gfx950_gluon`` driver. +TODO: add details once API has settled """ import math @@ -110,7 +103,8 @@ def pa_decode_sparse( two-loop's second (top-k) cache + index set; must be None otherwise. On gfx950 the DSv4 gluon driver handles this: a 3D ``unified_kv`` selects the - packed fp8_ds_mla / bf16 block cache (``extra_*`` = the two-loop), a 2D one the + packed fp8_dsv4_mla (584 B rows) / bf16 block cache (``extra_*`` = the + two-loop), a 2D one the uniform pool (``kv_scales`` present = fp8). ``kv_splits``/``skip_reduce`` are honored; ``block_h`` and fp16 ``q`` fall through to the triton path. @@ -134,7 +128,7 @@ def pa_decode_sparse( raise RuntimeError(f"pa_decode_sparse expects fp16/bf16 q, got {q.dtype}") # gfx950: route to the merged DSv4 sparse-MLA gluon driver. Format is inferred - # from the cache: 3D -> packed fp8_ds_mla / bf16 block cache (optional SWA+top-k + # from the cache: 3D -> packed fp8_dsv4_mla / bf16 block cache (optional SWA+top-k # two-loop via extra_*); 2D -> uniform pool (OCP fp8 + fp32 kv_scales, or bf16). # kv_splits and skip_reduce are honored here; block_h and fp16 q fall through to # the triton path below (the gluon kernel is bf16-only: bf16 LDS + bf16 MFMA). @@ -474,7 +468,7 @@ def _pa_decode_sparse_gfx950_gluon( fp8_fnuz=False, ): """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from cache.ndim: - 3D [nb, block, ...] -> packed fp8_ds_mla (uint8: 448 NoPE fp8 e4m3 OCP + + 3D [nb, block, 584] -> packed fp8_dsv4_mla (uint8: 448 NoPE fp8 e4m3 OCP + embedded UE8M0 per-64 scale + 64 RoPE bf16) or a bf16 block cache; pass extra_* for the SWA+top-k two-loop, else a single segment. @@ -510,7 +504,7 @@ def _pa_decode_sparse_gfx950_gluon( if cache.ndim == 2: # uniform pool: one fp8 gather over the whole head + separate fp32 scales, # or bf16. page_size=1 -> block_idx=slot, pos=0; scales ride the bf16 ptr. - UNIFORM = True + FLAT_POOL = True main_is_fp8 = cache.dtype == torch.uint8 if main_is_fp8: assert cache_scales is not None and cache_scales.dtype == torch.float32 @@ -532,8 +526,10 @@ def _pa_decode_sparse_gfx950_gluon( avg_main = indices.numel() / max(1, num_queries) # one segment; no extra avg_extra = 0.0 else: - # packed fp8_ds_mla [nb, block, 584] (embedded scale) or bf16 block cache. - UNIFORM = False + # packed fp8_dsv4_mla [nb, block, 584] (UE8M0 block trailer) or bf16 block + # cache. NB: vLLM's fp8_ds_mla also names the 656 B V3.2 layout, which is + # fp8_dsv32_mla and reaches the kernel through sparse_mla.py instead. + FLAT_POOL = False main_is_fp8 = cache.dtype == torch.uint8 main_bf16 = cache.view(torch.bfloat16) if main_is_fp8 else cache has_extra = ( @@ -556,12 +552,12 @@ def _pa_decode_sparse_gfx950_gluon( avg_extra = extra_indices.numel() / max(1, num_queries) if has_extra else 0.0 # Kernel-side cache-format tags (kernel shared with sparse_mla.py). - if UNIFORM: - main_fmt = "uniform" if main_is_fp8 else "bf16" + if FLAT_POOL: + main_fmt = "fp8_g64" if main_is_fp8 else "bf16" extra_fmt = main_fmt else: - main_fmt = "dsv4" if main_is_fp8 else "bf16" - extra_fmt = "dsv4" if extra_is_fp8 else "bf16" + main_fmt = "fp8_dsv4_mla" if main_is_fp8 else "bf16" + extra_fmt = "fp8_dsv4_mla" if extra_is_fp8 else "bf16" # Alignment hint for the page strides so row gathers can vectorize: the largest # power of 2 (<= 16) dividing both. @@ -656,7 +652,7 @@ def _pick_splits(hb): adaptive_splits = num_splits > 1 # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm - asm_deq = one_wg_per_cu and not UNIFORM and not fp8_fnuz and main_is_fp8 + asm_deq = one_wg_per_cu and not FLAT_POOL and not fp8_fnuz and main_is_fp8 # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear # workgroup id, so the axis order decides what shares an XCD's L2. diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 87cb6618db5..acbbacbb30b 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -79,20 +79,22 @@ def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): assert kv_scale.dtype == torch.float32 u8 = kv.view(torch.uint8) fnuz = kv.dtype == torch.float8_e4m3fnuz - return "tensor", u8, u8, kv_scale.reshape(1), 1, fnuz + return "fp8_scalar", u8, u8, kv_scale.reshape(1), 1, fnuz if ( 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 ): - # vLLM fp8_ds_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B + # fp8_dsv32_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B. + # This is vLLM's fp8_ds_mla on V3.2 / Kimi-K3; the same vLLM name also + # covers V4's 584 B layout (fp8_dsv4_mla), hence the explicit generation. u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) assert ( u8.stride(2) == 1 and u8.stride(1) == u8.shape[2] - ), "fp8_ds_mla rows must be contiguous 656-byte records" + ), "fp8_dsv32_mla rows must be contiguous 656-byte records" return ( - "dsmla", + "fp8_dsv32_mla", u8, u8.view(torch.bfloat16), u8.view(torch.float32), @@ -163,9 +165,9 @@ def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool ) if dot_precision == "bf16": return False - if fmt == "dsmla": + if fmt == "fp8_dsv32_mla": raise ValueError( - "dot_precision='fp8' does not support the fp8_ds_mla cache." + "dot_precision='fp8' does not support the fp8_dsv32_mla cache." "Use dot_precision='bf16'." ) if fmt == "bf16": @@ -205,7 +207,8 @@ def sparse_mla_fwd( the QK contraction becomes a single dot over kv_lora_rank. kv_buffer: the KV pool [nb, block, R] / [slots, 1, 1, R] / [slots, R] in bf16, the same shapes in fp8 (+ scalar - kv_scale), or [nb, block, 656] uint8 (vLLM ``fp8_ds_mla``). + kv_scale), or [nb, block, 656] uint8 (``fp8_dsv32_mla``, which is + vLLM's ``fp8_ds_mla`` on DeepSeek-V3.2 / Kimi-K3). kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts. kv_indices: flat int32 GLOBAL slot ids into the pool. softmax_scale: the layer's softmax scale. @@ -394,7 +397,7 @@ def sparse_mla_fwd( part_m, part_l, part_acc, - scl, # f32 side-channel: k_scale ("tensor") / f32 view ("dsmla") + scl, # f32 side-channel: k_scale ("fp8_scalar") / f32 view ("fp8_dsv32_mla") scl, float(softmax_scale), q.stride(0), From b286403238ea50bd32930eff0ee0f83bd5e440ee Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Date: Tue, 15 Sep 2026 09:39:12 -0500 Subject: [PATCH 47/51] Cagri/sparse mla unified api (#5551) --- .../{pa_decode_sparse.py => sparse_mla.py} | 110 ++-- .../ops/triton/attention/pa_decode_sparse.py | 111 ++-- aiter/ops/triton/attention/sparse_mla.py | 545 ++++++++++++------ .../op_benchmarks/triton/bench_sparse_mla.py | 4 +- .../attention/test_pa_decode_sparse.py | 9 +- .../triton_tests/attention/test_sparse_mla.py | 2 +- 6 files changed, 477 insertions(+), 304 deletions(-) rename aiter/ops/triton/_gluon_kernels/gfx950/attention/{pa_decode_sparse.py => sparse_mla.py} (96%) diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py similarity index 96% rename from aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py rename to aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py index 8df21a23319..b5df196d086 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_mla.py @@ -14,9 +14,9 @@ K-only rope and the QK contraction chains a second MFMA into the first. The KV buffer is the entire V in both geometries. -Cache formats (Fmt.KIND), per segment. The two packed kinds are both vLLM's -``--kv-cache-dtype fp8_ds_mla``, which is one name for two byte layouts -distinguished by model generation, so they carry the generation instead: +Cache formats (Fmt.KIND), per segment. vLLM calls both packed kinds +--kv-cache-dtype fp8_ds_mla, one name for two byte layouts told apart by model +generation, so the tags here carry the generation instead: "bf16" full-row bf16 "fp8_scalar" whole-row fp8 + a single per-tensor f32 scale (k_scale). @@ -91,19 +91,13 @@ def _cache_load( @gluon.jit -def _fp8_to_f32(x_u8, FP8_FNUZ: gl.constexpr): - # fnuz -> f32 has no native cvt on gfx950 (software unpack); fnuz -> bf16 is - # cheap and exact (3 mantissa bits), so route fnuz through bf16. - if FP8_FNUZ: - return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16).to(gl.float32) +def _fp8_to_f32(x_u8): return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.float32) @gluon.jit -def _fp8_to_bf16(x_u8, FP8_FNUZ: gl.constexpr): +def _fp8_to_bf16(x_u8): # Exact: fp8's 3 mantissa bits fit bf16's 8. - if FP8_FNUZ: - return x_u8.to(gl.float8e4b8, bitcast=True).to(gl.bfloat16) return x_u8.to(gl.float8e4nv, bitcast=True).to(gl.bfloat16) @@ -294,10 +288,7 @@ def __init__( KV_DIM, ROPE_DIM, ROPE_SEPARATE, - MFMA_K, NUM_WARPS, - GATHER_TW1, - LDS_PAD, UNI_TILE, HAS_INVALID, HEAD_ALIGNED, @@ -307,7 +298,6 @@ def __init__( IDX_CACHE="", ASYNC_LDS=False, RELAXED_LOAD=True, - ROPE_VEC=16, PAD_INTERVAL=1024, ): self.BLOCK_M = gl.constexpr(BLOCK_M) @@ -317,7 +307,10 @@ def __init__( self.ROPE_SEPARATE = gl.constexpr(ROPE_SEPARATE) self.QK_DIM = gl.constexpr(KV_DIM + (ROPE_DIM if ROPE_SEPARATE else 0)) self.NUM_WARPS = gl.constexpr(NUM_WARPS) + # Threads the row gather spends on the head dim + GATHER_TW1 = 32 self.GATHER_TW1 = gl.constexpr(GATHER_TW1) + LDS_PAD = 16 if FP8_MFMA else 8 self.LDS_PAD = gl.constexpr(LDS_PAD) self.UNI_TILE = gl.constexpr(UNI_TILE) self.HAS_INVALID = gl.constexpr(HAS_INVALID) @@ -328,12 +321,9 @@ def __init__( self.IDX_CACHE = gl.constexpr(IDX_CACHE) self.ASYNC_LDS = gl.constexpr(ASYNC_LDS) self.RELAXED_LOAD = gl.constexpr(RELAXED_LOAD) + ROPE_VEC = 16 self.ROPE_VEC = gl.constexpr(ROPE_VEC) - # K=32 is the only shape the backend offers for plain fp8 operands, and - # it divides KV_DIM, ROPE_DIM and BLOCK_K alike, so one layout still - # covers both dots. - if FP8_MFMA: - MFMA_K = 32 + MFMA_K = 32 if FP8_MFMA else 16 self.MFMA_K = gl.constexpr(MFMA_K) self.qk_layout = gl.constexpr( @@ -358,9 +348,9 @@ def __init__( self.p_layout = gl.constexpr(gl.DotOperandLayout(0, self.pv_layout, KW)) self.v_layout = gl.constexpr(gl.DotOperandLayout(1, self.pv_layout, KW)) - # 16 uint8 per thread = 128-bit gather loads - # GATHER_TW1 = threads spent on the head dim: 32 requests a whole - # 512 B token row per instruction, at the cost of a longer slot vector. + # 16 uint8 per thread = 128-bit gather loads. GATHER_TW1 lanes span a row, + # so one gather covers GATHER_TW1 * 16 = 512 B of it, at the cost of a + # longer slot vector. GSPT = 16 self.gather_l = gl.constexpr( gl.BlockedLayout( @@ -462,8 +452,8 @@ class Fmt: IS_FP8: gl.constexpr # pipeline select: prefetched fp8 loop vs bf16 loop BLOCK_SIZE: gl.constexpr USE_BUFFER_LOAD: gl.constexpr - FP8_FNUZ: gl.constexpr - ASM_DEQ: gl.constexpr + # dsv4 E8M0 dequant: "none" | "upcast" | "asm" (gathers the int16 view) + DEQ: gl.constexpr NOPE_CHUNK: gl.constexpr CHUNK_AXIS: gl.constexpr NOPE_DIM: gl.constexpr # fp8 payload width (448 dsv4; KV_DIM elsewhere) @@ -487,8 +477,7 @@ def __init__( KIND, BLOCK_SIZE, USE_BUFFER_LOAD, - FP8_FNUZ, - ASM_DEQ, + DEQ, NOPE_DIM, NOPE_CHUNK, CHUNK_AXIS, @@ -497,8 +486,7 @@ def __init__( self.IS_FP8 = gl.constexpr(KIND != "bf16") self.BLOCK_SIZE = gl.constexpr(BLOCK_SIZE) self.USE_BUFFER_LOAD = gl.constexpr(USE_BUFFER_LOAD) - self.FP8_FNUZ = gl.constexpr(FP8_FNUZ) - self.ASM_DEQ = gl.constexpr(ASM_DEQ) + self.DEQ = gl.constexpr(DEQ if KIND == "fp8_dsv4_mla" else "none") self.NOPE_CHUNK = gl.constexpr(NOPE_CHUNK) self.CHUNK_AXIS = gl.constexpr(CHUNK_AXIS) self.NOPE_DIM = gl.constexpr(NOPE_DIM) @@ -594,7 +582,7 @@ def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): """Dequant one fp8 slab into kv_smem[:, off:off+W]. Dequant stays in f32 (gfx950 has no bf16 multiply). sc: raw UE8M0 byte (dsv4), f32 scale (uniform/dsmla), or unused ("fp8_scalar": bare fp8 -> bf16 convert).""" - if fmt.ASM_DEQ: + if fmt.DEQ == "asm": # x_u8 is the int16 view here (see _gather_full), so its column count is W8/2. W8: gl.constexpr = x_u8.shape[1] * 2 # Adjacent fp8 columns share a scale (groups are even), so dropping @@ -607,14 +595,21 @@ def _deq_store(x_u8, sc, kv_smem, off, cfg, fmt, AXIS: gl.constexpr): else: kv_smem.slice(off, x_u8.shape[0], dim=0).store(val) else: - if fmt.KIND == "fp8_scalar": - val = _fp8_to_bf16(x_u8, fmt.FP8_FNUZ) + if fmt.DEQ == "upcast": + # Upstream fused fp8 x E8M0 -> bf16 upcast; the driver only asks for + # it when cdna4.scaled_upcast exists, else DEQ="asm" (_deq_asm). + # sc is the raw E8M0 byte, already in x_u8's shape and layout. + val = gl.amd.cdna4.scaled_upcast( + x_u8.to(gl.float8e4nv, bitcast=True), sc, gl.bfloat16 + ) + elif fmt.KIND == "fp8_scalar": + val = _fp8_to_bf16(x_u8) else: if fmt.KIND == "fp8_g64" or fmt.KIND == "fp8_dsv32_mla": scale = sc else: scale = gl.exp2(sc.to(gl.float32) - 127.0) - val = (_fp8_to_f32(x_u8, fmt.FP8_FNUZ) * scale).to(gl.bfloat16) + val = (_fp8_to_f32(x_u8) * scale).to(gl.bfloat16) if AXIS == 1: kv_smem.slice(off, x_u8.shape[1], dim=1).store(val) else: @@ -824,7 +819,7 @@ def _async_segment( """Direct-to-LDS tile walk, one LDS buffer. vmcnt is one in-order FIFO, so the wait goes before any plain load in the body: a - plain load outstanding past the newest copy group would drag `wait_group` into + plain load outstanding past the newest copy group would drag wait_group into waiting for that copy too. The body has no masked variant: the slot read is clamped and the partial tail falls out of the score mask (UNI_TILE). """ @@ -1023,7 +1018,7 @@ def _gather_full( fmt.USE_BUFFER_LOAD, CACHE=cfg.GATHER_CACHE, ) - if fmt.ASM_DEQ: + if fmt.DEQ == "asm": # 2-byte elements: <2 x i16> = 4 packed fp8 per VGPR out of one # dword load. Same byte as nope_row (both even), addressed through # the bf16 view; the layout convert is a rename (shared dim-0 tiling). @@ -1083,7 +1078,7 @@ def _stage(cfg, seg, x_u8, sc, k_rope, kv_smem, rope_smem): elif fmt.KIND == "fp8_dsv32_mla": rope_smem.store(k_rope) elif fmt.KIND == "fp8_scalar" and cfg.ROPE_SEPARATE: - rope_smem.store(_fp8_to_bf16(k_rope, fmt.FP8_FNUZ)) + rope_smem.store(_fp8_to_bf16(k_rope)) # "fp8_g64": the whole head is one fp8 tile; nothing else to store. @@ -1331,7 +1326,7 @@ def _decode_tile( other=127, CACHE=cfg.GATHER_CACHE, ) - if fmt.ASM_DEQ: + if fmt.DEQ == "asm": row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) @@ -1378,7 +1373,7 @@ def _decode_tile( fmt.USE_BUFFER_LOAD, CACHE=cfg.GATHER_CACHE, ) - if fmt.ASM_DEQ: + if fmt.DEQ == "asm": row16 = gl.convert_layout( nope_row >> 1, gl.SliceLayout(1, cfg.gather16_l) ) @@ -1668,14 +1663,14 @@ def _process_segment( return m_i, l_i, acc -_pa_decode_sparse_repr = make_kernel_repr( - "_pa_decode_sparse", +_sparse_mla_repr = make_kernel_repr( + "_sparse_mla", ["BLOCK_M", "BLOCK_K", "HEAD_SIZE", "NUM_SPLITS", "MAIN_FMT", "ROPE_SEPARATE"], ) -@gluon.jit(repr=_pa_decode_sparse_repr) -def _pa_decode_sparse( +@gluon.jit(repr=_sparse_mla_repr) +def _sparse_mla( # Shapes below: C = queries, H = num_heads, S = HEAD_SIZE (the V width), # R = ROPE_DIM, nnz = total gathered tokens in a segment's index list. q_ptr, # [C, H, S (+R when ROPE_SEPARATE)] bf16 @@ -1734,9 +1729,6 @@ def _pa_decode_sparse( BLOCK_K: gl.constexpr, NUM_SPLITS: gl.constexpr, HEAD_ALIGNED: gl.constexpr, - MFMA_K: gl.constexpr, - GATHER_TW1: gl.constexpr, - LDS_PAD: gl.constexpr, # NOPE_CHUNK: extent of one dequant piece along CHUNK_AXIS (0 = rows, # 1 = columns); >= the tile's extent means one shot. NOPE_CHUNK: gl.constexpr, @@ -1751,14 +1743,13 @@ def _pa_decode_sparse( MAIN_SPLITS: gl.constexpr, # ADAPTIVE_SPLITS: re-decide the useful split count per query at runtime. ADAPTIVE_SPLITS: gl.constexpr, - ASM_DEQ: gl.constexpr, + DEQ: gl.constexpr, # see Fmt.DEQ # Per-cache buffer/global gate: buffer_load carries a 32-bit offset (2 GB), # and the two caches are sized independently. MAIN_USE_BUFFER_LOAD: gl.constexpr, EXTRA_USE_BUFFER_LOAD: gl.constexpr, IDX_BUFFER_LOAD: gl.constexpr, HAS_INVALID: gl.constexpr, - FP8_FNUZ: gl.constexpr, FP8_MFMA: gl.constexpr = False, # q already quantized to e4m3 by the caller, plus the scalar f32 scale it # was quantized with. This is the calling convention aiter's asm @@ -1773,7 +1764,6 @@ def _pa_decode_sparse( IDX_CACHE: gl.constexpr = _NO_CACHE, ASYNC_LDS: gl.constexpr = False, RELAXED_LOAD: gl.constexpr = True, - ROPE_VEC: gl.constexpr = 16, PAD_INTERVAL: gl.constexpr = 1024, ): """One program = (query, split, head-block). Two-loop: main (SWA) then @@ -1814,10 +1804,6 @@ def _pa_decode_sparse( "fp8_dsv4_mla/fp8_dsv32_mla rows carry a rope tail; ROPE_DIM=0 is " "inconsistent", ) - gl.static_assert( - (not ASM_DEQ) or MAIN_FMT == "fp8_dsv4_mla" or EXTRA_FMT == "fp8_dsv4_mla", - "ASM_DEQ is the dsv4 E8M0 dequant", - ) # The fp8 path needs one positive scalar scale per cache, since that is what # folds outside the loop, and OCP e4m3 code points, which is what the matrix # core reads. @@ -1826,9 +1812,9 @@ def _pa_decode_sparse( or (MAIN_FMT == "fp8_scalar" and (not HAS_EXTRA or EXTRA_FMT == "fp8_scalar")), "FP8_MFMA requires the per-tensor fp8 format on every segment", ) - gl.static_assert(not (FP8_MFMA and FP8_FNUZ), "FP8_MFMA is OCP e4m3 only") # The direct-to-LDS path stages raw code points, needs the clamped (branch-free) # tail, and carries no per-tile validity vector. + # TODO(has_invalid): the sentinel restriction is conservative, not fundamental. gl.static_assert( (not ASYNC_LDS) or (FP8_MFMA and UNI_TILE and not HAS_INVALID and not HAS_EXTRA), @@ -1855,10 +1841,7 @@ def _pa_decode_sparse( HEAD_SIZE, ROPE_DIM, ROPE_SEPARATE, - MFMA_K, NUM_WARPS, - GATHER_TW1, - LDS_PAD, UNI_TILE, HAS_INVALID, HEAD_ALIGNED, @@ -1868,7 +1851,6 @@ def _pa_decode_sparse( IDX_CACHE, ASYNC_LDS, RELAXED_LOAD, - ROPE_VEC, PAD_INTERVAL, ) main_fmt = Fmt( @@ -1876,8 +1858,7 @@ def _pa_decode_sparse( MAIN_FMT, MAIN_BLOCK_SIZE, MAIN_USE_BUFFER_LOAD, - FP8_FNUZ, - ASM_DEQ and MAIN_FMT == "fp8_dsv4_mla", + DEQ, NOPE_DIM, NOPE_CHUNK, CHUNK_AXIS, @@ -1887,8 +1868,7 @@ def _pa_decode_sparse( EXTRA_FMT, EXTRA_BLOCK_SIZE, EXTRA_USE_BUFFER_LOAD, - FP8_FNUZ, - ASM_DEQ and EXTRA_FMT == "fp8_dsv4_mla", + DEQ, NOPE_DIM, NOPE_CHUNK, CHUNK_AXIS, @@ -2215,14 +2195,14 @@ def _pa_decode_sparse( ) -_pa_decode_sparse_reduce_repr = make_kernel_repr( - "_pa_decode_sparse_reduce", +_sparse_mla_reduce_repr = make_kernel_repr( + "_sparse_mla_reduce", ["BLOCK_M", "HEAD_SIZE", "NUM_SPLITS"], ) -@gluon.jit(repr=_pa_decode_sparse_reduce_repr) -def _pa_decode_sparse_reduce( +@gluon.jit(repr=_sparse_mla_reduce_repr) +def _sparse_mla_reduce( part_m_ptr, part_l_ptr, part_acc_ptr, diff --git a/aiter/ops/triton/attention/pa_decode_sparse.py b/aiter/ops/triton/attention/pa_decode_sparse.py index 964564ee1e9..423060c670c 100644 --- a/aiter/ops/triton/attention/pa_decode_sparse.py +++ b/aiter/ops/triton/attention/pa_decode_sparse.py @@ -12,11 +12,11 @@ import torch import triton -from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( - _pa_decode_sparse as _pa_decode_sparse_gfx950, +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla as _sparse_mla_gfx950, ) -from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( - _pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950, +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla_reduce as _sparse_mla_reduce_gfx950, ) from aiter.ops.triton._gluon_kernels.gfx1250.attention.pa_decode_sparse import ( _pa_decode_sparse as gluon_pa_decode_sparse, @@ -34,14 +34,23 @@ from aiter.ops.triton.utils.common_utils import max_addressable_bytes from aiter.ops.triton.utils.device_info import get_num_sms from aiter.ops.triton.utils.logger import AiterTritonLogger +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype DEVICE_ARCH = arch_info.get_arch() +# Fused fp8 x E8M0 -> bf16 upcast check (gluon cdna4.scaled_upcast) +try: + from triton.experimental.gluon.language.amd import cdna4 as _cdna4 + + _HAS_SCALED_UPCAST = hasattr(_cdna4, "scaled_upcast") +except ImportError: + _HAS_SCALED_UPCAST = False + _LOGGER = AiterTritonLogger() _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() def _check_out(out, q, dtype): @@ -126,6 +135,10 @@ def pa_decode_sparse( raise RuntimeError("pa_decode_sparse requires CUDA/HIP tensors") if q.dtype not in (torch.bfloat16, torch.float16): raise RuntimeError(f"pa_decode_sparse expects fp16/bf16 q, got {q.dtype}") + _LOGGER.info( + f"PA_DECODE_SPARSE: q={tuple(q.shape)} unified_kv={tuple(unified_kv.shape)} " + f"{unified_kv.dtype} kv_indices={tuple(kv_indices.shape)}" + ) # gfx950: route to the merged DSv4 sparse-MLA gluon driver. Format is inferred # from the cache: 3D -> packed fp8_dsv4_mla / bf16 block cache (optional SWA+top-k @@ -133,21 +146,15 @@ def pa_decode_sparse( # kv_splits and skip_reduce are honored here; block_h and fp16 q fall through to # the triton path below (the gluon kernel is bf16-only: bf16 LDS + bf16 MFMA). if DEVICE_ARCH == "gfx950" and block_h is None and q.dtype == torch.bfloat16: + # gfx950 (CDNA4) reads OCP e4m3 natively. fnuz is the gfx942 encoding, + # so it never appears here and falls through to the triton path below. if unified_kv.ndim == 3: - _ok = kv_scales is None and ( - unified_kv.dtype == torch.uint8 or unified_kv.dtype == q.dtype - ) + # packed / bf16 block cache: it carries its own scales, if any + _ok = kv_scales is None and unified_kv.dtype in (torch.uint8, q.dtype) + elif kv_scales is not None: + _ok = unified_kv.dtype in (torch.float8_e4m3fn, torch.uint8) else: - _fp8 = unified_kv.dtype in ( - torch.float8_e4m3fn, - torch.float8_e4m3fnuz, - torch.uint8, - ) - _ok = (kv_scales is not None and _fp8) or ( - kv_scales is None and unified_kv.dtype == q.dtype - ) - # fnuz vs OCP e4m3 (2D fp8 only) selects the in-kernel dequant bias. - fp8_fnuz = unified_kv.ndim == 2 and unified_kv.dtype == torch.float8_e4m3fnuz + _ok = unified_kv.dtype == q.dtype if _ok: cache = ( unified_kv.view(torch.uint8) @@ -168,7 +175,6 @@ def pa_decode_sparse( kv_splits=kv_splits, skip_reduce=skip_reduce, has_invalid=bool(has_invalid), - fp8_fnuz=fp8_fnuz, out=out, ) @@ -202,9 +208,6 @@ def pa_decode_sparse( ) T, H, D = q.shape - _LOGGER.info( - f"PA_DECODE_SPARSE T={T} H={H} D={D} " f"total_indices={kv_indices.shape[0]}" - ) out = _check_out(out, q, q.dtype) assert kv_indices.dtype == torch.int32 and kv_indices.is_contiguous() @@ -465,7 +468,6 @@ def _pa_decode_sparse_gfx950_gluon( skip_reduce=False, out=None, has_invalid=False, - fp8_fnuz=False, ): """Merged gfx950 gluon DSv4 sparse-MLA decode driver. Format from cache.ndim: 3D [nb, block, 584] -> packed fp8_dsv4_mla (uint8: 448 NoPE fp8 e4m3 OCP + @@ -481,13 +483,8 @@ def _pa_decode_sparse_gfx950_gluon( # Tuned launch config (gfx950 / MI355). BLOCK_M = heads per MFMA M-tile, and 16 # is both the MFMA M and the DSv4 head count; BLOCK_K = KV tile; num_warps = # BLOCK_K // 16, because warps tile the dot-N and MFMA N = 16. - BLOCK_M, BLOCK_K, MFMA_K = 16, 64, 16 + BLOCK_M, BLOCK_K = 16, 64 num_warps = max(1, BLOCK_K // 16) - # Threads the NoPE gather spends on the head dim, 16 B each: 32 requests a whole - # 512 B token row per instruction instead of four 128 B quarters, which is what - # a scattered top-k gather wants. The cost is a longer per-lane slot vector. - gather_tw1 = 32 - lds_pad = 8 # bf16 elements of padding after each kv_smem row NOPE_DIM, ROPE_DIM = 448, 64 MAX_BYTES = 2**31 - 1 @@ -582,29 +579,12 @@ def _pa_decode_sparse_gfx950_gluon( heads_blocks = (num_heads + BLOCK_M - 1) // BLOCK_M out = _check_out(out, q, torch.bfloat16) - def _pick_splits(hb): - if kv_splits is not None: - return max(1, int(kv_splits)) - return _decode_num_splits_occ(num_queries, hb, avg_main, avg_extra, BLOCK_K) - - num_splits = _pick_splits(heads_blocks) - - tiles = max( - math.ceil(avg_main / BLOCK_K) if avg_main > 0 else 1, - math.ceil(avg_extra / BLOCK_K) if avg_extra > 0 else 1, - ) - if ( - num_heads % 8 == 0 - and tiles <= 2 - and num_queries * heads_blocks * num_splits <= get_num_sms() - ): - hb8 = (num_heads + 7) // 8 - ns8 = _pick_splits(hb8) - if ns8 >= num_splits and num_queries * hb8 * ns8 <= get_num_sms(): - BLOCK_M = 8 - HEAD_ALIGNED = num_heads % BLOCK_M == 0 - heads_blocks = hb8 - num_splits = ns8 + if kv_splits is not None: + num_splits = max(1, int(kv_splits)) + else: + num_splits = _decode_num_splits_occ( + num_queries, heads_blocks, avg_main, avg_extra, BLOCK_K + ) # Q is read once per query without split-K, and re-read by every split q_cache = ".cg" if num_splits == 1 else "" @@ -631,8 +611,11 @@ def _pick_splits(hb): part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # Dequant chunking - col_reps = head_dim // (gather_tw1 * 16) + # Dequant chunking. The gather layout puts 32 of a wave's 64 lanes along a row + # (16 B each, 512 B) and the other 32 on a second row; col_reps is how many + # such spans each lane holds, which is what makes a column split a free + # register rename. + col_reps = head_dim // 512 chunk_axis = 1 if col_reps >= 4 else 0 nope_chunk = max(1, BLOCK_K // 4) if chunk_axis == 0 else min(128, head_dim) @@ -651,13 +634,19 @@ def _pick_splits(hb): # than one split to give back. adaptive_splits = num_splits > 1 - # Fuse the fp8 x E8M0 dequant into v_cvt_scalef32_pk_bf16_fp8 via inline asm - asm_deq = one_wg_per_cu and not FLAT_POOL and not fp8_fnuz and main_is_fp8 + # Fuse the dsv4 dequant into v_cvt_scalef32_pk_bf16_fp8. The asm fallback + # gathers an extra int16 tile, so it only pays off at one workgroup per CU. + if _HAS_SCALED_UPCAST: + deq = "upcast" + elif one_wg_per_cu: + deq = "asm" + else: + deq = "none" # Grid dim 0 varies fastest and XCD assignment is round-robin over the linear # workgroup id, so the axis order decides what shares an XCD's L2. grid = (num_queries, num_splits, heads_blocks) - _pa_decode_sparse_gfx950[grid]( + _sparse_mla_gfx950[grid]( q, cache, main_bf16, @@ -706,9 +695,6 @@ def _pick_splits(hb): BLOCK_K=BLOCK_K, NUM_SPLITS=num_splits, HEAD_ALIGNED=HEAD_ALIGNED, - MFMA_K=MFMA_K, - GATHER_TW1=gather_tw1, - LDS_PAD=lds_pad, NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", @@ -720,12 +706,11 @@ def _pick_splits(hb): UNI_TILE=True, MAIN_SPLITS=main_splits, ADAPTIVE_SPLITS=adaptive_splits, - ASM_DEQ=asm_deq, + DEQ=deq, MAIN_USE_BUFFER_LOAD=main_use_buffer_load, EXTRA_USE_BUFFER_LOAD=extra_use_buffer_load, IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, - FP8_FNUZ=fp8_fnuz, num_warps=num_warps, waves_per_eu=waves_per_eu, ) @@ -737,7 +722,7 @@ def _pick_splits(hb): # One head per reduce workgroup rgrid = (num_queries, num_heads) - _pa_decode_sparse_reduce_gfx950[rgrid]( + _sparse_mla_reduce_gfx950[rgrid]( part_m, part_l, part_acc, diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index acbbacbb30b..71667761562 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -1,23 +1,15 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Sparse MLA attention (gfx950 gluon): ``kv_lora_rank`` latent + appended -decoupled rope, token-granular top-k gather. - -Prefill and decode are the same MQA operator on this path, one program per query -token over that token's own gathered KV, so this serves both phases; the launcher -picks the tile from what the grid supplies. -""" - import math import torch -from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( - _pa_decode_sparse as _pa_decode_sparse_gfx950, +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla as _sparse_mla_gfx950, ) -from aiter.ops.triton._gluon_kernels.gfx950.attention.pa_decode_sparse import ( - _pa_decode_sparse_reduce as _pa_decode_sparse_reduce_gfx950, +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_mla import ( + _sparse_mla_reduce as _sparse_mla_reduce_gfx950, ) from aiter.ops.triton.attention.pa_decode_sparse import ( _as_int32_contiguous_1d, @@ -29,18 +21,6 @@ _LOGGER = AiterTritonLogger() -# buffer_load carries a signed 32-bit offset; caches past this span gather -# through 64-bit addresses -MAX_BYTES = 2**31 - 1 - -# Tuned launch config (gfx950 / MI355). num_warps = BLOCK_K // 16 (warps tile -# the dot-N, MFMA N=16); GATHER_TW1=32 requests a whole 512 B token row per -# gather instruction. -_BLOCK_K = 64 -_MFMA_K = 16 -_GATHER_TW1 = 32 -_LDS_PAD = 8 - def _check_out(out, q, kv_lora_rank): """Caller-supplied output buffer, or a fresh one. A buffer wider than @@ -56,55 +36,32 @@ def _check_out(out, q, kv_lora_rank): return out -def _infer_cache_format(kv, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale): - """-> (fmt, cache, alt_ptr, scl_ptr, block_size, fp8_fnuz). Pointer roles - follow the kernel's Seg contract (see its docstring).""" +def _cache_pointers(fmt, kv, d_qk, kv_scale): + """-> (cache, alt_ptr, scl_ptr, block_size) for an MLA-route format. Pointer + roles follow the kernel's Seg contract (see its docstring).""" if kv.ndim == 4: # [slots, 1, 1, R], the asm mla_decode_fwd view - assert kv.shape[1] == 1 and kv.shape[2] == 1 kv = kv.reshape(kv.shape[0], kv.shape[3]) elif kv.ndim == 3 and kv.shape[2] == d_qk: # vLLM paged cache [nb, block_size, R]: indices are global slot ids and # rows are contiguous, so the flat view is stride-identical. - assert kv.stride(2) == 1 and kv.stride(1) == kv.shape[2] * kv.stride(2) + if not (kv.stride(2) == 1 and kv.stride(1) == kv.shape[2]): + raise ValueError("a [nb, block, R] pool must have contiguous rows") kv = kv.reshape(-1, d_qk) - - fp8_dtypes = (torch.uint8, torch.float8_e4m3fn, torch.float8_e4m3fnuz) - if kv.ndim == 2 and kv.shape[1] == d_qk: - if kv.dtype == torch.bfloat16: - return "bf16", kv, kv, None, 1, False - assert kv.dtype in fp8_dtypes, f"unsupported flat cache dtype {kv.dtype}" - assert ( - kv_scale is not None - ), "flat fp8 cache needs the per-tensor kv_scale (layer._k_scale)" - assert kv_scale.dtype == torch.float32 + if fmt == "bf16": + return kv, kv, None, 1 + if fmt == "fp8_scalar": u8 = kv.view(torch.uint8) - fnuz = kv.dtype == torch.float8_e4m3fnuz - return "fp8_scalar", u8, u8, kv_scale.reshape(1), 1, fnuz - if ( - 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 - ): - # fp8_dsv32_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B. - # This is vLLM's fp8_ds_mla on V3.2 / Kimi-K3; the same vLLM name also - # covers V4's 584 B layout (fp8_dsv4_mla), hence the explicit generation. - u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) - assert ( - u8.stride(2) == 1 and u8.stride(1) == u8.shape[2] - ), "fp8_dsv32_mla rows must be contiguous 656-byte records" - return ( - "fp8_dsv32_mla", - u8, - u8.view(torch.bfloat16), - u8.view(torch.float32), - u8.shape[1], - False, - ) - raise ValueError(f"unrecognized sparse-MLA kv cache: {tuple(kv.shape)} {kv.dtype}") + return u8, u8, kv_scale.reshape(1), 1 + # fp8_dsv32_mla: 512 fp8 | 4 f32 per-128 scales | 64 bf16 rope = 656 B. + u8 = kv if kv.dtype == torch.uint8 else kv.view(torch.uint8) + if not (u8.stride(2) == 1 and u8.stride(1) == u8.shape[2]): + raise ValueError("fp8_dsv32_mla rows must be contiguous records") + return u8, u8.view(torch.bfloat16), u8.view(torch.float32), u8.shape[1] -def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int: +def _mla_num_splits( + num_queries: int, heads_blocks: int, avg_topk: float, block_k: int = 64 +) -> int: """Split-K count for the sparse-MLA decode. Below one workgroup per CU, split to fill the machine but never past 8. @@ -112,18 +69,12 @@ def _mla_num_splits(num_queries: int, heads_blocks: int, avg_topk: float) -> int num_sms = get_num_sms() base_wg = max(1, num_queries * heads_blocks) cta_cap = max(1, (2 * num_sms) // base_wg) - tiles = max(1, math.ceil(avg_topk / _BLOCK_K)) + tiles = max(1, math.ceil(avg_topk / block_k)) if base_wg >= num_sms: return max(1, min(cta_cap, tiles // 16)) return max(1, min(cta_cap, tiles, 8)) -# Direct-to-LDS tile pipeline; wins wherever the launch gives a CU more than one -# workgroup to hide the copy behind. -_ASYNC_LDS_DEFAULT = True -_ASYNC_ROPE_VEC = 16 # bytes/lane in the rope copy - - def _async_launch_config( fp8_dots: bool, has_invalid: bool, @@ -134,6 +85,7 @@ def _async_launch_config( use_buffer_load: bool, uni_tile: bool = True, has_extra: bool = False, + block_k: int = 64, ) -> tuple[bool, int, int]: """-> (ASYNC_LDS, BLOCK_K, waves_per_eu) for this launch. @@ -143,8 +95,14 @@ def _async_launch_config( workgroups/CU) the large tile wins instead, by halving the cross-warp softmax exchange rate per token. """ - enabled = _ASYNC_LDS_DEFAULT - enabled = enabled and fp8_dots and uni_tile and not has_invalid and not has_extra + ASYNC_LDS_DEFAULT = True + enabled = ( + ASYNC_LDS_DEFAULT + and fp8_dots + and uni_tile + and not has_invalid + and not has_extra + ) workgroups = num_queries * heads_blocks * max(1, num_splits) num_sms = get_num_sms() if enabled and workgroups >= 4 * num_sms: @@ -152,13 +110,10 @@ def _async_launch_config( waves_per_eu = 2 if use_buffer_load and workgroups <= num_sms: waves_per_eu = 1 - return enabled, (128 if enabled else _BLOCK_K), waves_per_eu - + return enabled, (128 if enabled else block_k), waves_per_eu -_E4M3_MAX = 448.0 - -def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool: +def _resolve_dot_precision(dot_precision: str, fmt: str) -> bool: if dot_precision not in ("bf16", "fp8"): raise ValueError( f"dot_precision must be 'bf16' or 'fp8', got {dot_precision!r}" @@ -172,17 +127,208 @@ def _resolve_dot_precision(dot_precision: str, fmt: str, fp8_fnuz: bool) -> bool ) if fmt == "bf16": raise ValueError("dot_precision='fp8' needs an fp8 cache.") - if fp8_fnuz: - raise ValueError("dot_precision='fp8' is OCP e4m3 only.") return True +_DSV4_ROW = 448 + 2 * 64 + 8 # 584 B: fp8 nope | bf16 rope | 8 B UE8M0 trailer + + +def _dsv32_row(kv_lora_rank, qk_rope_head_dim): + return kv_lora_rank + 4 * (kv_lora_rank // 128) + 2 * qk_rope_head_dim + + +def _classify_flat(kv, width, slots, kv_scale, what): + """Flat pool rows are one QK row per slot; dtype and kv_scale pick the tag.""" + if kv.dtype == torch.bfloat16: + return "bf16" # kv_scale, if any, is ignored + if kv.dtype == torch.float8_e4m3fnuz: + raise ValueError( + f"{what}: float8_e4m3fnuz is the gfx942 encoding; gfx950 reads OCP e4m3" + ) + if kv.element_size() != 1: + raise ValueError(f"{what}: unsupported cache dtype {kv.dtype}") + if kv_scale is None: + raise ValueError( + f"{what}: a flat fp8 cache needs kv_scale, [1] f32 (fp8_scalar) or " + f"[slots, D // 64] f32 (fp8_g64)" + ) + if kv_scale.dtype != torch.float32: + raise ValueError(f"{what}: kv_scale must be f32, got {kv_scale.dtype}") + if kv_scale.numel() == 1: + return "fp8_scalar" + if width % 64 or tuple(kv_scale.shape) != (slots, width // 64): + raise ValueError( + f"{what}: kv_scale {tuple(kv_scale.shape)} is neither [1] (fp8_scalar) " + f"nor [slots, D // 64] = [{slots}, {width // 64}] (fp8_g64)" + ) + return "fp8_g64" + + +def _classify_cache(q, kv, kv_lora_rank, qk_rope_head_dim, kv_scale, what="kv_buffer"): + """kv -> one of the kernel's format tags, or raise. + + Record width, dtype and the kv_scale shape are the only things that tell the + formats apart, so whatever they do not pin down is rejected here instead of + being decoded as the wrong layout. Host metadata only: safe under graph + capture. + """ + d_qk = q.shape[-1] + if kv.device != q.device: + raise ValueError(f"{what} is on {kv.device}, q is on {q.device}") + if kv.ndim == 4: + if kv.shape[1] != 1 or kv.shape[2] != 1: + raise ValueError( + f"{what}: a 4-D cache must be [slots, 1, 1, R], got {tuple(kv.shape)}" + ) + if kv.ndim in (2, 4): + width = kv.shape[-1] + if width != d_qk: + raise ValueError( + f"{what}: flat cache rows are {width} wide but q is {d_qk}; a flat " + f"pool stores one QK row (kv_lora_rank + qk_rope_head_dim) per slot" + ) + return _classify_flat(kv, width, kv.shape[0], kv_scale, what) + if kv.ndim == 3: + nb, block, width = kv.shape + if width == d_qk: + # [nb, block, R]: a flat pool stored in blocks. Same rules, except the + # per-64 scale vector is only defined for the 2-D pool. + fmt = _classify_flat(kv, width, nb * block, kv_scale, what) + if fmt == "fp8_g64": + raise ValueError(f"{what}: fp8_g64 needs a 2-D [slots, D] pool") + return fmt + if kv_scale is not None: + raise ValueError( + f"{what}: packed caches carry their own scales; kv_scale must be None" + ) + one_byte = kv.element_size() == 1 + if one_byte and width == _DSV4_ROW: + return "fp8_dsv4_mla" + dsv32 = _dsv32_row(kv_lora_rank, qk_rope_head_dim) + if one_byte and qk_rope_head_dim > 0 and width == dsv32: + return "fp8_dsv32_mla" + raise ValueError( + f"{what}: unrecognized 3-D cache, {width}-wide {kv.dtype} records for q " + f"width {d_qk}. Expected {d_qk} (bf16 or fp8_scalar block pool), " + f"{_DSV4_ROW} uint8 (fp8_dsv4_mla) or {dsv32} uint8 (fp8_dsv32_mla with " + f"kv_lora_rank={kv_lora_rank}, qk_rope_head_dim={qk_rope_head_dim})" + ) + raise ValueError(f"{what} must be 2-D, 3-D or 4-D, got {kv.ndim}-D") + + +def _check_geometry(fmt, d_qk, kv_lora_rank, qk_rope_head_dim): + if kv_lora_rank <= 0 or qk_rope_head_dim < 0: + raise ValueError( + f"kv_lora_rank must be > 0 and qk_rope_head_dim >= 0, got " + f"{kv_lora_rank} / {qk_rope_head_dim}" + ) + if fmt == "fp8_dsv4_mla": + if d_qk != 512: + raise ValueError( + f"fp8_dsv4_mla needs q width 512 (448 nope + 64 rope), got {d_qk}" + ) + return + if fmt == "fp8_g64": + return # the whole row is the head; geometry args are not read + if d_qk != kv_lora_rank + qk_rope_head_dim: + hint = ( + " A bf16 or fp8_scalar cache carries no rope information: pass " + "qk_rope_head_dim=0 for rope-inside-the-row or rope-free models." + if fmt in ("bf16", "fp8_scalar") + else "" + ) + raise ValueError( + f"q width {d_qk} != kv_lora_rank {kv_lora_rank} + qk_rope_head_dim " + f"{qk_rope_head_dim}.{hint}" + ) + if fmt == "fp8_dsv32_mla" and kv_lora_rank % 128: + raise ValueError( + f"fp8_dsv32_mla stores one scale per 128 latent columns; kv_lora_rank " + f"{kv_lora_rank} is not a multiple of 128" + ) + + +def _check_index_stream(indptr, indices, num_queries, device, what=""): + """Both are flat integer tensors, so only the length of indptr can catch a + swapped pair. No device reads.""" + for name, t in ((f"{what}kv_indptr", indptr), (f"{what}kv_indices", indices)): + if not torch.is_tensor(t): + raise ValueError(f"{name} must be a tensor, got {type(t).__name__}") + if t.dtype not in (torch.int32, torch.int64): + raise ValueError(f"{name} must be int32 (int64 accepted), got {t.dtype}") + if t.ndim != 1: + raise ValueError(f"{name} must be 1-D, got {tuple(t.shape)}") + if t.device != device: + raise ValueError(f"{name} is on {t.device}, q is on {device}") + if indptr.numel() != num_queries + 1: + raise ValueError( + f"{what}kv_indptr must have C + 1 = {num_queries + 1} entries, got " + f"{indptr.numel()}. kv_indptr and kv_indices are both flat int32; " + "check their order." + ) + + +def _forward_paged( + q, + kv, + kv_indptr, + kv_indices, + softmax_scale, + kv_scale, + kv_splits, + skip_reduce, + has_invalid, + dot_precision, + q_scale, + out, + return_lse, + attn_sink, + extra_kv, + extra_indptr, + extra_indices, +): + """dsv4 and the SWA+top-k two-loop, until the two launchers merge.""" + from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse + + unsupported = [ + name + for name, asked in ( + ("dot_precision='fp8'", dot_precision != "bf16"), + ("return_lse", return_lse), + ) + if asked + ] + if unsupported: + raise ValueError( + f"{', '.join(unsupported)} is not supported for dsv4 / two-loop " + "caches yet." + ) + res = pa_decode_sparse( + q, + kv, + kv_indices, + kv_indptr, + attn_sink, + softmax_scale, + kv_scales=kv_scale, + kv_splits=kv_splits, + has_invalid=has_invalid, + skip_reduce=skip_reduce, + extra_cache=extra_kv, + extra_indices=extra_indices, + extra_indptr=extra_indptr, + out=out, + ) + return res if isinstance(res, tuple) else (res, None) + + def sparse_mla_fwd( q: torch.Tensor, kv_buffer: torch.Tensor, kv_indptr: torch.Tensor, kv_indices: torch.Tensor, softmax_scale: float, + *, kv_scale: torch.Tensor | None = None, kv_lora_rank: int = 512, qk_rope_head_dim: int = 64, @@ -193,55 +339,149 @@ def sparse_mla_fwd( q_scale: torch.Tensor | None = None, out: torch.Tensor | None = None, return_lse: bool = False, + attn_sink: torch.Tensor | None = None, + extra_kv: torch.Tensor | None = None, + extra_indptr: torch.Tensor | None = None, + extra_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Sparse (top-k gathered) MLA attention. - ``qk_rope_head_dim`` selects the geometry: separated rope (DeepSeek-V3.2, - GLM-5.1/5.2), where the query is the latent plus an appended rope, or - rope-free (GLM-5.3-Flash), where the query is the latent alone. + Everything after softmax_scale is keyword-only: kv_indptr and kv_indices are + both flat int32, so a positional mix-up would type-check and return garbage. + + Supported KV cache formats. The format is inferred from kv_buffer's shape, + dtype and kv_scale; each row gives what the caller has to pass. R is the + QK width, kv_lora_rank + qk_rope_head_dim. + + format kv_buffer kv_scale geometry args + bf16 [slots, R], [nb, block, R], None as the model + [slots, 1, 1, R]; bf16 + fp8_scalar same shapes, fp8 or uint8 [1] f32 as the model + (GLM-5.x; the only format + that runs dot_precision="fp8") + fp8_g64 [slots, D] fp8 or uint8 [slots, D//64] kv_lora_rank=D, + (uniform pool, rope inside) f32 qk_rope_head_dim=0 + fp8_dsv32_mla [nb, block, 656] uint8 None 512 / 64 + (DeepSeek-V3.2, Kimi-K3; vLLM + fp8_ds_mla: 512 fp8 | 4 f32 + per-128 | 64 bf16 rope) + fp8_dsv4_mla [nb, block, 584] uint8 None ignored + (DeepSeek-V4; vLLM fp8_ds_mla: + 448 fp8 | 64 bf16 rope + 8 B + UE8M0 trailer per block) + + Geometry, set by qk_rope_head_dim: separated rope (DeepSeek-V3.2, + GLM-5.1/5.2), where the query is the latent plus an appended rope and V is + the latent; rope-free (GLM-5.3-Flash), qk_rope_head_dim=0, where the query + is the latent alone. DeepSeek-V4 keeps its rope inside the 512-wide row + (448 nope + 64 rope) with V the whole row; that layout is fixed by the cache + format, so the geometry args (kv_lora_rank and qk_rope_head_dim) are not + read for it. A bf16 pool carries no rope information, so a rope-inside or + rope-free model on a bf16 cache must pass qk_rope_head_dim=0. Args: q: [C, H, kv_lora_rank + qk_rope_head_dim] queries, one row per query token (prefill and decode alike). qk_rope_head_dim: width of the appended rope. 0 means rope-free, and the QK contraction becomes a single dot over kv_lora_rank. - kv_buffer: the KV pool [nb, block, R] / [slots, 1, 1, R] / - [slots, R] in bf16, the same shapes in fp8 (+ scalar - kv_scale), or [nb, block, 656] uint8 (``fp8_dsv32_mla``, which is - vLLM's ``fp8_ds_mla`` on DeepSeek-V3.2 / Kimi-K3). - kv_indptr: ``[C + 1]`` int32 prefix sum of per-query index counts. + kv_buffer: the KV pool, one of the formats in the table above. + kv_indptr: [C + 1] int32 prefix sum of per-query index counts. kv_indices: flat int32 GLOBAL slot ids into the pool. softmax_scale: the layer's softmax scale. - kv_scale: ``[1]`` f32 per-tensor cache scale; - required for the flat fp8 format. + kv_scale: f32 cache scale, shape per the table; the packed formats + carry their own scales and take None. kv_splits: split-K override; default follows the occupancy policy. - skip_reduce: with split-K active, return ``(part_acc, part_m, part_l)`` + skip_reduce: with split-K active, return (part_acc, part_m, part_l) instead of launching the combine. has_invalid: index stream carries -1 sentinels (masked out). Default False. + attn_sink: optional [H] f32 per-head sink: exp(sink) joins the softmax + denominator (and the LSE) as a virtual key. None means no sink, + which is not the same as a zero sink. + extra_kv, extra_indptr, extra_indices: a second segment attended in the + same pass, for the SWA-window + top-k two-loop. All three together. dot_precision: what the QK and PV matrix-core ops run in. "bf16" (default): the KV tile is dequantized to bf16 on its way into LDS and both dots are bf16. Works with every cache format. "fp8": the cache's own code points go to the fp8 matrix core with no dequant, and the per-tensor scale folds outside the tile loop. + tensor scale fp8 kv cache only. - q is adapted to the choice. bf16 q is quantized here when "fp8" is - asked for, fp8 q is passed straight through, and fp8 q under "bf16" - dots is widened in-kernel, which is exact. - q_scale: scalar f32. Required when q arrives already fp8 (the scale it - was quantized with; the aiter asm convention. Optional when q is bf16 and - dot_precision="fp8". - out: optional ``[C, H, >= kv_lora_rank]`` bf16 destination. + q is adapted to the choice. bf16 q is quantized in the kernel + prologue, one scale per (query, head-block) tile; fp8 q is passed + straight through, and fp8 q under "bf16" dots is widened in-kernel. + + q_scale: scalar f32, required when q arrives already fp8 (the scale it + was quantized with; the aiter asm convention). Ignored for bf16 q, + which the kernel quantizes itself when dot_precision="fp8". + out: optional [C, H, >= kv_lora_rank] bf16 destination. return_lse: also return the natural-log log-sum-exp, [C, H] f32, for merging partials across context-parallel ranks. A fully masked row reports -inf. Returns: - (out, lse), following aiter.mla.mla_decode_fwd: out is + (out, lse), out is [C, H, kv_lora_rank] bf16 (the latent V), lse is None unless return_lse. """ - assert q.ndim == 3, f"expected q=[b,h,d], got {q.shape}" + if q.ndim != 3: + raise ValueError(f"expected q=[C, H, d_qk], got {tuple(q.shape)}") + num_queries, num_heads, d_qk = q.shape + fmt = _classify_cache(q, kv_buffer, kv_lora_rank, qk_rope_head_dim, kv_scale) + _LOGGER.info( + f"SPARSE_MLA_FWD: q={tuple(q.shape)} kv_buffer={tuple(kv_buffer.shape)} " + f"{kv_buffer.dtype} kv_indices={tuple(kv_indices.shape)} fmt={fmt}" + ) + _check_geometry(fmt, d_qk, kv_lora_rank, qk_rope_head_dim) + _check_index_stream(kv_indptr, kv_indices, num_queries, q.device) + if attn_sink is not None and attn_sink.numel() != num_heads: + raise ValueError( + f"attn_sink must have one entry per head ({num_heads}), got " + f"{tuple(attn_sink.shape)}" + ) + extra = (extra_kv, extra_indptr, extra_indices) + has_extra = all(a is not None for a in extra) + if any(a is not None for a in extra) and not has_extra: + raise ValueError("extra_kv, extra_indptr and extra_indices go together") + if has_extra: + # The SWA-window + top-k two-loop: both segments are block caches the + # dsv4 decoder reads. A flat main pool would silently drop the extra one. + if kv_buffer.ndim != 3 or fmt not in ("fp8_dsv4_mla", "bf16"): + raise ValueError( + "the two-loop needs a 3-D block cache (fp8_dsv4_mla or bf16) as " + f"kv_buffer, got {fmt}" + ) + extra_fmt = _classify_cache( + q, extra_kv, kv_lora_rank, qk_rope_head_dim, None, what="extra_kv" + ) + if extra_kv.ndim != 3 or extra_fmt not in ("fp8_dsv4_mla", "bf16"): + raise ValueError( + f"extra_kv must be a 3-D fp8_dsv4_mla or bf16 block cache, got " + f"{extra_fmt}" + ) + _check_index_stream( + extra_indptr, extra_indices, num_queries, q.device, what="extra_" + ) + if fmt in ("fp8_dsv4_mla", "fp8_g64") or has_extra: + return _forward_paged( + q, + kv_buffer, + kv_indptr, + kv_indices, + softmax_scale, + kv_scale, + kv_splits, + skip_reduce, + has_invalid, + dot_precision, + q_scale, + out, + return_lse, + attn_sink, + extra_kv, + extra_indptr, + extra_indices, + ) assert arch_info.get_arch() == "gfx950", "sparse_mla_fwd is gfx950-only" q_is_fp8 = q.dtype == torch.float8_e4m3fn if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): @@ -254,60 +494,36 @@ def sparse_mla_fwd( ) ) 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. + # Caller-quantized q, tensor scale if q_scale is None: raise ValueError("fp8 q needs q_scale (the scale it was quantized with)") if q_scale.numel() != 1: - raise ValueError(f"q_scale must be a scalar, got {tuple(q_scale.shape)}") - q_scale = q_scale.reshape(1).to(torch.float32).contiguous() - num_queries, num_heads, d_qk = q.shape - assert ( - d_qk == kv_lora_rank + qk_rope_head_dim - ), f"q last dim {d_qk} != {kv_lora_rank} + {qk_rope_head_dim}" - assert qk_rope_head_dim >= 0, "qk_rope_head_dim cannot be negative" - _LOGGER.info( - f"SPARSE_MLA C={num_queries} H={num_heads} d_qk={d_qk} " - f"nnz={kv_indices.shape[0]}" - ) - - fmt, cache, alt, scl, block_size, fp8_fnuz = _infer_cache_format( - kv_buffer, d_qk, kv_lora_rank, qk_rope_head_dim, kv_scale - ) - fp8_dots = _resolve_dot_precision(dot_precision, fmt, fp8_fnuz) - if fp8_dots and not q_is_fp8: - # Quantize the way production does, one scaled_fp8_quant over - # [C, H*d_qk]. A caller-supplied q_scale is the layer's calibrated - # static scale; without one, take a per-tensor amax. - if q_scale is None: - q_scale = ( - (q.detach().float().abs().amax() / _E4M3_MAX) - .clamp_min(torch.finfo(torch.float32).tiny) - .reshape(1) + raise ValueError( + f"q_scale must be a single per-tensor value, got " + f"{tuple(q_scale.shape)}; per-head q scales are not supported" ) - q = ( - (q.float() / q_scale) - .clamp(-_E4M3_MAX, _E4M3_MAX) - .to(torch.float8_e4m3fn) - .contiguous() - ) - q_is_fp8 = True - elif q_scale is not None and not q_is_fp8: - raise ValueError( - "q_scale was given but q is bf16 and dot_precision='bf16', so " - "nothing would use it" - ) + q_scale = q_scale.reshape(1).to(torch.float32).contiguous() + cache, alt, scl, block_size = _cache_pointers(fmt, kv_buffer, d_qk, kv_scale) + fp8_dots = _resolve_dot_precision(dot_precision, fmt) + if not q_is_fp8: + # q_scale describes an fp8 q's encoding. With bf16 q the kernel quantizes + # per (query, head-block) tile when the dots are fp8, so a caller-supplied + # scale has nothing to apply to; callers pass layer._q_scale regardless. + q_scale = None kv_indices = _as_int32_contiguous_1d(kv_indices) kv_indptr = _as_int32_contiguous_1d(kv_indptr) - assert kv_indptr.numel() == num_queries + 1 - # No attention sink in MLA models; the kernel still wants a live pointer - # for the compile-time-elided argument slot. - attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) + has_sink = attn_sink is not None + if has_sink: + attn_sink = attn_sink.reshape(-1).to(torch.float32).contiguous() + else: + # The kernel still wants a live pointer for the compile-time-elided slot. + attn_sink = torch.empty(1, device=q.device, dtype=torch.float32) - # H < 16 runs natively at BLOCK_M = next_pow2(H) instead of padding heads. + # Tuned launch config (gfx950 / MI355). H < 16 runs natively at + # BLOCK_M = next_pow2(H) instead of padding heads block_m = 16 if num_heads >= 16 else max(8, 1 << (num_heads - 1).bit_length()) - num_warps = _BLOCK_K // 16 + block_k = 64 + num_warps = block_k // 16 num_rows = cache.shape[0] * block_size if cache.ndim >= 2 else cache.shape[0] avg_topk = kv_indices.numel() / max(1, num_queries) @@ -318,7 +534,8 @@ def sparse_mla_fwd( if int(cache.stride(0)) % a == 0: cs0_align = a break - + # buffer_load carries a signed 32-bit offset + MAX_BYTES = 2**31 - 1 use_buffer_load = max_addressable_bytes(cache) < MAX_BYTES idx_use_buffer_load = max_addressable_bytes(kv_indices) < MAX_BYTES @@ -337,7 +554,7 @@ def sparse_mla_fwd( if kv_splits is not None: num_splits = max(1, int(kv_splits)) else: - num_splits = _mla_num_splits(num_queries, heads_blocks, avg_topk) + num_splits = _mla_num_splits(num_queries, heads_blocks, avg_topk, block_k) if num_splits > 1: part_m = torch.empty( @@ -361,10 +578,13 @@ def sparse_mla_fwd( part_m = part_l = part_acc = out # unused placeholders (never dereferenced) pm_stride0 = pm_stride_s = pa_stride0 = pa_stride_s = pa_stride_h = 0 - # Dequant chunking - col_reps = kv_lora_rank // (_GATHER_TW1 * 16) + # Dequant chunking. The gather layout puts 32 of a wave's 64 lanes along a row + # (16 B each, 512 B) and the other 32 on a second row; col_reps is how many + # such spans each lane holds, which is what makes a column split a free + # register rename. + col_reps = kv_lora_rank // 512 chunk_axis = 1 if col_reps >= 4 else 0 - nope_chunk = max(1, _BLOCK_K // 4) if chunk_axis == 0 else min(128, kv_lora_rank) + nope_chunk = max(1, block_k // 4) if chunk_axis == 0 else min(128, kv_lora_rank) async_lds_on, block_k, waves_per_eu = _async_launch_config( fp8_dots, has_invalid, @@ -375,14 +595,13 @@ def sparse_mla_fwd( use_buffer_load, uni_tile=True, has_extra=False, + block_k=block_k, ) - # num_warps stays 4: warps tile the dot's N (MFMA N=16), so a larger BLOCK_K just - # gives each warp two N tiles. # Q is read once per query without split-K, and re-read by every split q_cache = ".cg" if num_splits == 1 else "" grid = (num_queries, num_splits, heads_blocks) - _pa_decode_sparse_gfx950[grid]( + _sparse_mla_gfx950[grid]( q, cache, alt, @@ -398,7 +617,7 @@ def sparse_mla_fwd( part_l, part_acc, scl, # f32 side-channel: k_scale ("fp8_scalar") / f32 view ("fp8_dsv32_mla") - scl, + scl, # extra segment's slot, unread (HAS_EXTRA=False) float(softmax_scale), q.stride(0), q.stride(1), @@ -415,7 +634,7 @@ def sparse_mla_fwd( pa_stride_h, num_heads, HAS_EXTRA=False, - HAS_SINK=False, + HAS_SINK=has_sink, MAIN_FMT=fmt, EXTRA_FMT=fmt, MAIN_BLOCK_SIZE=block_size, @@ -429,9 +648,6 @@ def sparse_mla_fwd( BLOCK_K=block_k, NUM_SPLITS=num_splits, HEAD_ALIGNED=head_aligned, - MFMA_K=_MFMA_K, - GATHER_TW1=_GATHER_TW1, - LDS_PAD=(_LDS_PAD * 2 if fp8_dots else _LDS_PAD), NOPE_CHUNK=nope_chunk, CHUNK_AXIS=chunk_axis, PART_STORE_CACHE="", @@ -440,18 +656,13 @@ def sparse_mla_fwd( Q_CACHE=q_cache, MAIN_SPLITS=num_splits, ADAPTIVE_SPLITS=num_splits > 1, - ASM_DEQ=False, + DEQ="none", MAIN_USE_BUFFER_LOAD=use_buffer_load, EXTRA_USE_BUFFER_LOAD=use_buffer_load, IDX_BUFFER_LOAD=idx_use_buffer_load, HAS_INVALID=has_invalid, - FP8_FNUZ=fp8_fnuz, FP8_MFMA=fp8_dots, ASYNC_LDS=async_lds_on, - ROPE_VEC=_ASYNC_ROPE_VEC, - # A 576 B row is not 128 B aligned: it spans five lines and shares its end - # line with its own rope read. ".cg" (sc0 nt) marks that evict-first and - # throws the sharing away, so this gather wants a plain cached load. GATHER_CACHE="", q_scl_ptr=q_scale, Q_FP8=q_is_fp8, @@ -468,7 +679,7 @@ def sparse_mla_fwd( # One head per reduce workgroup rgrid = (num_queries, num_heads) - _pa_decode_sparse_reduce_gfx950[rgrid]( + _sparse_mla_reduce_gfx950[rgrid]( part_m, part_l, part_acc, @@ -482,7 +693,7 @@ def sparse_mla_fwd( pa_stride_s, pa_stride_h, num_heads, - HAS_SINK=False, + HAS_SINK=has_sink, HEAD_SIZE=kv_lora_rank, BLOCK_M=1, NUM_SPLITS=num_splits, diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py index 736af74e19c..2150e0d6701 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_mla.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -31,8 +31,8 @@ E4M3_MAX = 448.0 # Both kernels the launch runs are named after the gluon kernel, the combine as -# _pa_decode_sparse_reduce, so one substring picks up both and nothing else. -KERNEL_MATCH = "_pa_decode_sparse" +# _sparse_mla_reduce, so one substring picks up both and nothing else. +KERNEL_MATCH = "_sparse_mla" FLUSH_BYTES = 512 << 20 _flush_buf = None diff --git a/op_tests/triton_tests/attention/test_pa_decode_sparse.py b/op_tests/triton_tests/attention/test_pa_decode_sparse.py index 72bdd8e7191..b4f89a6836b 100644 --- a/op_tests/triton_tests/attention/test_pa_decode_sparse.py +++ b/op_tests/triton_tests/attention/test_pa_decode_sparse.py @@ -9,6 +9,7 @@ from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import checkAllclose @@ -311,18 +312,14 @@ def test_pa_decode_sparse_vs_reference( ) -# --------------------------------------------------------------------------- -# FP8 KV cache quantization helpers -# --------------------------------------------------------------------------- - _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. - Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + Returns (kv_fp8, kv_scales) where kv_fp8 is _FP8_DTYPE and kv_scales is [total_pages, D // group_size] fp32. """ total_pages, D = unified_kv.shape diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index b6198c1b6b5..008426d38da 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -193,7 +193,7 @@ def test_global_load_path(): q, cache, ks, idx, ptr, truth = _build("tensor", 8, 16, 2048, live, ragged=False) big = torch.empty(3_800_000, D_QK, dtype=cache.dtype, device=cache.device) big[:live] = cache - assert smd.max_addressable_bytes(big) >= smd.MAX_BYTES + assert smd.max_addressable_bytes(big) >= 2**31 - 1 # past buffer_load's offset out, _ = sparse_mla_fwd(q, big, ptr, idx, sm, kv_scale=ks) e = rel_err(out, reference(q, truth.to(torch.bfloat16), idx, ptr, sm)) assert e < 2e-2, f"global load path: rel-err {e:.3e}" From 3a38c85613f99e4a52863a7414b75df72a5cb9f8 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 15 Sep 2026 14:50:20 +0000 Subject: [PATCH 48/51] formatting --- aiter/ops/triton/attention/sparse_mla.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_mla.py b/aiter/ops/triton/attention/sparse_mla.py index 71667761562..78ae6294d56 100644 --- a/aiter/ops/triton/attention/sparse_mla.py +++ b/aiter/ops/triton/attention/sparse_mla.py @@ -175,11 +175,10 @@ def _classify_cache(q, kv, kv_lora_rank, qk_rope_head_dim, kv_scale, what="kv_bu d_qk = q.shape[-1] if kv.device != q.device: raise ValueError(f"{what} is on {kv.device}, q is on {q.device}") - if kv.ndim == 4: - if kv.shape[1] != 1 or kv.shape[2] != 1: - raise ValueError( - f"{what}: a 4-D cache must be [slots, 1, 1, R], got {tuple(kv.shape)}" - ) + if kv.ndim == 4 and (kv.shape[1] != 1 or kv.shape[2] != 1): + raise ValueError( + f"{what}: a 4-D cache must be [slots, 1, 1, R], got {tuple(kv.shape)}" + ) if kv.ndim in (2, 4): width = kv.shape[-1] if width != d_qk: From 33c58f39a99bc858373a9b6deaf2e8ee2f2b0334 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 15 Sep 2026 19:19:54 +0000 Subject: [PATCH 49/51] fix broken fp8 dtype --- op_tests/triton_tests/attention/test_mla_v4_triton.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/op_tests/triton_tests/attention/test_mla_v4_triton.py b/op_tests/triton_tests/attention/test_mla_v4_triton.py index be966fa696e..96a91d7f30b 100644 --- a/op_tests/triton_tests/attention/test_mla_v4_triton.py +++ b/op_tests/triton_tests/attention/test_mla_v4_triton.py @@ -10,6 +10,7 @@ from aiter.benchmark_data_init import fill, make_generator from aiter.ops.triton.attention.pa_decode_sparse import pa_decode_sparse from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import ( benchmark, checkAllclose, @@ -455,13 +456,13 @@ def test_pa_decode_sparse_vs_reference( # --------------------------------------------------------------------------- _FP8_GROUP_SIZE = 64 -_FP8_DTYPE = torch.float8_e4m3fnuz +_FP8_DTYPE = get_fp8_e4m3_dtype() def _quantize_kv_fp8(unified_kv, group_size=_FP8_GROUP_SIZE): """Quantize bf16/fp16 unified_kv to (fp8, scales) with 1xGROUP_SIZE block scaling. - Returns (kv_fp8, kv_scales) where kv_fp8 is float8_e4m3fnuz and + Returns (kv_fp8, kv_scales) where kv_fp8 is _FP8_DTYPE and kv_scales is [total_pages, D // group_size] fp32. """ total_pages, D = unified_kv.shape From 0762106ec25d2c1ab076bb14e47820998fd61403 Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 15 Sep 2026 20:20:42 +0000 Subject: [PATCH 50/51] simplify comments --- .../op_benchmarks/triton/bench_sparse_mla.py | 46 ++++++++----------- 1 file changed, 19 insertions(+), 27 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_mla.py b/op_tests/op_benchmarks/triton/bench_sparse_mla.py index 2150e0d6701..a21ec66546e 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_mla.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -3,8 +3,8 @@ """Benchmark for sparse_mla_fwd (gfx950 gluon, separated-rope MLA). -The cache is flushed between iterations by default, which matters: a decode shape -runs up to twice as fast when the loop is allowed to re-read its KV out of cache. +The cache is flushed between iterations by default. Leaving it warm lets the loop +re-read its KV and flatters decode shapes badly. Usage: python op_tests/op_benchmarks/triton/bench_sparse_mla.py @@ -19,8 +19,9 @@ from torch.autograd import DeviceType from torch.profiler import ProfilerActivity, profile -from aiter.ops.triton.attention.sparse_mla import _mla_num_splits, sparse_mla_fwd +from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from op_tests.op_benchmarks.triton.utils.benchmark_utils import ( get_caller_name_no_ext, ) @@ -28,10 +29,10 @@ KV_LORA_RANK = 512 QK_ROPE_HEAD_DIM = 64 D_QK = KV_LORA_RANK + QK_ROPE_HEAD_DIM -E4M3_MAX = 448.0 +E4M3_DTYPE = get_fp8_e4m3_dtype() +E4M3_MAX = torch.finfo(E4M3_DTYPE).max -# Both kernels the launch runs are named after the gluon kernel, the combine as -# _sparse_mla_reduce, so one substring picks up both and nothing else. +# Picks up the main kernel and _sparse_mla_reduce, nothing else. KERNEL_MATCH = "_sparse_mla" FLUSH_BYTES = 512 << 20 @@ -69,17 +70,17 @@ def device_time_ms(func, warmup=25, rep=100, flush=True): return total / rep / 1e3 -def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes, num_splits): - """Bytes the launch actually moves, counting rows as gathered.""" +def bytes_moved(num_tokens, num_heads, nnz, kv_elem_bytes): + """Bytes the launch actually moves, counting rows as gathered. + + Split-K partials are left out: the split count is the kernel's own decision, + so when it does split, that traffic lands in the time but not here. + """ kv = nnz * D_QK * kv_elem_bytes # the gather, and the bulk of it idx = nnz * 4 # int32 index stream, read once q = num_tokens * num_heads * D_QK * kv_elem_bytes out = num_tokens * num_heads * KV_LORA_RANK * 2 - partials = 0 - if num_splits > 1: - # [tokens, splits, heads, lora] bf16 accumulator plus the m/l pair - partials = 2 * num_tokens * num_splits * num_heads * (KV_LORA_RANK * 2 + 8) - return kv + idx + q + out + partials + return kv + idx + q + out def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): @@ -88,9 +89,6 @@ def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): Every sequence owns its own slice of the pool, so there is no cross-request KV sharing to inflate the cache hit rate. Per-token KV length is min(pos + 1, topk), which is what the index converter emits. - - A single query token per sequence is a decode step, so it sits at the end of - the context. More than one means prefill, and they walk positions 0..n-1. """ torch.manual_seed(0) gen = torch.Generator().manual_seed(1) @@ -126,9 +124,9 @@ def build_case(num_seqs, num_tokens, num_heads, context, topk, device="cuda"): # fp8 dots want the flat per-tensor format; bf16 dots read kv as it is. scale = (kv.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) - kv_fp8 = ( - (kv.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) - ).view(torch.uint8) + kv_fp8 = ((kv.float() / scale).clamp(-E4M3_MAX, E4M3_MAX).to(E4M3_DTYPE)).view( + torch.uint8 + ) return q, kv, kv_fp8, scale.to(device), indices, indptr.to(device) @@ -185,7 +183,7 @@ def bench_sparse_mla( # Hand q over already quantized, the way production does. cache, scale = kv_fp8, kv_scale q_scale = (q.float().abs().amax() / E4M3_MAX).clamp_min(1e-30).reshape(1) - q = (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn) + q = (q.float() / q_scale).clamp(-E4M3_MAX, E4M3_MAX).to(E4M3_DTYPE) else: cache, scale, q_scale = kv, None, None @@ -209,13 +207,7 @@ def func(): num_tokens = q.shape[0] # QK reads the whole row, PV only the latent half flops = 2.0 * num_heads * nnz * (D_QK + KV_LORA_RANK) - moved = bytes_moved( - num_tokens, - num_heads, - nnz, - 1 if dots == "fp8" else 2, - _mla_num_splits(num_tokens, 1, nnz / num_tokens), - ) + moved = bytes_moved(num_tokens, num_heads, nnz, 1 if dots == "fp8" else 2) if metric == "time": return time_ms From e28ca27bfbdbdc98eefa6485dea8a097c870e62f Mon Sep 17 00:00:00 2001 From: Mehmet Cagri Kaymak Date: Tue, 15 Sep 2026 20:24:38 +0000 Subject: [PATCH 51/51] fix hardcoded fp8 max limit --- op_tests/triton_tests/attention/test_sparse_mla.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/op_tests/triton_tests/attention/test_sparse_mla.py b/op_tests/triton_tests/attention/test_sparse_mla.py index 008426d38da..7ac5026a566 100644 --- a/op_tests/triton_tests/attention/test_sparse_mla.py +++ b/op_tests/triton_tests/attention/test_sparse_mla.py @@ -11,12 +11,15 @@ import torch from aiter.ops.triton.utils._triton import arch_info +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype if arch_info.get_arch() == "gfx950": import aiter.ops.triton.attention.sparse_mla as smd from aiter.ops.triton.attention.sparse_mla import sparse_mla_fwd -FP8_MAX = 448.0 +# The packed fp8_ds_mla record is OCP e4m3 by definition of the format. +FP8_DTYPE = get_fp8_e4m3_dtype() +FP8_MAX = torch.finfo(FP8_DTYPE).max KV_LORA, ROPE = 512, 64 D_QK = KV_LORA + ROPE @@ -29,7 +32,7 @@ def _skip_unless_gfx950(): def quantize_flat_fp8(kv): """vLLM's production layout: whole row fp8 with one per-tensor scale.""" scale = (kv.float().abs().amax() / FP8_MAX).clamp_min(1e-30).reshape(1) - q8 = (kv.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + q8 = (kv.float() / scale).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) return q8.view(torch.uint8), scale.to(torch.float32) @@ -40,7 +43,7 @@ def pack_ds_mla(kv, block_size=64): cache = torch.zeros(nb * block_size, 656, dtype=torch.uint8, device=kv.device) x = kv[:, :KV_LORA].float().reshape(T, 4, 128) scale = (x.abs().amax(-1) / FP8_MAX).clamp_min(1e-30) - q8 = (x / scale[..., None]).to(torch.float8_e4m3fn).view(torch.uint8) + q8 = (x / scale[..., None]).to(FP8_DTYPE).view(torch.uint8) cache[:T, :KV_LORA] = q8.reshape(T, KV_LORA) cache[:T, 512:528].view(torch.float32).copy_(scale) cache[:T, 528:].view(torch.bfloat16).copy_(kv[:, KV_LORA:]) @@ -48,12 +51,12 @@ def pack_ds_mla(kv, block_size=64): def dequant_flat_fp8(u8, scale): - return u8.view(torch.float8_e4m3fn).float() * scale + return u8.view(FP8_DTYPE).float() * scale def dequant_ds_mla(cache): flat = cache.reshape(-1, 656) - x = flat[:, :KV_LORA].view(torch.float8_e4m3fn).float().reshape(-1, 4, 128) + x = flat[:, :KV_LORA].view(FP8_DTYPE).float().reshape(-1, 4, 128) sc = flat[:, 512:528].view(torch.float32) nope = (x * sc[..., None]).reshape(-1, KV_LORA) rope = flat[:, 528:].view(torch.bfloat16).float()