From 3a427ef10b3268c639120a6de6367a9e92659280 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Fri, 14 Aug 2026 14:43:24 -0500 Subject: [PATCH 01/15] [TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward Adds the training backward for the DeepSeek-V4 sparse MLA attention on gfx950 (CDNA4). This is the counterpart to the DSv4 sparse prefill forward added in #3833; aiter had no sparse-MLA backward before this. Same op contract as the forward's has_pe=False path: shared-KV GQA with K == V == kv as one dense 512-wide tensor, RoPE applied in place caller-side, scale 1/sqrt(512), attn_sink folded into the softmax denominator only, and topk_indices == -1 masked out. from aiter.ops.triton.gluon.mla_gluon import mla_gluon from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dsv4, ) o, lse = mla_gluon(..., has_pe=False, attn_sink=sink, return_lse=True) dq, dkv, d_sink = sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk, attn_sink=sink) mla_gluon's lse is already sink-inclusive (the sink is folded into e_max/e_sum before lse = e_max + log(e_sum)), which is the convention this backward expects, so the two compose directly with no adaptation. Five kernels plus one torch reduction: delta triton rowsum(O*dO), streams bf16 and accumulates fp32 dQ gluon one LDS read of the gathered KV feeds both the S and dP MFMAs; also emits this chunk's dS / P dKV-interm gluon contracts over all heads inside one MFMA pair, Q/dO transposed once into registers, D split across grid.y CSR build torch sort + searchsorted on int16-narrowed keys dKV gather triton atomic-free -- the scatter is inverted so each KV row gathers its own contributors d_sink torch 26 us num_kv >= T is supported, so a compressed pool (kv rows T..num_kv-1) works. R_CHUNK splits the rank dimension and defaults to unchunked; it exists only to bound the interm intermediate (T*topk*512 bf16, 2.0 GiB at T=4096 topk=512) and costs a dQ read-modify-write between chunks plus one CSR build per chunk. Performance (MI355X, T=4096 H=128 topk=512, realistic SWA(128)+pool top-k), against the Primus-Turbo FlyDSL backward on identical tensors in one process: ours flydsl per-kernel sum 3.380 ms 3.223 ms 407 vs 426 TFLOPS end-to-end 3.238 ms 3.471 ms 424 vs 396 TFLOPS Read that as parity: the two differ mainly in how much each allocates per call. Correctness: cos > 0.999 on dq / dkv / d_sink against torch autograd through an independent fp32 reference forward, across H in {64,128}, pool and no-pool, chunked and unchunked, with and without attn_sink. Note the test skips on non-gfx950, so it will not exercise anything unless CI has a CDNA4 runner. It has been run on MI355X. --- .../attention/sparse_attention_dsv4_bwd.py | 172 +++++ .../attention/sparse_attention_dsv4_bwd.py | 168 +++++ aiter/ops/triton/gluon/README.md | 30 + .../gluon/sparse_attention_dsv4_bwd_gluon.py | 667 ++++++++++++++++++ .../test_sparse_attention_dsv4_bwd.py | 137 ++++ 5 files changed, 1174 insertions(+) create mode 100644 aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py create mode 100644 aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py create mode 100644 aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py create mode 100644 op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..cafe9f75b6 --- /dev/null +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Triton kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +``_delta_v4_kernel`` + ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the + bf16 inputs and accumulates in fp32, so it moves exactly the working set. + +``_bwd_dkv_gather_acc_v4_be`` + ``build_inverted_topk_fast`` + Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is + inverted into a CSR gather (each output KV row collects its own contributors), so no atomics + are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load + and cuts the trip count on the long runs a realistic top-k produces. + +Public entry: ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4``. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _delta_v4_kernel( + O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) + dO_ptr, # [n_rows, D] bf16 + Delta_ptr, # [n_rows] fp32 + n_rows, + D: tl.constexpr, + BLOCK_R: tl.constexpr, +): + """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + mask = rows < n_rows + offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] + o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) + + +def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): + # BLOCK_R=8 / num_warps=8 measured best (0.173 ms, 6.21 TB/s = 78% peak at T4096 H128); + # the whole sweep plateaus at 0.173-0.187 once a lane loads >= 8 bf16, i.e. once the load + # is a dwordx4. Below that (BLOCK_R=2 nw=8, 2 bf16/lane) it falls off a cliff to 3.12 TB/s. + """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. + + ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. + """ + assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() + T, H, D = o.shape + n_rows = T * H + if out is None: + out = torch.empty(T, H, dtype=torch.float32, device=o.device) + _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( + o, + do, + out, + n_rows, + D=D, + BLOCK_R=BLOCK_R, + num_warps=num_warps, + ) + return out + + +@triton.jit +def _bwd_dkv_gather_acc_v4_be( + Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] + InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers + InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token + dKV_acc_ptr, # [num_kv, D] fp32 — accumulator + stride_interm_r: tl.int64, + stride_acc_t: tl.int64, + D: tl.constexpr, + BLOCK_E: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. + + Fixes two things about ``_bwd_dkv_gather_acc_v4``, which walks the run one entry at a time + with a bare ``tl.arange(0, D)``: + + * **load width.** A [D] block over 256 threads is 2 bf16 = 4 B per lane -- a dword. The + [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane instead, so the loads + become dwordx4. The gather is issue-bound, so this is the dominant term. + * **loop trip count.** The run is consumed BLOCK_E entries at a time rather than one, and + ``tl.sum`` over the entry axis folds them. The realistic topk gives run lengths up to + ~3000 (pool rows), so the serial walk was the other half of the problem. + + ``ACCUMULATE=False`` skips the read-modify-write of the destination, valid when the caller + does not chunk (each KV row is then written by exactly one CTA). + """ + k = tl.program_id(0) + offs_d = tl.arange(0, D) + offs_e = tl.arange(0, BLOCK_E) + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + acc_base = k.to(tl.int64) * stride_acc_t + + if ACCUMULATE: + acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) + else: + acc = tl.zeros([D], dtype=tl.float32) + + for i0 in range(start, end, BLOCK_E): + idx = i0 + offs_e + m = idx < end + entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) + vals = tl.load( + Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], + mask=m[:, None], + other=0.0, + ) + acc += tl.sum(vals.to(tl.float32), axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_d, acc) + + +def build_inverted_topk_fast(topk_indices_slice, num_kv): + """CSR inverted index over ``num_kv`` KV rows. Bit-identical to the reference below. + + One stable sort yields both the permutation (``inv_data``) and the sorted keys; + ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k, + which is exactly what ``cumsum(bincount(flat+1))`` computes. Invalid (-1) entries sort to + the front, so ``inv_ptr[0]`` starts past them and they are never visited. + + Two things make this ~3x faster than the reference: + * the sort key is narrowed to int16 when ``num_kv`` fits, so the radix sort makes 2 + byte-passes instead of 8; + * ``searchsorted`` replaces the separate ``bincount`` + ``cumsum`` passes. + + Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. + """ + flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid + if num_kv < 32767: # int16 range, -1 included + keys = flat_kv.to(torch.int16) + ar = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) + else: + keys = flat_kv.to(torch.int32) + ar = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) + sorted_vals, inv_data = torch.sort(keys, stable=True) + inv_ptr = torch.searchsorted(sorted_vals, ar).to(torch.int32) + return inv_ptr, inv_data.to(torch.int32) + + +def dkv_gather_acc_be( + interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True +): + # BLOCK_E=64 / num_warps=8 measured best (0.345 ms, 6.29 TB/s = 79% peak at T4096 H128 + # topk512 SWA+pool). Time falls monotonically with BLOCK_E across the whole sweep + # (4->64: 1.021, 0.715, 0.505, 0.407, 0.345), i.e. both the load width AND the trip count + # on the ~3000-entry pool runs were binding. Reference was 1.491 ms at 1.45 TB/s. + """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. + + Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. + """ + _, _, D = interm.shape + num_kv = dkv_acc.shape[0] + _bwd_dkv_gather_acc_v4_be[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D=D, + BLOCK_E=BLOCK_E, + ACCUMULATE=accumulate, + num_warps=num_warps, + ) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..603cb31275 --- /dev/null +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +Counterpart to the DSv4 sparse prefill forward. The op is the official V4 form: shared-KV GQA +where ``K == V == kv`` is a single dense 512-wide tensor, RoPE already applied in place +caller-side, scale ``1/sqrt(512)``, ``attn_sink`` folded into the softmax denominator only, and +``topk_indices == -1`` masked out. + + P = exp(Q@kv^T * scale - lse) + dP = dO@kv^T + delta = rowsum(O * dO) + dS = P * (dP - delta) * scale + dQ = dS @ kv + dKV = scatter_add over top-k of sum_h ( dS*Q + P*dO ) + +Pipeline (five kernels + one torch reduction), per rank chunk: + + delta triton rowsum(O*dO) + dQ gluon also emits this chunk's dS / P + dKV-interm gluon interm[t, slot, d] = sum_h (dS*Q + P*dO) + CSR build torch inverted top-k index (sort + searchsorted) + dKV gather triton reduce interm over the top-k mapping, atomic-free + d_sink torch 26 us, not worth a kernel + +``lse`` and ``o`` come from the forward. The merged Gluon prefill kernel produces both:: + + from aiter.ops.triton.gluon.mla_gluon import mla_gluon + o, lse = mla_gluon(..., has_pe=False, attn_sink=sink, return_lse=True) + +Its ``lse`` is sink-inclusive (the sink is folded into ``e_max``/``e_sum`` before +``lse = e_max + log(e_sum)``), which is the convention this backward expects. + +Measured on MI355X at ``T=4096 H=128 topk=512`` with a realistic SWA(128)+pool top-k: +delta 0.178 / dQ 1.391 / interm 1.152 / CSR build 0.130 / gather 0.503 / d_sink 0.026 ms, +3.380 ms total = 407 TFLOPS. +""" + +import torch + +from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( + build_inverted_topk_fast, + delta_v4, + dkv_gather_acc_be, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( + sparse_mla_bwd_dkv_interm_v4_bd as _dkv_interm_gluon, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( + sparse_mla_bwd_dq_gluon as _dq_gluon, +) +from aiter.ops.triton.utils._triton import arch_info + +_BLOCK_H_DQ = 64 +_TILE_K_DQ = 32 +_BD_DKV = 256 +_TILE_K_DKV = 128 + + +def sparse_mla_bwd_dsv4( + q, + kv, + do, + o, + lse, + topk_indices, + attn_sink=None, + scale=None, + R_CHUNK=None, +): + """Backward for the DSv4 sparse-MLA prefill attention. gfx950 (CDNA4) only. + + Args: + q: [T, H, 512] bf16 + kv: [num_kv, 512] bf16, K == V. ``num_kv >= T``; rows ``T..num_kv-1`` are + the compressed pool, which ``topk_indices`` may reference. + do: [T, H, 512] bf16, gradient of the attention output + o: [T, H, 512] bf16, the forward output + lse: [T, H] fp32, sink-inclusive log-sum-exp from the forward + topk_indices: [T, TOPK] int32, -1 marks an invalid slot + attn_sink: [H] fp32 per-head sink bias, or None + scale: softmax scale, defaults to ``1/sqrt(512)`` + R_CHUNK: split the rank dimension into chunks of this width. ``None`` (default) + runs unchunked, which is what you want. Chunking exists only to bound the + ``interm`` intermediate, which is ``T*TOPK*512`` bf16 (2.0 GiB at + T=4096, TOPK=512); it costs a dQ read-modify-write between chunks and one + CSR build per chunk. Any multiple of 32 is accepted. + + Returns: + dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) + """ + assert ( + arch_info.get_arch() == "gfx950" + ), f"sparse_mla_bwd_dsv4 requires gfx950 (CDNA4), got {arch_info.get_arch()}" + + T, H, D = q.shape + TOPK = topk_indices.shape[1] + num_kv = kv.shape[0] + assert D == 512, f"DSv4 sparse-MLA backward is fixed to head_dim 512, got {D}" + assert kv.shape[-1] == D and do.shape == q.shape and o.shape == q.shape + assert num_kv >= T, f"num_kv ({num_kv}) must be >= T ({T})" + assert q.is_contiguous() and kv.is_contiguous() and do.is_contiguous() + assert o.is_contiguous() and topk_indices.is_contiguous() + + if scale is None: + scale = 1.0 / (D**0.5) + if R_CHUNK is None: + R_CHUNK = TOPK + lse = lse.float().contiguous() + + # Both mfma tiles must divide the chunk width; step down rather than making it the + # caller's problem, so a small R_CHUNK still works. + tk_dkv = next( + (t for t in (_TILE_K_DKV, 64, 32) if t <= R_CHUNK and R_CHUNK % t == 0), None + ) + tk_dq = next( + (t for t in (_TILE_K_DQ, 32) if t <= R_CHUNK and R_CHUNK % t == 0), None + ) + assert ( + tk_dkv is not None and tk_dq is not None + ), f"R_CHUNK={R_CHUNK} must be a multiple of 32 (it is the mfma tile width)" + + delta = delta_v4(o, do) + + dq = torch.empty_like(q) + dkv_acc = torch.zeros(num_kv, D, dtype=torch.float32, device=q.device) + chunk_dS = torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device) + interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) + + for r in range(0, TOPK, R_CHUNK): + _dq_gluon( + q, + kv, + do, + topk_indices, + lse, + delta, + dq, + chunk_dS, + chunk_P, + scale, + r, + R_CHUNK, + BLOCK_H=_BLOCK_H_DQ, + TILE_K=tk_dq, + is_first_chunk=(r == 0), + ) + _dkv_interm_gluon( + q, do, chunk_dS, chunk_P, R_CHUNK, BD=_BD_DKV, TILE_K=tk_dkv, interm=interm + ) + inv_ptr, inv_data = build_inverted_topk_fast( + topk_indices[:, r : r + R_CHUNK], num_kv + ) + dkv_gather_acc_be(interm, inv_ptr, inv_data, dkv_acc) + + dkv = dkv_acc.to(kv.dtype) + + d_sink = None + if attn_sink is not None: + # d_sink[h] = -sum_t exp(sink[h] - lse[t,h]) * delta[t,h] + d_sink = -(torch.exp(attn_sink[None, :].float() - lse) * delta).sum(dim=0) + + return dq, dkv, d_sink + + +__all__ = ["sparse_mla_bwd_dsv4"] diff --git a/aiter/ops/triton/gluon/README.md b/aiter/ops/triton/gluon/README.md index d6477dc63a..fd693926b9 100644 --- a/aiter/ops/triton/gluon/README.md +++ b/aiter/ops/triton/gluon/README.md @@ -57,6 +57,12 @@ Some features (e.g., scheduling hints like `sched_barrier`) require the [AMD Glu python op_tests/triton_tests/
test_pa_decode_gluon.py TBDTBDTBD + + sparse_attention_
dsv4_bwd_gluonDSv4 Sparse
MLA BackwardCDNA4 + Q/KV/dO/O: bf16, K == V
head_dim = 512 (dense)
lse: fp32, sink-inclusive
num_kv ≥ T (pool ok)
topk % 32 == 0
gfx950 only + python op_tests/triton_tests/
attention/test_sparse_
attention_dsv4_bwd.py + ~407
TFLOPS—— + @@ -247,6 +253,30 @@ python op_tests/test_mla.py -c 10000 100000 -b 1 3 4 -n 16,1 -d bf16 -kvd bf16 - | 100K | 3 | 85 | 19 | 88.77 | 3.89 | | 100K | 4 | 64 | 25 | 106.96 | 4.31 | +### `sparse_attention_dsv4_bwd_gluon.py` — DeepSeek V4 Sparse MLA Backward + +**Public entry:** `aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None)` -> `(dq, dkv, d_sink)` + +Training backward for the DSv4 sparse prefill attention (the `has_pe=False` form of `mla_gluon`). Same op contract: `K == V == kv` as one dense 512-wide tensor, RoPE applied in place caller-side, scale `1/sqrt(512)`, `attn_sink` in the softmax denominator only, `topk_indices == -1` masked. + +`o` and `lse` come from the forward; `mla_gluon(..., has_pe=False, return_lse=True)` produces both, and its `lse` is already sink-inclusive, which is what this backward expects. + +Five kernels plus one torch reduction: + +| phase | impl | notes | +|---|---|---| +| `delta = rowsum(O*dO)` | triton | streams bf16, accumulates fp32 | +| dQ | **gluon** | one LDS read of the gathered KV feeds both the `S` and `dP` MFMAs; also emits this chunk's `dS`/`P` | +| dKV-interm | **gluon** | contracts over all heads inside one MFMA pair, Q/dO transposed once into registers, D split across `grid.y` | +| CSR build | torch | `sort` + `searchsorted` on int16-narrowed keys | +| dKV gather | triton | atomic-free: the scatter is inverted so each KV row gathers its own contributors | +| `d_sink` | torch | 26 us | + +`R_CHUNK` splits the rank dimension. Leave it `None` (unchunked) unless memory forces otherwise: it exists only to bound the `interm` intermediate (`T*topk*512` bf16, 2.0 GiB at T=4096 topk=512) and costs a dQ read-modify-write between chunks plus one CSR build per chunk. + +**Measured** (MI355X, T=4096 H=128 topk=512, SWA(128)+pool top-k): 3.38 ms / 407 TFLOPS as a per-kernel sum, 3.24 ms / 424 TFLOPS end-to-end. + + ### `pa_decode_gluon.py` — Paged Attention Decode **Function:** `pa_decode_gluon(output, query, key_cache, value_cache, context_lengths, block_tables, softmax_scale, query_length, max_context_partition_num, context_partition_size, compute_type, query_scale, key_scale, value_scale, ...)` diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py new file mode 100644 index 0000000000..16bffd4273 --- /dev/null +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py @@ -0,0 +1,667 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Gluon kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +Two kernels, both operating on the official V4 form (``K == V == kv``, one dense 512-wide +tensor, RoPE already applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` in the +softmax denominator only, ``topk == -1`` masked): + +``_dq_v4_kernel`` + Per (query token, head block): ``S = Q@kv^T``, ``P = exp(S - lse)``, ``dP = dO@kv^T``, + ``dS = P*(dP - delta)*scale``, ``dQ += dS@kv``. Three MFMAs per tile; the gathered KV tile + is read from LDS once and feeds both the ``S`` and ``dP`` MFMAs. Also emits the ``dS`` / ``P`` + chunks the dKV-interm kernel consumes. + +``_dkv_interm_v4_bd_kernel`` + ``interm[t, slot, d] = sum_h ( dS[t,h,slot]*Q[t,h,d] + P[t,h,slot]*dO[t,h,d] )``, contracting + over ALL heads inside one MFMA pair so nothing accumulates across a loop over heads. Q and dO + are transposed once into registers and D is split across ``grid.y``, which is what keeps them + read once instead of ``topk/TILE_K`` times. + +Public entry: ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4``. +""" + +import torch +import triton +import triton.language as tl +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def _dq_v4_kernel( + Q_ptr, # [T, H, D] bf16 + KV_ptr, # [T, D] bf16 (K == V) + dO_ptr, # [T, H, D] bf16 + TopK_ptr, # [T, TOPK_padded] int32 + LSE_ptr, # [T, H] fp32 (sink-inclusive) + Delta_ptr, # [T, H] fp32 + dQ_ptr, # [T, H, D] bf16 (RMW across chunks) + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_dq_t: tl.int64, + stride_dq_h: tl.int64, + stride_topk_t: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + R_START: tl.int32, + R_CHUNK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D: gl.constexpr, + IS_FIRST_CHUNK: gl.constexpr, +): + gl.static_assert(TILE_K % 32 == 0, "16x16x32 needs TILE_K multiple of 32") + gl.static_assert(D % 32 == 0, "16x16x32 needs D multiple of 32") + + # ---- single 16x16x32 MFMA parent (score contracts D; accumulate contracts TILE_K) ---- + mma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 32], + transposed=True, + warps_per_cta=[4, 1], + ) + qa: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mma, k_width=8) + qb: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mma, k_width=8) + + # ---- blocked layouts for global loads ---- + _q_tpw_k: gl.constexpr = min(64, D // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( # [BLOCK_H, D] (Q, dO, dQ) + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + _kv_tpw_m: gl.constexpr = min(64, D // 8) + _kv_tpw_n: gl.constexpr = 64 // _kv_tpw_m + blk_kv: gl.constexpr = gl.BlockedLayout( # [D, TILE_K] + size_per_thread=[8, 1], + threads_per_warp=[_kv_tpw_m, _kv_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + sh_kv: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], [D, TILE_K], [0, 1] + ) + + # ---- program ids ---- + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + # ---- Q / dO offsets + load (register, convert to dot operand) ---- + offs_h_q = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_d_q = gl.arange(0, D, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs = ( + q_base + + offs_h_q[:, None].to(tl.int64) * stride_q_h + + offs_d_q[None, :].to(tl.int64) + ) + do_base = token_idx.to(tl.int64) * stride_do_t + do_offs = ( + do_base + + offs_h_q[:, None].to(tl.int64) * stride_do_h + + offs_d_q[None, :].to(tl.int64) + ) + + q_blk = gl.amd.cdna4.buffer_load( + ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None], other=0.0 + ) + do_blk = gl.amd.cdna4.buffer_load( + ptr=dO_ptr, offsets=do_offs.to(tl.int32), mask=mask_h_q[:, None], other=0.0 + ) + Q_dot = gl.convert_layout(q_blk, qa) + dO_dot = gl.convert_layout(do_blk, qa) + + # ---- topk / KV offsets ---- + topk_base = token_idx.to(tl.int64) * stride_topk_t + R_START + stride_kv_t_i32: tl.int32 = stride_kv_t.to(tl.int32) + offs_tile_kv = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_kv)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mma)) + offs_d_kv = gl.arange(0, D, layout=gl.SliceLayout(1, blk_kv)) + + smem_kv = gl.allocate_shared_memory( + KV_ptr.dtype.element_ty, [2, D, TILE_K], layout=sh_kv + ) + + dQ_acc = gl.zeros([BLOCK_H, D], dtype=gl.float32, layout=mma) + + offs_h_s = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + mask_h_s = offs_h_s < num_heads + lse = gl.amd.cdna4.buffer_load( + ptr=LSE_ptr, + offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), + mask=mask_h_s, + other=0.0, + ) + delta = gl.amd.cdna4.buffer_load( + ptr=Delta_ptr, + offsets=(token_idx * num_heads + offs_h_s).to(tl.int32), + mask=mask_h_s, + other=0.0, + ) + + # ---- prologue: gather kv tile 0 ---- + topk_pos_kv = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_kv).to(tl.int32), + mask=offs_tile_kv < R_CHUNK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + offs_tile_mma).to(tl.int32), + mask=offs_tile_mma < R_CHUNK, + other=-1, + ) + valid_kv = topk_pos_kv != -1 + valid_mma = topk_pos_mma != -1 + safe_kv = gl.where(valid_kv, topk_pos_kv, 0) + kv_offs = safe_kv[None, :] * stride_kv_t_i32 + offs_d_kv[:, None] + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_kv.index(0), ptr=KV_ptr, offsets=kv_offs, mask=valid_kv[None, :] + ) + gl.amd.cdna4.async_copy.commit_group() + + ds_base = ( + token_idx.to(tl.int64) * stride_ds_t + + hg_idx.to(tl.int64) * BLOCK_H * stride_ds_h + ) + offs_h_dsp = gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + offs_tile_dsp = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mma)) + mask_h_dsp = (hg_offset + offs_h_dsp) < num_heads + + cur_buf = 0 + for t in range(NUM_TILES - 1): + next_offs_kv = (t + 1) * TILE_K + offs_tile_kv + next_offs_mma = (t + 1) * TILE_K + offs_tile_mma + topk_pos_kv_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_kv).to(tl.int32), + mask=next_offs_kv < R_CHUNK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=TopK_ptr, + offsets=(topk_base + next_offs_mma).to(tl.int32), + mask=next_offs_mma < R_CHUNK, + other=-1, + ) + valid_kv_next = (next_offs_kv < R_CHUNK) & (topk_pos_kv_next != -1) + valid_mma_next = (next_offs_mma < R_CHUNK) & (topk_pos_mma_next != -1) + safe_kv_next = gl.where(valid_kv_next, topk_pos_kv_next, 0) + + next_buf = 1 - cur_buf + kv_offs_next = safe_kv_next[None, :] * stride_kv_t_i32 + offs_d_kv[:, None] + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_kv.index(next_buf), + ptr=KV_ptr, + offsets=kv_offs_next, + mask=valid_kv_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + + kv_smem_cur = smem_kv.index(cur_buf) + # score K (direct); V (permuted) read LATE, before the accumulate + K_T_dot = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem_cur, qb) + + S = gl.amd.cdna4.mfma( + Q_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, qa) + K_v_dot = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem_cur.permute([1, 0]), qb + ) # load V LATE + dQ_acc = gl.amd.cdna4.mfma(dS_dot, K_v_dot, dQ_acc) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ( + ds_base + + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + + col[None, :].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, + ptr=dS_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + cur_buf = next_buf + valid_mma = valid_mma_next + + # ---- epilogue: last tile ---- + gl.amd.cdna4.async_copy.wait_group(0) + t = NUM_TILES - 1 + kv_smem_cur = smem_kv.index(cur_buf) + K_T_dot = gl.amd.cdna4.async_copy.load_shared_relaxed(kv_smem_cur, qb) + + S = gl.amd.cdna4.mfma( + Q_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + S = S * scale + offs_h_mma = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, mma)) + valid_mask = valid_mma[None, :] & (offs_h_mma < num_heads)[:, None] + S = gl.where(valid_mask, S, float("-inf")) + + P = gl.exp(S - lse[:, None]) + P = gl.where(valid_mask, P, 0.0) + dP = gl.amd.cdna4.mfma( + dO_dot, K_T_dot, gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mma) + ) + dS = P * (dP - delta[:, None]) * scale + dS = gl.where(valid_mask, dS, 0.0) + + dS_bf = dS.to(KV_ptr.dtype.element_ty) + dS_dot = gl.convert_layout(dS_bf, qa) + K_v_dot = gl.amd.cdna4.async_copy.load_shared_relaxed( + kv_smem_cur.permute([1, 0]), qb + ) + dQ_acc = gl.amd.cdna4.mfma(dS_dot, K_v_dot, dQ_acc) + + col = t * TILE_K + offs_tile_dsp + dsp_offs = ( + ds_base + + offs_h_dsp[:, None].to(tl.int64) * stride_ds_h + + col[None, :].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dS_bf, + ptr=dS_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + gl.amd.cdna4.buffer_store( + stored_value=P.to(KV_ptr.dtype.element_ty), + ptr=P_ptr, + offsets=dsp_offs.to(tl.int32), + mask=mask_h_dsp[:, None], + ) + + # ---- store dQ (RMW across chunks) ---- + dq_base = token_idx.to(tl.int64) * stride_dq_t + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_q)) + offs_d_o = gl.arange(0, D, layout=gl.SliceLayout(0, blk_q)) + mask_h_o = offs_h_o < num_heads + dq_offs = ( + dq_base + + offs_h_o[:, None].to(tl.int64) * stride_dq_h + + offs_d_o[None, :].to(tl.int64) + ) + dq_blk = gl.convert_layout(dQ_acc.to(dQ_ptr.dtype.element_ty), blk_q) + if not IS_FIRST_CHUNK: + prev = gl.amd.cdna4.buffer_load( + ptr=dQ_ptr, offsets=dq_offs.to(tl.int32), mask=mask_h_o[:, None], other=0.0 + ) + dq_blk = (dq_blk.to(gl.float32) + prev.to(gl.float32)).to( + dQ_ptr.dtype.element_ty + ) + gl.amd.cdna4.buffer_store( + stored_value=dq_blk, + ptr=dQ_ptr, + offsets=dq_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + +def sparse_mla_bwd_dq_gluon( + q, + kv, + do, + topk, + lse, + delta, + dq, + chunk_dS, + chunk_P, + scale, + r_start, + R_CHUNK, + BLOCK_H=64, + TILE_K=32, + is_first_chunk=True, +): + """Launch the dQ kernel for one rank chunk. Writes ``dq`` (RMW when not the first chunk) + plus this chunk's ``chunk_dS`` / ``chunk_P``.""" + T, H, D = q.shape + _dq_v4_kernel[(T, triton.cdiv(H, BLOCK_H))]( + q, + kv, + do, + topk, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + H, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D=D, + IS_FIRST_CHUNK=is_first_chunk, + num_warps=4, + waves_per_eu=1, + ) + + +@gluon.jit +def _dkv_interm_v4_bd_kernel( + Q_ptr, # [T, H, D] bf16 + dO_ptr, # [T, H, D] bf16 + dS_ptr, # [T, H, R_CHUNK] bf16 + P_ptr, # [T, H, R_CHUNK] bf16 + Interm_ptr, # [T, R_CHUNK, D] bf16 + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_do_t: tl.int64, + stride_do_h: tl.int64, + stride_ds_t: tl.int64, + stride_ds_h: tl.int64, + stride_interm_t: tl.int64, + stride_interm_r: tl.int64, + num_heads: tl.int32, + R_CHUNK: gl.constexpr, + TILE_K: gl.constexpr, + NH: gl.constexpr, + BD: gl.constexpr, + D: gl.constexpr, + MFMA_K: gl.constexpr, + DUAL_STAGE: gl.constexpr, + PREFETCH: gl.constexpr, +): + """Grid (T, D//BD). NH is the padded head count and the mfma contraction dim.""" + # instr_shape[2]=32: on gfx950 v_mfma_f32_16x16x32_bf16 does 2x the FLOPs of the 16-deep + # form in the same 16 cycles. The old kernel used 16 and it did not matter there because it + # was bandwidth-saturated at 7.3 TB/s; once the traffic is halved the matrix rate binds. + mfma: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, MFMA_K], + transposed=True, + warps_per_cta=[4, 1], + ) + _q_tpw_k: gl.constexpr = min(64, BD // 8) + _q_tpw_m: gl.constexpr = 64 // _q_tpw_k + blk_q: gl.constexpr = gl.BlockedLayout( # [H, BD] global load + size_per_thread=[1, 8], + threads_per_warp=[_q_tpw_m, _q_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_ds: gl.constexpr = gl.BlockedLayout( # [H, TILE_K] global load + size_per_thread=[1, 4], + threads_per_warp=[16, 4], + warps_per_cta=[4, 1], + order=[1, 0], + ) + sh_q: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], [NH, BD], [1, 0] + ) + + dot_a: gl.constexpr = gl.DotOperandLayout(operand_index=0, parent=mfma, k_width=8) + dot_b: gl.constexpr = gl.DotOperandLayout(operand_index=1, parent=mfma, k_width=8) + + token_idx = gl.program_id(axis=0) + dblk = gl.program_id(axis=1) + d_off = dblk * BD + NUM_TILES: gl.constexpr = R_CHUNK // TILE_K + + q_base = token_idx.to(tl.int64) * stride_q_t + do_base = token_idx.to(tl.int64) * stride_do_t + ds_base = token_idx.to(tl.int64) * stride_ds_t + interm_base = token_idx.to(tl.int64) * stride_interm_t + + # ---- prologue: stage Q, transpose to [BD, H] registers; reuse the buffer for dO ---- + offs_h_q = gl.arange(0, NH, layout=gl.SliceLayout(1, blk_q)) + offs_d_q = d_off + gl.arange(0, BD, layout=gl.SliceLayout(0, blk_q)) + mask_h_q = offs_h_q < num_heads + q_offs = ( + q_base + + offs_h_q[:, None].to(tl.int64) * stride_q_h + + offs_d_q[None, :].to(tl.int64) + ) + do_offs = ( + do_base + + offs_h_q[:, None].to(tl.int64) * stride_do_h + + offs_d_q[None, :].to(tl.int64) + ) + + if DUAL_STAGE: + # two buffers -> BOTH HBM->LDS copies in flight behind ONE drain. Costs 2x LDS + # (128 KB at BD=256, so occ-1) but halves the exposed prologue latency. + smem_q = gl.allocate_shared_memory( + Q_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + smem_do = gl.allocate_shared_memory( + dO_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_q, ptr=Q_ptr, offsets=q_offs.to(tl.int32), mask=mask_h_q[:, None] + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_do, + ptr=dO_ptr, + offsets=do_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_q.permute([1, 0]).load(dot_a) # [BD, H], register-resident + dO_T = smem_do.permute([1, 0]).load(dot_a) + else: + # one buffer, re-used after a barrier: half the LDS (occ-2 at BD=256) but the two + # HBM round trips serialize. + smem_stage = gl.allocate_shared_memory( + Q_ptr.dtype.element_ty, [NH, BD], layout=sh_q + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_stage, + ptr=Q_ptr, + offsets=q_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + Q_T = smem_stage.permute([1, 0]).load(dot_a) + gl.barrier() # all warps done reading before reuse + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_stage, + ptr=dO_ptr, + offsets=do_offs.to(tl.int32), + mask=mask_h_q[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(0) + dO_T = smem_stage.permute([1, 0]).load(dot_a) + + # ---- main loop: rank tile inner, head contraction folded into the mfma ---- + offs_h_ds = gl.arange(0, NH, layout=gl.SliceLayout(1, blk_ds)) + offs_k_ds = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_ds)) + mask_h_ds = offs_h_ds < num_heads + offs_d_st = d_off + gl.arange(0, BD, layout=gl.SliceLayout(1, mfma)) + offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) + + # dS/P prefetch one rank tile ahead: without it the HBM load latency is fully exposed every + # iteration, which is what caps the restructured kernel once it stops being BW-bound. The + # last iteration re-loads tile NUM_TILES-1 rather than branching -- one redundant 8 KB load + # in NUM_TILES, cheaper than peeling the loop. + offs0 = ( + ds_base + + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + + offs_k_ds[None, :].to(tl.int64) + ) + dS_nxt = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=offs0.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_nxt = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=offs0.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + + for t in range(NUM_TILES): + if PREFETCH: + dS_blk = dS_nxt + P_blk = P_nxt + # re-load the last tile rather than branch: one redundant 8 KB load in NUM_TILES + t_nxt = min(t + 1, NUM_TILES - 1) + col_n = t_nxt * TILE_K + offs_k_ds + offs_n = ( + ds_base + + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + + col_n[None, :].to(tl.int64) + ) + dS_nxt = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, + offsets=offs_n.to(tl.int32), + mask=mask_h_ds[:, None], + other=0.0, + ) + P_nxt = gl.amd.cdna4.buffer_load( + ptr=P_ptr, + offsets=offs_n.to(tl.int32), + mask=mask_h_ds[:, None], + other=0.0, + ) + else: + col_c = t * TILE_K + offs_k_ds + offs_c = ( + ds_base + + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + + col_c[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, + offsets=offs_c.to(tl.int32), + mask=mask_h_ds[:, None], + other=0.0, + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, + offsets=offs_c.to(tl.int32), + mask=mask_h_ds[:, None], + other=0.0, + ) + + dS_dot = gl.convert_layout(dS_blk, dot_b) + P_dot = gl.convert_layout(P_blk, dot_b) + + dKV = gl.zeros([BD, TILE_K], dtype=gl.float32, layout=mfma) + dKV = gl.amd.cdna4.mfma(Q_T, dS_dot, dKV) + dKV = gl.amd.cdna4.mfma(dO_T, P_dot, dKV) + + col_st = t * TILE_K + offs_col_st + interm_offs = ( + interm_base + + col_st[None, :].to(tl.int64) * stride_interm_r + + offs_d_st[:, None].to(tl.int64) + ) + gl.amd.cdna4.buffer_store( + stored_value=dKV.to(Interm_ptr.dtype.element_ty), + ptr=Interm_ptr, + offsets=interm_offs.to(tl.int32), + ) + + +def sparse_mla_bwd_dkv_interm_v4_bd( + q, + do, + chunk_dS, + chunk_P, + R_CHUNK, + BD=256, + TILE_K=128, + MFMA_K=32, + DUAL_STAGE=1, + PREFETCH=0, + H_POW2=None, + num_warps=4, + interm=None, +): + """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. + + Defaults measured at T=4096 H=128 topk=512 (MI355X): 1.170 ms vs 1.631 for + ``dkv_interm_v4`` = 1.39x. Sweep notes: + * BD=256 beats 128 (traffic: dS/P costs D/BD x). + * MFMA_K=32 is worth ~14% at the best config -- it did NOT matter in the old kernel, + which was bandwidth-saturated; it does here. + * PREFETCH=0: prefetching dS/P helps at BD=128 but is consistently WORSE at BD=256 + (1.181 -> 1.306), where the extra live registers cost occupancy. + * DUAL_STAGE=1 (both prologue copies behind one drain) is a small consistent win. + """ + T, H, D = q.shape + assert R_CHUNK % TILE_K == 0 + assert D % BD == 0 + h_pow2 = H_POW2 or triton.next_power_of_2(H) + if interm is None: + interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) + _dkv_interm_v4_bd_kernel[(T, D // BD)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + H, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + NH=h_pow2, + BD=BD, + D=D, + MFMA_K=MFMA_K, + DUAL_STAGE=DUAL_STAGE, + PREFETCH=PREFETCH, + num_warps=num_warps, + ) + return interm diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..fe71d0dd33 --- /dev/null +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness for the DSv4 sparse-MLA training backward against torch autograd. + +The reference is a differentiable fp32 re-implementation of the V4 forward, differentiated by +autograd -- an independent path from the kernels, not a re-expression of them. It is a per-token +python loop, so the shapes here are deliberately small; the kernels are exercised at production +shapes (T=4096, H=128, topk=512) out of tree. +""" + +import pytest +import torch + +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 +from aiter.ops.triton.utils._triton import arch_info + +D = 512 +COS_TOL = 0.999 + +pytestmark = pytest.mark.skipif( + arch_info.get_arch() != "gfx950", + reason="DSv4 sparse-MLA backward is gfx950 (CDNA4) only", +) + + +def _ref_fwd_diff(q, kv, attn_sink, topk, scale): + """Differentiable fp32 V4 forward. q/kv/attn_sink require grad. Returns O [T, H, D].""" + outs = [] + for t in range(q.shape[0]): + idx = topk[t].long() + valid = idx != -1 + k = kv[idx.clamp(min=0)] + k = torch.where(valid[:, None], k, torch.zeros_like(k)) + s = (q[t] @ k.t()) * scale + s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) + if attn_sink is None: + m = s.max(dim=1).values + p = torch.where( + valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) + ) + denom = p.sum(dim=1) + else: + m = torch.maximum(s.max(dim=1).values, attn_sink) + p = torch.where( + valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) + ) + denom = p.sum(dim=1) + torch.exp(attn_sink - m) + outs.append((p @ k) / denom[:, None]) + return torch.stack(outs, dim=0) + + +def _ref_fwd(q, kv, attn_sink, topk, scale): + """Non-differentiable forward giving the kernel inputs O (bf16) and sink-inclusive lse.""" + with torch.no_grad(): + o = _ref_fwd_diff(q.float(), kv.float(), attn_sink, topk, scale) + lse = torch.empty(q.shape[0], q.shape[1], device=q.device, dtype=torch.float32) + for t in range(q.shape[0]): + idx = topk[t].long() + valid = idx != -1 + k = kv.float()[idx.clamp(min=0)] + k = torch.where(valid[:, None], k, torch.zeros_like(k)) + s = (q.float()[t] @ k.t()) * scale + s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) + m = s.max(dim=1).values + p = torch.where( + valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) + ) + denom = p.sum(dim=1) + if attn_sink is not None: + denom = denom + torch.exp(attn_sink - m) + lse[t] = m + torch.log(denom) + return o.to(torch.bfloat16).contiguous(), lse + + +def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.float().reshape(-1), b.float().reshape(-1), dim=0 + ).item() + + +@pytest.mark.parametrize( + "T, H, topk, npool, has_sink, r_chunk", + [ + (128, 64, 128, 0, True, None), # no pool, unchunked + (128, 128, 128, 0, True, None), # H=128 + (128, 128, 128, 64, True, None), # compressed pool (num_kv > T) + (128, 64, 128, 0, False, None), # no attn_sink + (128, 128, 128, 64, True, 64), # chunked: dQ RMW + per-chunk CSR build + ], +) +def test_sparse_mla_bwd_dsv4(T, H, topk, npool, has_sink, r_chunk): + torch.manual_seed(0) + dev = "cuda" + num_kv = T + npool + scale = 1.0 / (D**0.5) + + q = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + kv = torch.randn(num_kv, D, device=dev, dtype=torch.bfloat16) + do = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + sink = (torch.randn(H, device=dev, dtype=torch.float32) * 0.1) if has_sink else None + + indices = torch.randint(0, num_kv, (T, topk), dtype=torch.int32, device=dev) + invalid = torch.rand(T, topk, device=dev) < 0.1 + indices = torch.where(invalid, torch.full_like(indices, -1), indices).contiguous() + + o, lse = _ref_fwd(q, kv, sink, indices, scale) + + dq, dkv, d_sink = sparse_mla_bwd_dsv4( + q, kv, do, o, lse, indices, attn_sink=sink, scale=scale, R_CHUNK=r_chunk + ) + + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) if has_sink else None + _ref_fwd_diff(qg, kvg, sg, indices, scale).backward(do.float()) + + assert _cos(dq, qg.grad) > COS_TOL, f"dq cos {_cos(dq, qg.grad)}" + assert _cos(dkv, kvg.grad) > COS_TOL, f"dkv cos {_cos(dkv, kvg.grad)}" + if has_sink: + assert _cos(d_sink, sg.grad) > COS_TOL, f"d_sink cos {_cos(d_sink, sg.grad)}" + else: + assert d_sink is None + + +def test_sparse_mla_bwd_dsv4_rejects_bad_chunk(): + """R_CHUNK must be a multiple of the mfma tile width.""" + dev = "cuda" + T, H, topk = 64, 64, 64 + q = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + kv = torch.randn(T, D, device=dev, dtype=torch.bfloat16) + do = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + o = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) + lse = torch.randn(T, H, device=dev, dtype=torch.float32) + idx = torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev) + with pytest.raises(AssertionError, match="multiple of 32"): + sparse_mla_bwd_dsv4(q, kv, do, o, lse, idx, R_CHUNK=48) From 21b8588af0434ff95cb9802a259cd3fb31dec329 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 11:14:11 -0500 Subject: [PATCH 02/15] Drop the _be suffix from the dKV gather The suffix distinguished this kernel from an earlier one-entry-at-a-time gather that is not in this tree, so in-tree it only reads as an unexplained abbreviation. `_bwd_dkv_gather_acc_v4_be` -> `_bwd_dkv_gather_acc_v4`, `dkv_gather_acc_be` -> `dkv_gather_acc`. The kernel docstring compared itself against that absent kernel, which after the rename would have been self-referential; it now states the same reasoning as a design note -- carrying BLOCK_E entries per iteration widens the load to dwordx4 and cuts the trip count on the long runs a realistic top-k produces, and the gather needs both. No functional change. 6/6 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 27 ++++++++++--------- .../attention/sparse_attention_dsv4_bwd.py | 4 +-- 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py index cafe9f75b6..630ae5d371 100644 --- a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py @@ -7,7 +7,7 @@ ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the bf16 inputs and accumulates in fp32, so it moves exactly the working set. -``_bwd_dkv_gather_acc_v4_be`` + ``build_inverted_topk_fast`` +``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk_fast`` Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is inverted into a CSR gather (each output KV row collects its own contributors), so no atomics are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load @@ -66,7 +66,7 @@ def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): @triton.jit -def _bwd_dkv_gather_acc_v4_be( +def _bwd_dkv_gather_acc_v4( Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token @@ -79,15 +79,15 @@ def _bwd_dkv_gather_acc_v4_be( ): """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. - Fixes two things about ``_bwd_dkv_gather_acc_v4``, which walks the run one entry at a time - with a bare ``tl.arange(0, D)``: + Carrying ``BLOCK_E`` entries per iteration rather than walking the run one entry at a time + buys two separate things, and this gather needs both: - * **load width.** A [D] block over 256 threads is 2 bf16 = 4 B per lane -- a dword. The - [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane instead, so the loads - become dwordx4. The gather is issue-bound, so this is the dominant term. - * **loop trip count.** The run is consumed BLOCK_E entries at a time rather than one, and - ``tl.sum`` over the entry axis folds them. The realistic topk gives run lengths up to - ~3000 (pool rows), so the serial walk was the other half of the problem. + * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per + lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane + instead, so the loads become dwordx4. The gather is issue-bound, so this dominates. + * **loop trip count.** The run is consumed BLOCK_E entries at a time and ``tl.sum`` over + the entry axis folds them. A realistic top-k gives run lengths up to ~3000 (pool rows), + so the serial walk was the other half of the problem. ``ACCUMULATE=False`` skips the read-modify-write of the destination, valid when the caller does not chunk (each KV row is then written by exactly one CTA). @@ -145,20 +145,21 @@ def build_inverted_topk_fast(topk_indices_slice, num_kv): return inv_ptr, inv_data.to(torch.int32) -def dkv_gather_acc_be( +def dkv_gather_acc( interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True ): # BLOCK_E=64 / num_warps=8 measured best (0.345 ms, 6.29 TB/s = 79% peak at T4096 H128 # topk512 SWA+pool). Time falls monotonically with BLOCK_E across the whole sweep # (4->64: 1.021, 0.715, 0.505, 0.407, 0.345), i.e. both the load width AND the trip count - # on the ~3000-entry pool runs were binding. Reference was 1.491 ms at 1.45 TB/s. + # on the ~3000-entry pool runs were binding. The one-entry-at-a-time walk this replaced + # was 1.491 ms at 1.45 TB/s. """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. """ _, _, D = interm.shape num_kv = dkv_acc.shape[0] - _bwd_dkv_gather_acc_v4_be[(num_kv,)]( + _bwd_dkv_gather_acc_v4[(num_kv,)]( interm, inv_ptr, inv_data, diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 603cb31275..c344df5e2e 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -42,7 +42,7 @@ from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( build_inverted_topk_fast, delta_v4, - dkv_gather_acc_be, + dkv_gather_acc, ) from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( sparse_mla_bwd_dkv_interm_v4_bd as _dkv_interm_gluon, @@ -153,7 +153,7 @@ def sparse_mla_bwd_dsv4( inv_ptr, inv_data = build_inverted_topk_fast( topk_indices[:, r : r + R_CHUNK], num_kv ) - dkv_gather_acc_be(interm, inv_ptr, inv_data, dkv_acc) + dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc) dkv = dkv_acc.to(kv.dtype) From 2ba7a9bd5218f76a6ee9b7e74c66b2730b486fa9 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 11:27:41 -0500 Subject: [PATCH 03/15] Name things after what they are, not what they were Three symbols carried a suffix that only made sense against an earlier variant that was never part of this tree, so in-tree they read as unexplained abbreviations: build_inverted_topk_fast -> build_inverted_topk _dkv_interm_v4_bd_kernel -> _dkv_interm_v4_kernel sparse_mla_bwd_dkv_interm_v4_bd -> sparse_mla_bwd_dkv_interm_v4 Two of their docstrings pointed at those absent variants ("bit-identical to the reference below", "vs 1.631 for dkv_interm_v4", "did NOT matter in the old kernel"). Nothing below or elsewhere in the tree matched, so the comparisons are restated as standalone reasoning: why int16 keys plus searchsorted are the right formulation, and why MFMA_K only pays once the D split takes the kernel off the bandwidth ceiling. Also `ar` -> `row_ids` in the CSR build, with a note that its dtype must match the sort key -- that is why it is constructed inside each branch rather than once above them. No functional change. 6/6 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 20 +++++++++++-------- .../attention/sparse_attention_dsv4_bwd.py | 6 +++--- .../gluon/sparse_attention_dsv4_bwd_gluon.py | 15 +++++++------- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py index 630ae5d371..4272f25e5d 100644 --- a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py @@ -7,7 +7,7 @@ ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the bf16 inputs and accumulates in fp32, so it moves exactly the working set. -``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk_fast`` +``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk`` Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is inverted into a CSR gather (each output KV row collects its own contributors), so no atomics are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load @@ -118,30 +118,34 @@ def _bwd_dkv_gather_acc_v4( tl.store(dKV_acc_ptr + acc_base + offs_d, acc) -def build_inverted_topk_fast(topk_indices_slice, num_kv): - """CSR inverted index over ``num_kv`` KV rows. Bit-identical to the reference below. +def build_inverted_topk(topk_indices_slice, num_kv): + """CSR inverted index over ``num_kv`` KV rows. One stable sort yields both the permutation (``inv_data``) and the sorted keys; ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k, which is exactly what ``cumsum(bincount(flat+1))`` computes. Invalid (-1) entries sort to the front, so ``inv_ptr[0]`` starts past them and they are never visited. - Two things make this ~3x faster than the reference: + Two details carry most of the cost, and the obvious formulation gets both wrong -- writing + this as ``bincount`` + ``cumsum`` over int64 keys measured 3x slower for the same output: * the sort key is narrowed to int16 when ``num_kv`` fits, so the radix sort makes 2 byte-passes instead of 8; - * ``searchsorted`` replaces the separate ``bincount`` + ``cumsum`` passes. + * ``searchsorted`` does the job of the separate ``bincount`` + ``cumsum`` passes. Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. """ + # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. + # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by + # accident. flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid if num_kv < 32767: # int16 range, -1 included keys = flat_kv.to(torch.int16) - ar = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) else: keys = flat_kv.to(torch.int32) - ar = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) sorted_vals, inv_data = torch.sort(keys, stable=True) - inv_ptr = torch.searchsorted(sorted_vals, ar).to(torch.int32) + inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) return inv_ptr, inv_data.to(torch.int32) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index c344df5e2e..6048af67d2 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -40,12 +40,12 @@ import torch from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( - build_inverted_topk_fast, + build_inverted_topk, delta_v4, dkv_gather_acc, ) from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( - sparse_mla_bwd_dkv_interm_v4_bd as _dkv_interm_gluon, + sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, ) from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( sparse_mla_bwd_dq_gluon as _dq_gluon, @@ -150,7 +150,7 @@ def sparse_mla_bwd_dsv4( _dkv_interm_gluon( q, do, chunk_dS, chunk_P, R_CHUNK, BD=_BD_DKV, TILE_K=tk_dkv, interm=interm ) - inv_ptr, inv_data = build_inverted_topk_fast( + inv_ptr, inv_data = build_inverted_topk( topk_indices[:, r : r + R_CHUNK], num_kv ) dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc) diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py index 16bffd4273..6b7eea0064 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py @@ -13,7 +13,7 @@ is read from LDS once and feeds both the ``S`` and ``dP`` MFMAs. Also emits the ``dS`` / ``P`` chunks the dKV-interm kernel consumes. -``_dkv_interm_v4_bd_kernel`` +``_dkv_interm_v4_kernel`` ``interm[t, slot, d] = sum_h ( dS[t,h,slot]*Q[t,h,d] + P[t,h,slot]*dO[t,h,d] )``, contracting over ALL heads inside one MFMA pair so nothing accumulates across a loop over heads. Q and dO are transposed once into registers and D is split across ``grid.y``, which is what keeps them @@ -392,7 +392,7 @@ def sparse_mla_bwd_dq_gluon( @gluon.jit -def _dkv_interm_v4_bd_kernel( +def _dkv_interm_v4_kernel( Q_ptr, # [T, H, D] bf16 dO_ptr, # [T, H, D] bf16 dS_ptr, # [T, H, R_CHUNK] bf16 @@ -607,7 +607,7 @@ def _dkv_interm_v4_bd_kernel( ) -def sparse_mla_bwd_dkv_interm_v4_bd( +def sparse_mla_bwd_dkv_interm_v4( q, do, chunk_dS, @@ -624,11 +624,10 @@ def sparse_mla_bwd_dkv_interm_v4_bd( ): """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. - Defaults measured at T=4096 H=128 topk=512 (MI355X): 1.170 ms vs 1.631 for - ``dkv_interm_v4`` = 1.39x. Sweep notes: + Defaults measured at T=4096 H=128 topk=512 on gfx950 (MI355X): 1.170 ms. Sweep notes: * BD=256 beats 128 (traffic: dS/P costs D/BD x). - * MFMA_K=32 is worth ~14% at the best config -- it did NOT matter in the old kernel, - which was bandwidth-saturated; it does here. + * MFMA_K=32 is worth ~14% at the best config. It only pays once the kernel is off the + bandwidth ceiling, which is what splitting D across grid.y buys. * PREFETCH=0: prefetching dS/P helps at BD=128 but is consistently WORSE at BD=256 (1.181 -> 1.306), where the extra live registers cost occupancy. * DUAL_STAGE=1 (both prologue copies behind one drain) is a small consistent win. @@ -639,7 +638,7 @@ def sparse_mla_bwd_dkv_interm_v4_bd( h_pow2 = H_POW2 or triton.next_power_of_2(H) if interm is None: interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) - _dkv_interm_v4_bd_kernel[(T, D // BD)]( + _dkv_interm_v4_kernel[(T, D // BD)]( q, do, chunk_dS, From 444524696e4cc767527195d9011298dacabd99c6 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 11:41:18 -0500 Subject: [PATCH 04/15] Fold the Triton kernels into the Gluon module; drop the _gluon suffix `_triton_kernels/` holds the Triton *implementation* of an op, which a public wrapper can select or fall back to -- `attention/pa_mqa_logits.py` picks between the Gluon kernel and the `_triton_kernels/` one at import time. Our delta and CSR-gather kernels are not that: there is no Triton backward to fall back to, they are the memory-bound phases of this one kernel. Filing them under `_triton_kernels/attention/sparse_attention_dsv4_bwd.py` put them directly next to `_triton_kernels/attention/sparse_attention_dsv4.py`, which IS the forward's Triton path, so the pair read as a Triton fwd/bwd that does not exist. They now live beside the Gluon kernels they serve, which is what every other Gluon module here does -- mla_gluon, pa_decode_gluon and pa_mqa_logits all keep their plain @triton.jit helpers inline and none import from `_triton_kernels/`. The `_gluon` suffix goes with it: the file is already in `gluon/`, four of the seven modules there carry no suffix, and `pa_mqa_logits` is the same basename in all three trees without one. Same for `sparse_mla_bwd_dq_gluon` -> `sparse_mla_bwd_dq`, which its sibling never had. No functional change. 6/6 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 177 ----------------- .../attention/sparse_attention_dsv4_bwd.py | 8 +- aiter/ops/triton/gluon/README.md | 4 +- ..._gluon.py => sparse_attention_dsv4_bwd.py} | 179 +++++++++++++++++- 4 files changed, 178 insertions(+), 190 deletions(-) delete mode 100644 aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py rename aiter/ops/triton/gluon/{sparse_attention_dsv4_bwd_gluon.py => sparse_attention_dsv4_bwd.py} (74%) diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py deleted file mode 100644 index 4272f25e5d..0000000000 --- a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py +++ /dev/null @@ -1,177 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Triton kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). - -``_delta_v4_kernel`` - ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the - bf16 inputs and accumulates in fp32, so it moves exactly the working set. - -``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk`` - Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is - inverted into a CSR gather (each output KV row collects its own contributors), so no atomics - are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load - and cuts the trip count on the long runs a realistic top-k produces. - -Public entry: ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4``. -""" - -import torch -import triton -import triton.language as tl - - -@triton.jit -def _delta_v4_kernel( - O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) - dO_ptr, # [n_rows, D] bf16 - Delta_ptr, # [n_rows] fp32 - n_rows, - D: tl.constexpr, - BLOCK_R: tl.constexpr, -): - """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" - pid = tl.program_id(0) - rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) - mask = rows < n_rows - offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] - o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) - d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) - tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) - - -def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): - # BLOCK_R=8 / num_warps=8 measured best (0.173 ms, 6.21 TB/s = 78% peak at T4096 H128); - # the whole sweep plateaus at 0.173-0.187 once a lane loads >= 8 bf16, i.e. once the load - # is a dwordx4. Below that (BLOCK_R=2 nw=8, 2 bf16/lane) it falls off a cliff to 3.12 TB/s. - """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. - - ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. - """ - assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() - T, H, D = o.shape - n_rows = T * H - if out is None: - out = torch.empty(T, H, dtype=torch.float32, device=o.device) - _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( - o, - do, - out, - n_rows, - D=D, - BLOCK_R=BLOCK_R, - num_warps=num_warps, - ) - return out - - -@triton.jit -def _bwd_dkv_gather_acc_v4( - Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] - InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers - InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token - dKV_acc_ptr, # [num_kv, D] fp32 — accumulator - stride_interm_r: tl.int64, - stride_acc_t: tl.int64, - D: tl.constexpr, - BLOCK_E: tl.constexpr, - ACCUMULATE: tl.constexpr, -): - """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. - - Carrying ``BLOCK_E`` entries per iteration rather than walking the run one entry at a time - buys two separate things, and this gather needs both: - - * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per - lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane - instead, so the loads become dwordx4. The gather is issue-bound, so this dominates. - * **loop trip count.** The run is consumed BLOCK_E entries at a time and ``tl.sum`` over - the entry axis folds them. A realistic top-k gives run lengths up to ~3000 (pool rows), - so the serial walk was the other half of the problem. - - ``ACCUMULATE=False`` skips the read-modify-write of the destination, valid when the caller - does not chunk (each KV row is then written by exactly one CTA). - """ - k = tl.program_id(0) - offs_d = tl.arange(0, D) - offs_e = tl.arange(0, BLOCK_E) - start = tl.load(InvPtr_ptr + k) - end = tl.load(InvPtr_ptr + k + 1) - acc_base = k.to(tl.int64) * stride_acc_t - - if ACCUMULATE: - acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) - else: - acc = tl.zeros([D], dtype=tl.float32) - - for i0 in range(start, end, BLOCK_E): - idx = i0 + offs_e - m = idx < end - entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) - vals = tl.load( - Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], - mask=m[:, None], - other=0.0, - ) - acc += tl.sum(vals.to(tl.float32), axis=0) - - tl.store(dKV_acc_ptr + acc_base + offs_d, acc) - - -def build_inverted_topk(topk_indices_slice, num_kv): - """CSR inverted index over ``num_kv`` KV rows. - - One stable sort yields both the permutation (``inv_data``) and the sorted keys; - ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k, - which is exactly what ``cumsum(bincount(flat+1))`` computes. Invalid (-1) entries sort to - the front, so ``inv_ptr[0]`` starts past them and they are never visited. - - Two details carry most of the cost, and the obvious formulation gets both wrong -- writing - this as ``bincount`` + ``cumsum`` over int64 keys measured 3x slower for the same output: - * the sort key is narrowed to int16 when ``num_kv`` fits, so the radix sort makes 2 - byte-passes instead of 8; - * ``searchsorted`` does the job of the separate ``bincount`` + ``cumsum`` passes. - - Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. - """ - # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. - # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by - # accident. - flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid - if num_kv < 32767: # int16 range, -1 included - keys = flat_kv.to(torch.int16) - row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) - else: - keys = flat_kv.to(torch.int32) - row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) - sorted_vals, inv_data = torch.sort(keys, stable=True) - inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) - return inv_ptr, inv_data.to(torch.int32) - - -def dkv_gather_acc( - interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True -): - # BLOCK_E=64 / num_warps=8 measured best (0.345 ms, 6.29 TB/s = 79% peak at T4096 H128 - # topk512 SWA+pool). Time falls monotonically with BLOCK_E across the whole sweep - # (4->64: 1.021, 0.715, 0.505, 0.407, 0.345), i.e. both the load width AND the trip count - # on the ~3000-entry pool runs were binding. The one-entry-at-a-time walk this replaced - # was 1.491 ms at 1.45 TB/s. - """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. - - Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. - """ - _, _, D = interm.shape - num_kv = dkv_acc.shape[0] - _bwd_dkv_gather_acc_v4[(num_kv,)]( - interm, - inv_ptr, - inv_data, - dkv_acc, - interm.stride(1), - dkv_acc.stride(0), - D=D, - BLOCK_E=BLOCK_E, - ACCUMULATE=accumulate, - num_warps=num_warps, - ) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 6048af67d2..b4c1a98f1a 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -39,16 +39,12 @@ import torch -from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( build_inverted_topk, delta_v4, dkv_gather_acc, -) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, -) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd_gluon import ( - sparse_mla_bwd_dq_gluon as _dq_gluon, + sparse_mla_bwd_dq as _dq_gluon, ) from aiter.ops.triton.utils._triton import arch_info diff --git a/aiter/ops/triton/gluon/README.md b/aiter/ops/triton/gluon/README.md index fd693926b9..7c0096bc95 100644 --- a/aiter/ops/triton/gluon/README.md +++ b/aiter/ops/triton/gluon/README.md @@ -58,7 +58,7 @@ Some features (e.g., scheduling hints like `sched_barrier`) require the [AMD Glu TBDTBDTBD - sparse_attention_
dsv4_bwd_gluonDSv4 Sparse
MLA BackwardCDNA4 + sparse_attention_
dsv4_bwdDSv4 Sparse
MLA BackwardCDNA4 Q/KV/dO/O: bf16, K == V
head_dim = 512 (dense)
lse: fp32, sink-inclusive
num_kv ≥ T (pool ok)
topk % 32 == 0
gfx950 only python op_tests/triton_tests/
attention/test_sparse_
attention_dsv4_bwd.py ~407
TFLOPS—— @@ -253,7 +253,7 @@ python op_tests/test_mla.py -c 10000 100000 -b 1 3 4 -n 16,1 -d bf16 -kvd bf16 - | 100K | 3 | 85 | 19 | 88.77 | 3.89 | | 100K | 4 | 64 | 25 | 106.96 | 4.31 | -### `sparse_attention_dsv4_bwd_gluon.py` — DeepSeek V4 Sparse MLA Backward +### `sparse_attention_dsv4_bwd.py` — DeepSeek V4 Sparse MLA Backward **Public entry:** `aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None)` -> `(dq, dkv, d_sink)` diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py similarity index 74% rename from aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py rename to aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py index 6b7eea0064..e30f5611a5 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Gluon kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). +"""Kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). -Two kernels, both operating on the official V4 form (``K == V == kv``, one dense 512-wide -tensor, RoPE already applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` in the -softmax denominator only, ``topk == -1`` masked): +All operate on the official V4 form (``K == V == kv``, one dense 512-wide tensor, RoPE already +applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` in the softmax denominator +only, ``topk == -1`` masked). The two MFMA phases are Gluon; the two memory-bound phases are +plain Triton and live here rather than under ``_triton_kernels/`` because there is no Triton +implementation of this backward to fall back to -- they are parts of this kernel, not an +alternative to it. ``_dq_v4_kernel`` Per (query token, head block): ``S = Q@kv^T``, ``P = exp(S - lse)``, ``dP = dO@kv^T``, @@ -19,6 +22,16 @@ are transposed once into registers and D is split across ``grid.y``, which is what keeps them read once instead of ``topk/TILE_K`` times. +``_delta_v4_kernel`` + ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the + bf16 inputs and accumulates in fp32, so it moves exactly the working set. + +``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk`` + Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is + inverted into a CSR gather (each output KV row collects its own contributors), so no atomics + are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load + and cuts the trip count on the long runs a realistic top-k produces. + Public entry: ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4``. """ @@ -338,7 +351,7 @@ def _dq_v4_kernel( ) -def sparse_mla_bwd_dq_gluon( +def sparse_mla_bwd_dq( q, kv, do, @@ -664,3 +677,159 @@ def sparse_mla_bwd_dkv_interm_v4( num_warps=num_warps, ) return interm + + +@triton.jit +def _delta_v4_kernel( + O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) + dO_ptr, # [n_rows, D] bf16 + Delta_ptr, # [n_rows] fp32 + n_rows, + D: tl.constexpr, + BLOCK_R: tl.constexpr, +): + """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + mask = rows < n_rows + offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] + o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) + + +def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): + # BLOCK_R=8 / num_warps=8 measured best (0.173 ms, 6.21 TB/s = 78% peak at T4096 H128); + # the whole sweep plateaus at 0.173-0.187 once a lane loads >= 8 bf16, i.e. once the load + # is a dwordx4. Below that (BLOCK_R=2 nw=8, 2 bf16/lane) it falls off a cliff to 3.12 TB/s. + """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. + + ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. + """ + assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() + T, H, D = o.shape + n_rows = T * H + if out is None: + out = torch.empty(T, H, dtype=torch.float32, device=o.device) + _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( + o, + do, + out, + n_rows, + D=D, + BLOCK_R=BLOCK_R, + num_warps=num_warps, + ) + return out + + +@triton.jit +def _bwd_dkv_gather_acc_v4( + Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] + InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers + InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token + dKV_acc_ptr, # [num_kv, D] fp32 — accumulator + stride_interm_r: tl.int64, + stride_acc_t: tl.int64, + D: tl.constexpr, + BLOCK_E: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. + + Carrying ``BLOCK_E`` entries per iteration rather than walking the run one entry at a time + buys two separate things, and this gather needs both: + + * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per + lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane + instead, so the loads become dwordx4. The gather is issue-bound, so this dominates. + * **loop trip count.** The run is consumed BLOCK_E entries at a time and ``tl.sum`` over + the entry axis folds them. A realistic top-k gives run lengths up to ~3000 (pool rows), + so the serial walk was the other half of the problem. + + ``ACCUMULATE=False`` skips the read-modify-write of the destination, valid when the caller + does not chunk (each KV row is then written by exactly one CTA). + """ + k = tl.program_id(0) + offs_d = tl.arange(0, D) + offs_e = tl.arange(0, BLOCK_E) + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + acc_base = k.to(tl.int64) * stride_acc_t + + if ACCUMULATE: + acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) + else: + acc = tl.zeros([D], dtype=tl.float32) + + for i0 in range(start, end, BLOCK_E): + idx = i0 + offs_e + m = idx < end + entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) + vals = tl.load( + Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], + mask=m[:, None], + other=0.0, + ) + acc += tl.sum(vals.to(tl.float32), axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_d, acc) + + +def build_inverted_topk(topk_indices_slice, num_kv): + """CSR inverted index over ``num_kv`` KV rows. + + One stable sort yields both the permutation (``inv_data``) and the sorted keys; + ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k, + which is exactly what ``cumsum(bincount(flat+1))`` computes. Invalid (-1) entries sort to + the front, so ``inv_ptr[0]`` starts past them and they are never visited. + + Two details carry most of the cost, and the obvious formulation gets both wrong -- writing + this as ``bincount`` + ``cumsum`` over int64 keys measured 3x slower for the same output: + * the sort key is narrowed to int16 when ``num_kv`` fits, so the radix sort makes 2 + byte-passes instead of 8; + * ``searchsorted`` does the job of the separate ``bincount`` + ``cumsum`` passes. + + Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. + """ + # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. + # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by + # accident. + flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid + if num_kv < 32767: # int16 range, -1 included + keys = flat_kv.to(torch.int16) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) + else: + keys = flat_kv.to(torch.int32) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) + sorted_vals, inv_data = torch.sort(keys, stable=True) + inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) + return inv_ptr, inv_data.to(torch.int32) + + +def dkv_gather_acc( + interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True +): + # BLOCK_E=64 / num_warps=8 measured best (0.345 ms, 6.29 TB/s = 79% peak at T4096 H128 + # topk512 SWA+pool). Time falls monotonically with BLOCK_E across the whole sweep + # (4->64: 1.021, 0.715, 0.505, 0.407, 0.345), i.e. both the load width AND the trip count + # on the ~3000-entry pool runs were binding. The one-entry-at-a-time walk this replaced + # was 1.491 ms at 1.45 TB/s. + """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. + + Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. + """ + _, _, D = interm.shape + num_kv = dkv_acc.shape[0] + _bwd_dkv_gather_acc_v4[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D=D, + BLOCK_E=BLOCK_E, + ACCUMULATE=accumulate, + num_warps=num_warps, + ) From 569edf21370bae7949b3b6abb8dc9a3b81b15306 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 11:59:58 -0500 Subject: [PATCH 05/15] Reject an R_CHUNK that does not divide TOPK R_CHUNK was only checked for being a multiple of 32, so a value like 192 against TOPK=512 was accepted and left a 128-wide tail chunk. That is wrong in three places at once, and none of them raise: * both Gluon kernels take the chunk width as a constexpr and mask the top-k load against R_CHUNK rather than against TOPK, so the tail iteration reads past the end of each row -- into the next token's indices for tokens 0..T-2, and past the tensor entirely for the last one, which faults the GPU (measured at T=128 H=64 topk=128 R_CHUNK=96); * the CSR build is handed `topk_indices[:, r:r+R_CHUNK]`, which torch clamps to the narrower tail, so its entries are encoded against that width; * the gather consumes those entries against interm, which is still R_CHUNK wide, so every entry past the first token addresses the wrong row. The default path (R_CHUNK=None -> TOPK) and the previously tested chunked case both divide, so this never surfaced. Handling a narrower tail instead would cost a second kernel compile per shape for a knob whose only job is bounding the interm buffer, so the divisor is required. Adds a test for the case that clears the tile-width check but not this one (96 against TOPK=128). 7/7 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 16 +++++++- .../test_sparse_attention_dsv4_bwd.py | 38 ++++++++++++++----- 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index b4c1a98f1a..b796565401 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -81,7 +81,9 @@ def sparse_mla_bwd_dsv4( runs unchunked, which is what you want. Chunking exists only to bound the ``interm`` intermediate, which is ``T*TOPK*512`` bf16 (2.0 GiB at T=4096, TOPK=512); it costs a dQ read-modify-write between chunks and one - CSR build per chunk. Any multiple of 32 is accepted. + CSR build per chunk. Must be a multiple of 32 (the mfma tile width) and + must divide ``TOPK`` -- a partial tail chunk is rejected rather than + handled, since the chunk width is a kernel constexpr. Returns: dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) @@ -117,6 +119,18 @@ def sparse_mla_bwd_dsv4( tk_dkv is not None and tk_dq is not None ), f"R_CHUNK={R_CHUNK} must be a multiple of 32 (it is the mfma tile width)" + # A tail chunk narrower than R_CHUNK is wrong in three places at once. The kernels take + # R_CHUNK as a constexpr and mask the top-k load against it rather than against TOPK, so + # they read past the end of each row -- into the next token's indices for tokens 0..T-2, + # and past the tensor entirely for the last one (measured: that faults the GPU). The CSR + # build meanwhile gets a torch-clamped, narrower slice, and the gather then indexes interm, + # which is still R_CHUNK wide, with entries encoded against that narrower width. Require + # the divisor rather than paying a second compile for a narrower tail. + assert TOPK % R_CHUNK == 0, ( + f"R_CHUNK={R_CHUNK} must divide TOPK={TOPK} -- a partial tail chunk reads past the end " + "of each top-k row and desynchronizes the CSR index space from interm" + ) + delta = delta_v4(o, do) dq = torch.empty_like(q) diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py index fe71d0dd33..93cd133028 100644 --- a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -123,15 +123,35 @@ def test_sparse_mla_bwd_dsv4(T, H, topk, npool, has_sink, r_chunk): assert d_sink is None +def _dummy_inputs(T=64, H=64, topk=64, dev="cuda"): + return dict( + q=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + kv=torch.randn(T, D, device=dev, dtype=torch.bfloat16), + do=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + o=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + lse=torch.randn(T, H, device=dev, dtype=torch.float32), + idx=torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev), + ) + + def test_sparse_mla_bwd_dsv4_rejects_bad_chunk(): """R_CHUNK must be a multiple of the mfma tile width.""" - dev = "cuda" - T, H, topk = 64, 64, 64 - q = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) - kv = torch.randn(T, D, device=dev, dtype=torch.bfloat16) - do = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) - o = torch.randn(T, H, D, device=dev, dtype=torch.bfloat16) - lse = torch.randn(T, H, device=dev, dtype=torch.float32) - idx = torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev) + t = _dummy_inputs(topk=64) with pytest.raises(AssertionError, match="multiple of 32"): - sparse_mla_bwd_dsv4(q, kv, do, o, lse, idx, R_CHUNK=48) + sparse_mla_bwd_dsv4( + t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"], R_CHUNK=48 + ) + + +def test_sparse_mla_bwd_dsv4_rejects_indivisible_chunk(): + """A chunk width that is a valid tile multiple but does not divide TOPK is still rejected. + + 96 = 3*32 so it clears the tile-width check, but 128 % 96 == 32 would leave a 32-wide tail + chunk. The kernels take the chunk width as a constexpr and would read past the end of each + top-k row, silently, so this is an error rather than a handled case. + """ + t = _dummy_inputs(topk=128) + with pytest.raises(AssertionError, match="must divide TOPK"): + sparse_mla_bwd_dsv4( + t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"], R_CHUNK=96 + ) From 747abc8c2cf6351304bdbc34edd94b8335d3df2d Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 12:13:11 -0500 Subject: [PATCH 06/15] Add kernel reprs and fill in the input validation Both from review feedback on the PR. Kernel reprs follow `_gluon_kernels/gfx950/attention/pa_decode_sparse.py`: a `make_kernel_repr` object per kernel naming its constexprs, passed as `repr=` to the jit decorator. Note the DSv4 forward does not do this -- neither the Triton path nor `mla_gluon` -- but the same-arch Gluon decode does and so do ~20 other kernel modules, so this reads as an omission there rather than a deliberate choice. It earns its keep on a kernel whose PR is about performance: a trace now says which tile configuration ran instead of leaving it to be inferred. _dq_v4_kernel_R_CHUNK_512_BLOCK_H_64_TILE_K_32_D_512_IS_FIRST_CHUNK_1 _delta_v4_kernel_D_512_BLOCK_R_8 Validation follows `attention/pa_prefill_sparse.py`, which is the forward's public entry: RuntimeError with the offending value for dtype problems, assert for shape and contiguity. The compound asserts are split so a failure says which tensor was wrong, and `lse`, `attn_sink` and the `topk_indices` dtype are now checked at all -- previously nothing looked at them. 7/7 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 30 +++++++++- .../triton/gluon/sparse_attention_dsv4_bwd.py | 56 +++++++++++++++++-- 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index b796565401..f30b8c8683 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -92,14 +92,38 @@ def sparse_mla_bwd_dsv4( arch_info.get_arch() == "gfx950" ), f"sparse_mla_bwd_dsv4 requires gfx950 (CDNA4), got {arch_info.get_arch()}" + if q.dtype != torch.bfloat16: + raise RuntimeError(f"sparse_mla_bwd_dsv4 expects bf16 q, got {q.dtype}") + for name, t in (("kv", kv), ("do", do), ("o", o)): + if t.dtype != q.dtype: + raise RuntimeError(f"{name} dtype mismatch: {name}={t.dtype}, q={q.dtype}") + T, H, D = q.shape TOPK = topk_indices.shape[1] num_kv = kv.shape[0] assert D == 512, f"DSv4 sparse-MLA backward is fixed to head_dim 512, got {D}" - assert kv.shape[-1] == D and do.shape == q.shape and o.shape == q.shape + assert kv.shape[-1] == D, f"kv must be [num_kv, {D}], got {tuple(kv.shape)}" + assert ( + do.shape == q.shape + ), f"do must match q {tuple(q.shape)}, got {tuple(do.shape)}" + assert o.shape == q.shape, f"o must match q {tuple(q.shape)}, got {tuple(o.shape)}" + assert lse.shape == (T, H), f"lse must be [{T}, {H}], got {tuple(lse.shape)}" assert num_kv >= T, f"num_kv ({num_kv}) must be >= T ({T})" - assert q.is_contiguous() and kv.is_contiguous() and do.is_contiguous() - assert o.is_contiguous() and topk_indices.is_contiguous() + if attn_sink is not None: + assert attn_sink.shape == ( + H, + ), f"attn_sink must be [{H}], got {tuple(attn_sink.shape)}" + assert ( + topk_indices.dtype == torch.int32 + ), f"topk_indices must be int32, got {topk_indices.dtype}" + for name, t in ( + ("q", q), + ("kv", kv), + ("do", do), + ("o", o), + ("topk_indices", topk_indices), + ): + assert t.is_contiguous(), f"{name} must be contiguous" if scale is None: scale = 1.0 / (D**0.5) diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py index e30f5611a5..0ecd37555c 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py @@ -41,8 +41,22 @@ from triton.experimental import gluon from triton.experimental.gluon import language as gl +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr -@gluon.jit + +_dq_v4_kernel_repr = make_kernel_repr( + "_dq_v4_kernel", + [ + "R_CHUNK", + "BLOCK_H", + "TILE_K", + "D", + "IS_FIRST_CHUNK", + ], +) + + +@gluon.jit(repr=_dq_v4_kernel_repr) def _dq_v4_kernel( Q_ptr, # [T, H, D] bf16 KV_ptr, # [T, D] bf16 (K == V) @@ -404,7 +418,22 @@ def sparse_mla_bwd_dq( ) -@gluon.jit +_dkv_interm_v4_kernel_repr = make_kernel_repr( + "_dkv_interm_v4_kernel", + [ + "R_CHUNK", + "TILE_K", + "NH", + "BD", + "D", + "MFMA_K", + "DUAL_STAGE", + "PREFETCH", + ], +) + + +@gluon.jit(repr=_dkv_interm_v4_kernel_repr) def _dkv_interm_v4_kernel( Q_ptr, # [T, H, D] bf16 dO_ptr, # [T, H, D] bf16 @@ -679,7 +708,16 @@ def sparse_mla_bwd_dkv_interm_v4( return interm -@triton.jit +_delta_v4_kernel_repr = make_kernel_repr( + "_delta_v4_kernel", + [ + "D", + "BLOCK_R", + ], +) + + +@triton.jit(repr=_delta_v4_kernel_repr) def _delta_v4_kernel( O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) dO_ptr, # [n_rows, D] bf16 @@ -723,7 +761,17 @@ def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): return out -@triton.jit +_bwd_dkv_gather_acc_v4_repr = make_kernel_repr( + "_bwd_dkv_gather_acc_v4", + [ + "D", + "BLOCK_E", + "ACCUMULATE", + ], +) + + +@triton.jit(repr=_bwd_dkv_gather_acc_v4_repr) def _bwd_dkv_gather_acc_v4( Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers From 3ee232270438361bc0d53a90d25a7b401b02a71c Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Tue, 18 Aug 2026 12:19:52 -0500 Subject: [PATCH 07/15] Fix the style regressions from the previous commit Ruff and Black both went red on fb3b7ed, entirely on changes I made rather than on anything in the kernels: * I001 -- the public wrapper's three `from ... import` statements were collapsed into one while rewriting the imports for the module merge. Ruff's own fix splits them back apart, so the original shape was correct; restored. * I001 -- placement of the new `make_kernel_repr` import. * C408 -- `dict()` call in the test helper added with the R_CHUNK test; now a literal. * Black wanted one fewer blank line after the same import. Also brings the module docstring in line with the rest of the file and with the forward, which writes `gfx950 (MI355X)` rather than the bare product name -- the architecture is what constrains the kernel, the SKU is what the wall-clock was measured on, so both belong. Verified against the CI's own versions (ruff 0.16.0, black stable) before pushing this time. 7/7 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 6 +++++- .../triton/gluon/sparse_attention_dsv4_bwd.py | 1 - .../attention/test_sparse_attention_dsv4_bwd.py | 16 ++++++++-------- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index f30b8c8683..349127c888 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -32,7 +32,7 @@ Its ``lse`` is sink-inclusive (the sink is folded into ``e_max``/``e_sum`` before ``lse = e_max + log(e_sum)``), which is the convention this backward expects. -Measured on MI355X at ``T=4096 H=128 topk=512`` with a realistic SWA(128)+pool top-k: +Measured on gfx950 (MI355X) at ``T=4096 H=128 topk=512`` with a realistic SWA(128)+pool top-k: delta 0.178 / dQ 1.391 / interm 1.152 / CSR build 0.130 / gather 0.503 / d_sink 0.026 ms, 3.380 ms total = 407 TFLOPS. """ @@ -43,7 +43,11 @@ build_inverted_topk, delta_v4, dkv_gather_acc, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dq as _dq_gluon, ) from aiter.ops.triton.utils._triton import arch_info diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py index 0ecd37555c..3334587fd8 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py @@ -43,7 +43,6 @@ from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr - _dq_v4_kernel_repr = make_kernel_repr( "_dq_v4_kernel", [ diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py index 93cd133028..539f65bcab 100644 --- a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -124,14 +124,14 @@ def test_sparse_mla_bwd_dsv4(T, H, topk, npool, has_sink, r_chunk): def _dummy_inputs(T=64, H=64, topk=64, dev="cuda"): - return dict( - q=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), - kv=torch.randn(T, D, device=dev, dtype=torch.bfloat16), - do=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), - o=torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), - lse=torch.randn(T, H, device=dev, dtype=torch.float32), - idx=torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev), - ) + return { + "q": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "kv": torch.randn(T, D, device=dev, dtype=torch.bfloat16), + "do": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "o": torch.randn(T, H, D, device=dev, dtype=torch.bfloat16), + "lse": torch.randn(T, H, device=dev, dtype=torch.float32), + "idx": torch.randint(0, T, (T, topk), dtype=torch.int32, device=dev), + } def test_sparse_mla_bwd_dsv4_rejects_bad_chunk(): From 17abe84b5ff4b1409043c74ec7e66f930f74c494 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Thu, 20 Aug 2026 21:43:14 -0500 Subject: [PATCH 08/15] Address review: add an op benchmark, drop dead knobs, trim the docstrings **Benchmark.** Review asked for a way to reproduce the performance in the PR description; the unit tests only establish correctness. `op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py` is modelled on the forward's bench_sparse_attention_dsv4.py. The top-k builder is the part worth attention. It is a sliding window plus a causally-visible compressed pool, not uniform random, because the two are not interchangeable for this op: a uniform top-k spreads contributors evenly over the KV rows, while the real distribution gives the pool rows runs of a few thousand, and the dKV gather is sized for exactly those runs. Benchmarking on uniform indices reports a gather roughly twice as fast as the real one. `--breakdown` times each phase separately, which is what produces the per-kernel table. A small-shape autograd check runs first so a broken kernel fails before it is timed. There is no FlyDSL comparison column: that backward has no single entry point -- it is driven phase by phase through compile() calls -- so wiring it in would mean carrying third-party launch code that aiter reviewers cannot run anyway. The comparison in the PR description was measured out of tree. **PREFETCH.** Never anything but 0: the launcher defaulted it off and the public wrapper never passed it, in any commit on this branch. It is a gl.constexpr, so the branch was folded at compile time and the prefetch path never reached the ISA -- removing it changes no generated code, and interm measures 1.136 ms against 1.138 before. The review suggested it would save a branch in the loop; that part does not hold, the value here is that the file no longer offers a knob that does nothing. What it would have done is software-pipeline the dS/P loads one rank tile ahead; that was slower at BD=256, which is the configuration we ship, because the extra live blocks cost occupancy. **ACCUMULATE.** Likewise always True from the caller. Rather than delete it, the wrapper now passes `accumulate=(r > 0)`: the accumulator is freshly zeroed, so the first chunk can write instead of reading it back. This is not a speedup -- the accumulator is 10.5 MB, about 1.7 us against a 3.2 ms pipeline, and end-to-end does not move outside noise. It makes the flag reachable and its docstring true. **Docstrings.** They lose the tuning history. What we measured on the way to a configuration is not something a reader of the shipped code needs; only the configuration itself is. Same rule applied to the CSR build, where a rejected bincount+cumsum formulation was being documented alongside the one actually used. 7/7 tests pass on gfx950 (MI355X). Benchmark reports 3.231 ms end-to-end and 3.217 ms as the per-kernel sum, 425 / 427 TFLOPS. --- .../attention/sparse_attention_dsv4_bwd.py | 5 +- .../triton/gluon/sparse_attention_dsv4_bwd.py | 124 ++---- .../triton/bench_sparse_attention_dsv4_bwd.py | 358 ++++++++++++++++++ 3 files changed, 391 insertions(+), 96 deletions(-) create mode 100644 op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 349127c888..8c86ce966e 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -191,7 +191,10 @@ def sparse_mla_bwd_dsv4( inv_ptr, inv_data = build_inverted_topk( topk_indices[:, r : r + R_CHUNK], num_kv ) - dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc) + # dkv_acc was just zeroed, so the first chunk can write instead of read-modify-write. + # Unchunked -- the default -- that is the only chunk, and it saves reading back the + # whole [num_kv, 512] fp32 accumulator. + dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc, accumulate=(r > 0)) dkv = dkv_acc.to(kv.dtype) diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py index 3334587fd8..520fd32b1d 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py @@ -427,7 +427,6 @@ def sparse_mla_bwd_dq( "D", "MFMA_K", "DUAL_STAGE", - "PREFETCH", ], ) @@ -455,7 +454,6 @@ def _dkv_interm_v4_kernel( D: gl.constexpr, MFMA_K: gl.constexpr, DUAL_STAGE: gl.constexpr, - PREFETCH: gl.constexpr, ): """Grid (T, D//BD). NH is the padded head count and the mfma contraction dim.""" # instr_shape[2]=32: on gfx950 v_mfma_f32_16x16x32_bf16 does 2x the FLOPs of the 16-deep @@ -568,65 +566,19 @@ def _dkv_interm_v4_kernel( offs_d_st = d_off + gl.arange(0, BD, layout=gl.SliceLayout(1, mfma)) offs_col_st = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma)) - # dS/P prefetch one rank tile ahead: without it the HBM load latency is fully exposed every - # iteration, which is what caps the restructured kernel once it stops being BW-bound. The - # last iteration re-loads tile NUM_TILES-1 rather than branching -- one redundant 8 KB load - # in NUM_TILES, cheaper than peeling the loop. - offs0 = ( - ds_base - + offs_h_ds[:, None].to(tl.int64) * stride_ds_h - + offs_k_ds[None, :].to(tl.int64) - ) - dS_nxt = gl.amd.cdna4.buffer_load( - ptr=dS_ptr, offsets=offs0.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 - ) - P_nxt = gl.amd.cdna4.buffer_load( - ptr=P_ptr, offsets=offs0.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 - ) - for t in range(NUM_TILES): - if PREFETCH: - dS_blk = dS_nxt - P_blk = P_nxt - # re-load the last tile rather than branch: one redundant 8 KB load in NUM_TILES - t_nxt = min(t + 1, NUM_TILES - 1) - col_n = t_nxt * TILE_K + offs_k_ds - offs_n = ( - ds_base - + offs_h_ds[:, None].to(tl.int64) * stride_ds_h - + col_n[None, :].to(tl.int64) - ) - dS_nxt = gl.amd.cdna4.buffer_load( - ptr=dS_ptr, - offsets=offs_n.to(tl.int32), - mask=mask_h_ds[:, None], - other=0.0, - ) - P_nxt = gl.amd.cdna4.buffer_load( - ptr=P_ptr, - offsets=offs_n.to(tl.int32), - mask=mask_h_ds[:, None], - other=0.0, - ) - else: - col_c = t * TILE_K + offs_k_ds - offs_c = ( - ds_base - + offs_h_ds[:, None].to(tl.int64) * stride_ds_h - + col_c[None, :].to(tl.int64) - ) - dS_blk = gl.amd.cdna4.buffer_load( - ptr=dS_ptr, - offsets=offs_c.to(tl.int32), - mask=mask_h_ds[:, None], - other=0.0, - ) - P_blk = gl.amd.cdna4.buffer_load( - ptr=P_ptr, - offsets=offs_c.to(tl.int32), - mask=mask_h_ds[:, None], - other=0.0, - ) + col_c = t * TILE_K + offs_k_ds + offs_c = ( + ds_base + + offs_h_ds[:, None].to(tl.int64) * stride_ds_h + + col_c[None, :].to(tl.int64) + ) + dS_blk = gl.amd.cdna4.buffer_load( + ptr=dS_ptr, offsets=offs_c.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) + P_blk = gl.amd.cdna4.buffer_load( + ptr=P_ptr, offsets=offs_c.to(tl.int32), mask=mask_h_ds[:, None], other=0.0 + ) dS_dot = gl.convert_layout(dS_blk, dot_b) P_dot = gl.convert_layout(P_blk, dot_b) @@ -658,20 +610,14 @@ def sparse_mla_bwd_dkv_interm_v4( TILE_K=128, MFMA_K=32, DUAL_STAGE=1, - PREFETCH=0, H_POW2=None, num_warps=4, interm=None, ): """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. - Defaults measured at T=4096 H=128 topk=512 on gfx950 (MI355X): 1.170 ms. Sweep notes: - * BD=256 beats 128 (traffic: dS/P costs D/BD x). - * MFMA_K=32 is worth ~14% at the best config. It only pays once the kernel is off the - bandwidth ceiling, which is what splitting D across grid.y buys. - * PREFETCH=0: prefetching dS/P helps at BD=128 but is consistently WORSE at BD=256 - (1.181 -> 1.306), where the extra live registers cost occupancy. - * DUAL_STAGE=1 (both prologue copies behind one drain) is a small consistent win. + ``BD`` splits D across ``grid.y``; dS/P are re-read once per D block, so a larger BD moves + less of them. ``MFMA_K=32`` is the CDNA4 16x16x32 depth. """ T, H, D = q.shape assert R_CHUNK % TILE_K == 0 @@ -701,7 +647,6 @@ def sparse_mla_bwd_dkv_interm_v4( D=D, MFMA_K=MFMA_K, DUAL_STAGE=DUAL_STAGE, - PREFETCH=PREFETCH, num_warps=num_warps, ) return interm @@ -736,9 +681,8 @@ def _delta_v4_kernel( def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): - # BLOCK_R=8 / num_warps=8 measured best (0.173 ms, 6.21 TB/s = 78% peak at T4096 H128); - # the whole sweep plateaus at 0.173-0.187 once a lane loads >= 8 bf16, i.e. once the load - # is a dwordx4. Below that (BLOCK_R=2 nw=8, 2 bf16/lane) it falls off a cliff to 3.12 TB/s. + # BLOCK_R=8 keeps each lane loading >= 8 bf16, i.e. a dwordx4; narrower blocks drop to a + # dword and the kernel loses most of its bandwidth. """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. @@ -784,18 +728,16 @@ def _bwd_dkv_gather_acc_v4( ): """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. - Carrying ``BLOCK_E`` entries per iteration rather than walking the run one entry at a time - buys two separate things, and this gather needs both: + ``BLOCK_E`` entries are carried per iteration, which the gather needs for two reasons: * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per - lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane - instead, so the loads become dwordx4. The gather is issue-bound, so this dominates. - * **loop trip count.** The run is consumed BLOCK_E entries at a time and ``tl.sum`` over - the entry axis folds them. A realistic top-k gives run lengths up to ~3000 (pool rows), - so the serial walk was the other half of the problem. - - ``ACCUMULATE=False`` skips the read-modify-write of the destination, valid when the caller - does not chunk (each KV row is then written by exactly one CTA). + lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane, + so the loads become dwordx4. The gather is issue-bound, so this dominates. + * **trip count.** ``tl.sum`` folds the entry axis, so the run is consumed BLOCK_E at a + time. A realistic top-k gives run lengths up to ~3000 on the pool rows. + + ``ACCUMULATE=False`` writes the destination instead of reading it back first. The caller + uses it for the first chunk, where the accumulator is still zero. """ k = tl.program_id(0) offs_d = tl.arange(0, D) @@ -827,15 +769,12 @@ def build_inverted_topk(topk_indices_slice, num_kv): """CSR inverted index over ``num_kv`` KV rows. One stable sort yields both the permutation (``inv_data``) and the sorted keys; - ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k, - which is exactly what ``cumsum(bincount(flat+1))`` computes. Invalid (-1) entries sort to - the front, so ``inv_ptr[0]`` starts past them and they are never visited. + ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k. + Invalid (-1) entries sort to the front, so ``inv_ptr[0]`` starts past them and they are + never visited. - Two details carry most of the cost, and the obvious formulation gets both wrong -- writing - this as ``bincount`` + ``cumsum`` over int64 keys measured 3x slower for the same output: - * the sort key is narrowed to int16 when ``num_kv`` fits, so the radix sort makes 2 - byte-passes instead of 8; - * ``searchsorted`` does the job of the separate ``bincount`` + ``cumsum`` passes. + The sort key is narrowed to int16 when ``num_kv`` fits, which is what keeps the radix sort + to two byte-passes. Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. """ @@ -857,11 +796,6 @@ def build_inverted_topk(topk_indices_slice, num_kv): def dkv_gather_acc( interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True ): - # BLOCK_E=64 / num_warps=8 measured best (0.345 ms, 6.29 TB/s = 79% peak at T4096 H128 - # topk512 SWA+pool). Time falls monotonically with BLOCK_E across the whole sweep - # (4->64: 1.021, 0.715, 0.505, 0.407, 0.345), i.e. both the load width AND the trip count - # on the ~3000-entry pool runs were binding. The one-entry-at-a-time walk this replaced - # was 1.491 ms at 1.45 TB/s. """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..98e978a4b2 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Benchmark for the DSv4 sparse MLA training backward (gfx950 / CDNA4). + +Reproduces the numbers quoted for `sparse_mla_bwd_dsv4`. The top-k is not uniform random: +it is a sliding window plus a causally-visible compressed pool, the distribution the V4 +indexer actually produces. That matters here -- a uniform top-k gives every KV row about the +same number of contributors, while the real one gives the pool rows runs of a few thousand, +and the dKV gather is sized for those runs. Benchmarking on uniform indices flatters the +gather by roughly 2x. + +`--breakdown` additionally times each phase on its own, which is how the per-kernel table in +the PR description was produced. + +Usage: + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py --breakdown + python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py --cfgs 4096,128,512 +""" + +import argparse + +import torch +import triton + +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( + build_inverted_topk, + delta_v4, + dkv_gather_acc, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( + sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, +) +from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( + sparse_mla_bwd_dq as _dq_gluon, +) +from aiter.ops.triton.utils._triton import arch_info + +D = 512 +SWA = 128 # sliding-window width +CR_POOL = ( + 4 # pool compression ratio: block b is visible to token t once (b+1)*CR_POOL-1 <= t +) + + +# --------------------------------------------------------------------------- +# Bench data builder +# --------------------------------------------------------------------------- +def _build_topk_swa_pool(T, topk, num_pool, device, generator): + """[T, topk] int32 -- SWA(128) window plus causally-visible pool ranks, -1 padded. + + KV row layout is ``[ per-token 0..T-1 | pool T..T+num_pool-1 ]``. Early tokens see fewer + than `topk - SWA` pool blocks, so their trailing slots stay -1, exactly as in production. + """ + idx = torch.arange(T, device=device) + off = torch.arange(SWA, device=device) + swa = idx[:, None] - (SWA - 1) + off[None, :] + swa = torch.where(swa >= 0, swa, torch.full_like(swa, -1)) + + n_pool = topk - SWA + if n_pool <= 0 or num_pool == 0: + return swa.to(torch.int32).contiguous() + assert n_pool <= num_pool, f"need {n_pool} pool ranks but only {num_pool} blocks" + + n_visible = torch.clamp((idx + 1) // CR_POOL, min=0, max=num_pool) + blk = torch.arange(num_pool, device=device) + visible = blk[None, :] < n_visible[:, None] + + # Pick n_pool visible blocks at random: score the visible ones in [0,1), push the rest to + # 2.0, take the smallest. + score = torch.rand(T, num_pool, device=device, generator=generator) + score = torch.where(visible, score, torch.full_like(score, 2.0)) + sel = score.argsort(dim=1)[:, :n_pool] + pool = torch.where(torch.gather(visible, 1, sel), T + sel, torch.full_like(sel, -1)) + return torch.cat([swa, pool], dim=1).to(torch.int32).contiguous() + + +def _build_case(T, H, topk, device, num_pool=1024, seed=0): + """Inputs for one configuration. + + `o` / `lse` are synthetic. No kernel branches on their values, so timing is unaffected and + this avoids a reference forward that would dominate the run. + """ + gen = torch.Generator(device=device) + gen.manual_seed(seed) + indices = _build_topk_swa_pool(T, topk, num_pool, device, gen) + num_kv = T + num_pool if topk > SWA else T + + def _bf16(*shape): + return torch.randn(*shape, device=device, dtype=torch.bfloat16, generator=gen) + + return { + "q": _bf16(T, H, D), + "kv": _bf16(num_kv, D), + "do": _bf16(T, H, D), + "o": _bf16(T, H, D), + "lse": torch.randn(T, H, device=device, dtype=torch.float32, generator=gen), + "sink": ( + torch.randn(H, device=device, dtype=torch.float32, generator=gen) * 0.1 + ).contiguous(), + "indices": indices, + "num_kv": num_kv, + } + + +# --------------------------------------------------------------------------- +# Timing +# --------------------------------------------------------------------------- +def _bench(fn, *, warmup=5, reps=20): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + ev0 = torch.cuda.Event(enable_timing=True) + ev1 = torch.cuda.Event(enable_timing=True) + ev0.record() + for _ in range(reps): + fn() + ev1.record() + torch.cuda.synchronize() + return ev0.elapsed_time(ev1) / reps + + +def _flops(T, H, topk): + # dQ contributes S, dP and dS@kv; dKV-interm contributes dS*Q and P*dO. Five 2*D-flop + # products per (token, head, top-k slot). + return 10.0 * D * T * H * topk + + +def _time_phases(case, scale): + """Per-phase timings, in the order the public entry runs them.""" + q, kv, do, o = case["q"], case["kv"], case["do"], case["o"] + lse, sink, indices = case["lse"], case["sink"], case["indices"] + T, H, _ = q.shape + topk = indices.shape[1] + num_kv = case["num_kv"] + + delta = delta_v4(o, do) + dq = torch.empty_like(q) + dkv_acc = torch.zeros(num_kv, D, dtype=torch.float32, device=q.device) + chunk_dS = torch.empty(T, H, topk, dtype=torch.bfloat16, device=q.device) + chunk_P = torch.empty(T, H, topk, dtype=torch.bfloat16, device=q.device) + interm = torch.empty(T, topk, D, dtype=torch.bfloat16, device=q.device) + inv_ptr, inv_data = build_inverted_topk(indices, num_kv) + + phases = [ + ("delta", lambda: delta_v4(o, do)), + ( + "dq", + lambda: _dq_gluon( + q, + kv, + do, + indices, + lse, + delta, + dq, + chunk_dS, + chunk_P, + scale, + 0, + topk, + BLOCK_H=64, + TILE_K=32, + is_first_chunk=True, + ), + ), + ( + "interm", + lambda: _dkv_interm_gluon( + q, do, chunk_dS, chunk_P, topk, BD=256, TILE_K=128, interm=interm + ), + ), + ("csr_build", lambda: build_inverted_topk(indices, num_kv)), + ("gather", lambda: dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc)), + ( + "d_sink", + lambda: -(torch.exp(sink[None, :].float() - lse) * delta).sum(dim=0), + ), + ] + return [(name, _bench(fn)) for name, fn in phases] + + +# --------------------------------------------------------------------------- +# Correctness gate +# --------------------------------------------------------------------------- +def _ref_scores(q_t, kv_f, idx, scale): + """Masked fp32 scores for one token: [H, topk], invalid slots -inf. Also returns the rows.""" + valid = idx != -1 + k = kv_f[idx.clamp(min=0)] + k = torch.where(valid[:, None], k, torch.zeros_like(k)) + s = (q_t @ k.t()) * scale + return torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))), k, valid + + +def _ref_forward(q, kv, sink, topk, scale): + """fp32 reference forward. Per-token python loop, so small shapes only. + + Returns `o` and the leaf tensors, so the caller can drive autograd through it. + """ + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) + outs = [] + for t in range(q.shape[0]): + s, k, valid = _ref_scores(qg[t], kvg, topk[t].long(), scale) + m = torch.maximum(s.max(dim=1).values, sg) + p = torch.where(valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s)) + denom = p.sum(dim=1) + torch.exp(sg - m) + outs.append((p @ k) / denom[:, None]) + return torch.stack(outs, dim=0), (qg, kvg, sg) + + +def _ref_lse(q, kv, sink, topk, scale): + """Sink-inclusive log-sum-exp matching the forward's convention.""" + with torch.no_grad(): + kv_f = kv.float() + rows = [] + for t in range(q.shape[0]): + s, _, _ = _ref_scores(q[t].float(), kv_f, topk[t].long(), scale) + rows.append(torch.logsumexp(torch.cat([s, sink[:, None]], dim=1), dim=1)) + return torch.stack(rows) + + +def check_correctness(device): + """Small-shape gate against autograd, so a broken kernel fails before it is timed.""" + print("\n========== CORRECTNESS ==========") + T, H, topk = 128, 64, 128 + scale = 1.0 / (D**0.5) + case = _build_case(T, H, topk, device, num_pool=64) + q, kv, do, sink, indices = ( + case["q"], + case["kv"], + case["do"], + case["sink"], + case["indices"], + ) + + ref_o, (qg, kvg, sg) = _ref_forward(q, kv, sink, indices, scale) + o = ref_o.detach().to(torch.bfloat16).contiguous() + lse = _ref_lse(q, kv, sink, indices, scale) + + dq, dkv, d_sink = sparse_mla_bwd_dsv4( + q, kv, do, o, lse, indices, attn_sink=sink, scale=scale + ) + ref_o.backward(do.float()) + + def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.float().reshape(-1), b.float().reshape(-1), dim=0 + ).item() + + for name, got, want in ( + ("dq", dq, qg.grad), + ("dkv", dkv, kvg.grad), + ("d_sink", d_sink, sg.grad), + ): + cos = _cos(got, want) + assert cos > 0.999, f"{name} cos={cos:.6f}" + print(f" {name:7s}: OK (cos={cos:.6f})") + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- +def _print_table(title, headers, rows): + def _fmt(x): + if isinstance(x, float): + return f"{x:.3f}" if x >= 1 or x == 0 else f"{x:.4f}" + return str(x) + + cells = [[_fmt(c) for c in r] for r in rows] + widths = [max(len(h), *(len(c[i]) for c in cells)) for i, h in enumerate(headers)] + print("| " + " | ".join(h.rjust(widths[i]) for i, h in enumerate(headers)) + " |") + print("| " + " | ".join("-" * widths[i] for i in range(len(headers))) + " |") + for c in cells: + print("| " + " | ".join(s.rjust(widths[i]) for i, s in enumerate(c)) + " |") + + +def run_bwd_bench(args, device): + print("\n========== BACKWARD ==========") + rows = [] + for T, H, topk in args.cfgs: + case = _build_case(T, H, topk, device) + scale = 1.0 / (D**0.5) + tflops = _flops(T, H, topk) / 1e12 + + ms = _bench( + lambda c=case, s=scale: sparse_mla_bwd_dsv4( + c["q"], + c["kv"], + c["do"], + c["o"], + c["lse"], + c["indices"], + attn_sink=c["sink"], + scale=s, + ) + ) + rows.append([T, H, case["num_kv"], topk, ms, tflops / (ms * 1e-3)]) + + if args.breakdown: + print(f"\n per-kernel, T={T} H={H} topk={topk}:") + phases = _time_phases(case, scale) + total = sum(t for _, t in phases) + for name, t in phases: + print(f" {name:10s} {t:7.3f} ms ({100 * t / total:4.1f}%)") + print( + f" {'SUM':10s} {total:7.3f} ms -> {tflops / (total * 1e-3):.0f} TFLOPS" + ) + + _print_table("BACKWARD", ["T", "H", "Kv", "topk", "ms", "TFLOPS"], rows) + + +def _parse_args(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--cfgs", + nargs="+", + type=str, + default=[ # (T, H, topk) + "4096,128,512", + "4096,128,1024", + "8192,128,512", + ], + ) + p.add_argument( + "--breakdown", + action="store_true", + help="also time each phase separately (the per-kernel table in the PR description)", + ) + p.add_argument("--skip-correctness", action="store_true") + args = p.parse_args() + args.cfgs = [tuple(int(x) for x in s.split(",")) for s in args.cfgs] + return args + + +def main(): + args = _parse_args() + device = "cuda" + if arch_info.get_arch() != "gfx950": + print(f"sparse_mla_bwd_dsv4 is gfx950 only; this is {arch_info.get_arch()}") + return + print( + f"GPU: {torch.cuda.get_device_name(0)} " + f"({torch.cuda.get_device_properties(0).multi_processor_count} CUs)" + ) + print(f"Triton: {triton.__version__}") + + if not args.skip_correctness: + check_correctness(device) + run_bwd_bench(args, device) + + +if __name__ == "__main__": + main() From 587618040c6bb7c37960cf6a37e3712d1eb2b2f6 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Fri, 21 Aug 2026 09:38:03 -0500 Subject: [PATCH 09/15] Check that topk_indices has one row per query token From review. Only `topk_indices.shape[1]` was read, so nothing rejected an index tensor with a different number of rows than `q`. The dQ grid is sized from `q`, and the kernel addresses each row as `token_idx * stride_topk_t`, so a shorter tensor is read past its end -- the same failure mode as the R_CHUNK case: plausible-looking indices for most tokens, then off the end. The other tensors that feed the same grid were already covered (`do` and `o` against `q.shape`, `lse` against `(T, H)`, `attn_sink` against `(H,)`); this was the remaining gap. 8/8 tests pass on gfx950 (MI355X). --- .../ops/triton/attention/sparse_attention_dsv4_bwd.py | 4 ++++ .../attention/test_sparse_attention_dsv4_bwd.py | 11 +++++++++++ 2 files changed, 15 insertions(+) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 8c86ce966e..d9b8fb1789 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -103,6 +103,10 @@ def sparse_mla_bwd_dsv4( raise RuntimeError(f"{name} dtype mismatch: {name}={t.dtype}, q={q.dtype}") T, H, D = q.shape + assert topk_indices.ndim == 2 and topk_indices.shape[0] == T, ( + f"topk_indices must be [T, TOPK] with T={T}, got {tuple(topk_indices.shape)} -- the dQ " + "grid is sized from q, so a shorter index tensor is read past its end" + ) TOPK = topk_indices.shape[1] num_kv = kv.shape[0] assert D == 512, f"DSv4 sparse-MLA backward is fixed to head_dim 512, got {D}" diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py index 539f65bcab..b9a5abc12b 100644 --- a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -143,6 +143,17 @@ def test_sparse_mla_bwd_dsv4_rejects_bad_chunk(): ) +def test_sparse_mla_bwd_dsv4_rejects_short_topk(): + """topk_indices must have one row per query token. + + The dQ grid is sized from `q`, so a shorter index tensor is read past its end rather than + producing a shape error anywhere downstream. + """ + t = _dummy_inputs(T=64, topk=64) + with pytest.raises(AssertionError, match="topk_indices must be"): + sparse_mla_bwd_dsv4(t["q"], t["kv"], t["do"], t["o"], t["lse"], t["idx"][:32]) + + def test_sparse_mla_bwd_dsv4_rejects_indivisible_chunk(): """A chunk width that is a valid tile multiple but does not divide TOPK is still rejected. From 7f82c52642b88b2b1d2f7f592afb5ee6612b8cae Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Fri, 21 Aug 2026 10:10:50 -0500 Subject: [PATCH 10/15] Move the Gluon kernels to _gluon_kernels/gfx950/attention Requested in review: `gluon/` is the older location and new Gluon kernels belong under `_gluon_kernels/`. The arch-namespaced path also fits this kernel better than the flat one did, since it is gfx950-only and asserts as much -- it now sits beside `_gluon_kernels/gfx950/attention/pa_decode_sparse.py`, which has the same shape: private kernels here, public entry under `attention/`. Imports in the public wrapper and the benchmark follow. No `__init__.py` is added; that tree uses namespace packages, as pa_decode_sparse already does. The `gluon/README.md` entry goes away with it. That file scopes itself to "all kernels in this directory", and the ops already living in `_gluon_kernels/` (pa_decode_sparse, fp8_mqa_logits) are not listed there either, so keeping the row would have been the odd one out. There is no README under `_gluon_kernels/` to move it to, and none is wanted; what it documented now lives in the PR description, which gained the full signature and the shapes the wrapper enforces: sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None) -> (dq, dkv, d_sink) q, do, o [T, H, 512] bf16, contiguous kv [num_kv, 512] bf16, K == V, num_kv >= T (rows T.. are the compressed pool) lse [T, H] fp32, sink-inclusive, from the forward topk_indices [T, TOPK] int32, -1 marks an invalid slot attn_sink [H] fp32, optional R_CHUNK multiple of 32 and a divisor of TOPK; None = unchunked head_dim is fixed at 512 and the op is gfx950-only; the public wrapper asserts both. Kernels are byte-identical apart from one docstring line that referred to the old location. 8/8 tests pass and the benchmark reports 3.226 ms / 426 TFLOPS on gfx950 (MI355X), unchanged. --- .../attention}/sparse_attention_dsv4_bwd.py | 4 +-- .../attention/sparse_attention_dsv4_bwd.py | 6 ++-- aiter/ops/triton/gluon/README.md | 30 ------------------- .../triton/bench_sparse_attention_dsv4_bwd.py | 8 ++--- 4 files changed, 9 insertions(+), 39 deletions(-) rename aiter/ops/triton/{gluon => _gluon_kernels/gfx950/attention}/sparse_attention_dsv4_bwd.py (99%) diff --git a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py similarity index 99% rename from aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py rename to aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py index 520fd32b1d..1bc1f20e5a 100644 --- a/aiter/ops/triton/gluon/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py @@ -6,8 +6,8 @@ All operate on the official V4 form (``K == V == kv``, one dense 512-wide tensor, RoPE already applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` in the softmax denominator only, ``topk == -1`` masked). The two MFMA phases are Gluon; the two memory-bound phases are -plain Triton and live here rather than under ``_triton_kernels/`` because there is no Triton -implementation of this backward to fall back to -- they are parts of this kernel, not an +plain Triton and live alongside them rather than under ``_triton_kernels/`` because there is no +Triton implementation of this backward to fall back to -- they are parts of this kernel, not an alternative to it. ``_dq_v4_kernel`` diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index d9b8fb1789..f2e77125d8 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -39,15 +39,15 @@ import torch -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( build_inverted_topk, delta_v4, dkv_gather_acc, ) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, ) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dq as _dq_gluon, ) from aiter.ops.triton.utils._triton import arch_info diff --git a/aiter/ops/triton/gluon/README.md b/aiter/ops/triton/gluon/README.md index 7c0096bc95..d6477dc63a 100644 --- a/aiter/ops/triton/gluon/README.md +++ b/aiter/ops/triton/gluon/README.md @@ -57,12 +57,6 @@ Some features (e.g., scheduling hints like `sched_barrier`) require the [AMD Glu python op_tests/triton_tests/
test_pa_decode_gluon.py TBDTBDTBD - - sparse_attention_
dsv4_bwdDSv4 Sparse
MLA BackwardCDNA4 - Q/KV/dO/O: bf16, K == V
head_dim = 512 (dense)
lse: fp32, sink-inclusive
num_kv ≥ T (pool ok)
topk % 32 == 0
gfx950 only - python op_tests/triton_tests/
attention/test_sparse_
attention_dsv4_bwd.py - ~407
TFLOPS—— - @@ -253,30 +247,6 @@ python op_tests/test_mla.py -c 10000 100000 -b 1 3 4 -n 16,1 -d bf16 -kvd bf16 - | 100K | 3 | 85 | 19 | 88.77 | 3.89 | | 100K | 4 | 64 | 25 | 106.96 | 4.31 | -### `sparse_attention_dsv4_bwd.py` — DeepSeek V4 Sparse MLA Backward - -**Public entry:** `aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None)` -> `(dq, dkv, d_sink)` - -Training backward for the DSv4 sparse prefill attention (the `has_pe=False` form of `mla_gluon`). Same op contract: `K == V == kv` as one dense 512-wide tensor, RoPE applied in place caller-side, scale `1/sqrt(512)`, `attn_sink` in the softmax denominator only, `topk_indices == -1` masked. - -`o` and `lse` come from the forward; `mla_gluon(..., has_pe=False, return_lse=True)` produces both, and its `lse` is already sink-inclusive, which is what this backward expects. - -Five kernels plus one torch reduction: - -| phase | impl | notes | -|---|---|---| -| `delta = rowsum(O*dO)` | triton | streams bf16, accumulates fp32 | -| dQ | **gluon** | one LDS read of the gathered KV feeds both the `S` and `dP` MFMAs; also emits this chunk's `dS`/`P` | -| dKV-interm | **gluon** | contracts over all heads inside one MFMA pair, Q/dO transposed once into registers, D split across `grid.y` | -| CSR build | torch | `sort` + `searchsorted` on int16-narrowed keys | -| dKV gather | triton | atomic-free: the scatter is inverted so each KV row gathers its own contributors | -| `d_sink` | torch | 26 us | - -`R_CHUNK` splits the rank dimension. Leave it `None` (unchunked) unless memory forces otherwise: it exists only to bound the `interm` intermediate (`T*topk*512` bf16, 2.0 GiB at T=4096 topk=512) and costs a dQ read-modify-write between chunks plus one CSR build per chunk. - -**Measured** (MI355X, T=4096 H=128 topk=512, SWA(128)+pool top-k): 3.38 ms / 407 TFLOPS as a per-kernel sum, 3.24 ms / 424 TFLOPS end-to-end. - - ### `pa_decode_gluon.py` — Paged Attention Decode **Function:** `pa_decode_gluon(output, query, key_cache, value_cache, context_lengths, block_tables, softmax_scale, query_length, max_context_partition_num, context_partition_size, compute_type, query_scale, key_scale, value_scale, ...)` diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py index 98e978a4b2..002c507a77 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -24,18 +24,18 @@ import torch import triton -from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( build_inverted_topk, delta_v4, dkv_gather_acc, ) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, ) -from aiter.ops.triton.gluon.sparse_attention_dsv4_bwd import ( +from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( sparse_mla_bwd_dq as _dq_gluon, ) +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 from aiter.ops.triton.utils._triton import arch_info D = 512 From 7e004f1125b632f8a7c1d13048eda95a2e8ae07a Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Fri, 21 Aug 2026 14:16:29 -0500 Subject: [PATCH 11/15] Separate the kernel layers, and stop duplicating the pipeline and the reference All from Bruno's review. **Layering.** Kernel-layer modules in this repo import no torch -- both neighbours in `_gluon_kernels/gfx950/attention/` have zero torch references, and so does everything under `_triton_kernels/attention/`. Ours was the exception. The two `@gluon.jit` kernels stay in `_gluon_kernels/gfx950/attention/`, the two `@triton.jit` kernels move to `_triton_kernels/attention/`, and the five launchers plus the CSR build and the sink gradient move up to `attention/sparse_attention_dsv4_bwd.py`. Both kernel modules are now torch-free with no public functions, matching their neighbours, so the kernels can be reached without pulling in torch. The Triton path is `_triton_kernels/attention/`, not `_triton_kernels/gfx950/attention/` as the comment suggested: that tree has no arch level. `pa_decode_sparse`, `fp8_mqa_logits` and `mla` all sit arch-namespaced under `_gluon_kernels/` and flat under `_triton_kernels/`. **Pipeline duplication.** The benchmark was restating the production orchestration -- phase order, tile widths, workspace, and a copy of the sink-gradient expression. It now shares the real thing: `plan_bwd()` validates and allocates, `bwd_phases()` yields `(name, thunk)` in execution order, `sparse_mla_bwd_dsv4` runs them, and the benchmark times them one at a time. The measurement is of the phases the op actually runs, and there is nothing left to drift. The sink gradient lives once, on the plan object. **Reference duplication.** The fp32 reference forward existed in both the test and the benchmark. It moves to `aiter.test_mha_common.sparse_mla_dsv4_ref`, beside the other reference implementations, and returns `(o, lse)` from one differentiable pass so backward values come from autograd through it. Both callers use it; the test file loses 46 lines. **Benchmark harness.** `triton.testing.do_bench` replaces the hand-rolled event loop -- 55 of the 67 benchmarks here already use it. Help text added to `--cfgs` and `--skip-correctness`. do_bench flushes L2 between repetitions, so the reported numbers move: 3.271 ms / 420 TFLOPS end-to-end at T=4096 H=128 topk=512, against 3.229 / 426 with the previous cache-warm loop. The per-phase figures move further for the same reason and now sum slightly above the end-to-end number. The PR description carries the new tables. 8/8 tests pass on gfx950 (MI355X). --- .../attention/sparse_attention_dsv4_bwd.py | 295 +----------- .../attention/sparse_attention_dsv4_bwd.py | 113 +++++ .../attention/sparse_attention_dsv4_bwd.py | 424 +++++++++++++++--- aiter/test_mha_common.py | 42 ++ .../triton/bench_sparse_attention_dsv4_bwd.py | 154 ++----- .../test_sparse_attention_dsv4_bwd.py | 70 +-- 6 files changed, 561 insertions(+), 537 deletions(-) create mode 100644 aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py diff --git a/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py index 1bc1f20e5a..7284a98cfc 100644 --- a/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py @@ -1,14 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). +"""Gluon kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). -All operate on the official V4 form (``K == V == kv``, one dense 512-wide tensor, RoPE already -applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` in the softmax denominator -only, ``topk == -1`` masked). The two MFMA phases are Gluon; the two memory-bound phases are -plain Triton and live alongside them rather than under ``_triton_kernels/`` because there is no -Triton implementation of this backward to fall back to -- they are parts of this kernel, not an -alternative to it. +The two MFMA phases. Both operate on the official V4 form: ``K == V == kv`` as one dense +512-wide tensor, RoPE already applied in place caller-side, scale ``1/sqrt(512)``, ``attn_sink`` +in the softmax denominator only, ``topk == -1`` masked. ``_dq_v4_kernel`` Per (query token, head block): ``S = Q@kv^T``, ``P = exp(S - lse)``, ``dP = dO@kv^T``, @@ -22,21 +19,10 @@ are transposed once into registers and D is split across ``grid.y``, which is what keeps them read once instead of ``topk/TILE_K`` times. -``_delta_v4_kernel`` - ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the - bf16 inputs and accumulates in fp32, so it moves exactly the working set. - -``_bwd_dkv_gather_acc_v4`` + ``build_inverted_topk`` - Reduce ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is - inverted into a CSR gather (each output KV row collects its own contributors), so no atomics - are needed. ``BLOCK_E`` entries are carried per loop iteration, which both widens the load - and cuts the trip count on the long runs a realistic top-k produces. - -Public entry: ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd.sparse_mla_bwd_dsv4``. +Launchers live in ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd``; this module stays +free of torch so the kernels can be called without it. """ -import torch -import triton import triton.language as tl from triton.experimental import gluon from triton.experimental.gluon import language as gl @@ -364,59 +350,6 @@ def _dq_v4_kernel( ) -def sparse_mla_bwd_dq( - q, - kv, - do, - topk, - lse, - delta, - dq, - chunk_dS, - chunk_P, - scale, - r_start, - R_CHUNK, - BLOCK_H=64, - TILE_K=32, - is_first_chunk=True, -): - """Launch the dQ kernel for one rank chunk. Writes ``dq`` (RMW when not the first chunk) - plus this chunk's ``chunk_dS`` / ``chunk_P``.""" - T, H, D = q.shape - _dq_v4_kernel[(T, triton.cdiv(H, BLOCK_H))]( - q, - kv, - do, - topk, - lse, - delta, - dq, - chunk_dS, - chunk_P, - q.stride(0), - q.stride(1), - kv.stride(0), - do.stride(0), - do.stride(1), - dq.stride(0), - dq.stride(1), - topk.stride(0), - chunk_dS.stride(0), - chunk_dS.stride(1), - scale, - H, - r_start, - R_CHUNK=R_CHUNK, - BLOCK_H=BLOCK_H, - TILE_K=TILE_K, - D=D, - IS_FIRST_CHUNK=is_first_chunk, - num_warps=4, - waves_per_eu=1, - ) - - _dkv_interm_v4_kernel_repr = make_kernel_repr( "_dkv_interm_v4_kernel", [ @@ -598,219 +531,3 @@ def _dkv_interm_v4_kernel( ptr=Interm_ptr, offsets=interm_offs.to(tl.int32), ) - - -def sparse_mla_bwd_dkv_interm_v4( - q, - do, - chunk_dS, - chunk_P, - R_CHUNK, - BD=256, - TILE_K=128, - MFMA_K=32, - DUAL_STAGE=1, - H_POW2=None, - num_warps=4, - interm=None, -): - """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. - - ``BD`` splits D across ``grid.y``; dS/P are re-read once per D block, so a larger BD moves - less of them. ``MFMA_K=32`` is the CDNA4 16x16x32 depth. - """ - T, H, D = q.shape - assert R_CHUNK % TILE_K == 0 - assert D % BD == 0 - h_pow2 = H_POW2 or triton.next_power_of_2(H) - if interm is None: - interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) - _dkv_interm_v4_kernel[(T, D // BD)]( - q, - do, - chunk_dS, - chunk_P, - interm, - q.stride(0), - q.stride(1), - do.stride(0), - do.stride(1), - chunk_dS.stride(0), - chunk_dS.stride(1), - interm.stride(0), - interm.stride(1), - H, - R_CHUNK=R_CHUNK, - TILE_K=TILE_K, - NH=h_pow2, - BD=BD, - D=D, - MFMA_K=MFMA_K, - DUAL_STAGE=DUAL_STAGE, - num_warps=num_warps, - ) - return interm - - -_delta_v4_kernel_repr = make_kernel_repr( - "_delta_v4_kernel", - [ - "D", - "BLOCK_R", - ], -) - - -@triton.jit(repr=_delta_v4_kernel_repr) -def _delta_v4_kernel( - O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) - dO_ptr, # [n_rows, D] bf16 - Delta_ptr, # [n_rows] fp32 - n_rows, - D: tl.constexpr, - BLOCK_R: tl.constexpr, -): - """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" - pid = tl.program_id(0) - rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) - mask = rows < n_rows - offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] - o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) - d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) - tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) - - -def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): - # BLOCK_R=8 keeps each lane loading >= 8 bf16, i.e. a dwordx4; narrower blocks drop to a - # dword and the kernel loses most of its bandwidth. - """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. - - ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. - """ - assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() - T, H, D = o.shape - n_rows = T * H - if out is None: - out = torch.empty(T, H, dtype=torch.float32, device=o.device) - _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( - o, - do, - out, - n_rows, - D=D, - BLOCK_R=BLOCK_R, - num_warps=num_warps, - ) - return out - - -_bwd_dkv_gather_acc_v4_repr = make_kernel_repr( - "_bwd_dkv_gather_acc_v4", - [ - "D", - "BLOCK_E", - "ACCUMULATE", - ], -) - - -@triton.jit(repr=_bwd_dkv_gather_acc_v4_repr) -def _bwd_dkv_gather_acc_v4( - Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] - InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers - InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token - dKV_acc_ptr, # [num_kv, D] fp32 — accumulator - stride_interm_r: tl.int64, - stride_acc_t: tl.int64, - D: tl.constexpr, - BLOCK_E: tl.constexpr, - ACCUMULATE: tl.constexpr, -): - """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. - - ``BLOCK_E`` entries are carried per iteration, which the gather needs for two reasons: - - * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per - lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane, - so the loads become dwordx4. The gather is issue-bound, so this dominates. - * **trip count.** ``tl.sum`` folds the entry axis, so the run is consumed BLOCK_E at a - time. A realistic top-k gives run lengths up to ~3000 on the pool rows. - - ``ACCUMULATE=False`` writes the destination instead of reading it back first. The caller - uses it for the first chunk, where the accumulator is still zero. - """ - k = tl.program_id(0) - offs_d = tl.arange(0, D) - offs_e = tl.arange(0, BLOCK_E) - start = tl.load(InvPtr_ptr + k) - end = tl.load(InvPtr_ptr + k + 1) - acc_base = k.to(tl.int64) * stride_acc_t - - if ACCUMULATE: - acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) - else: - acc = tl.zeros([D], dtype=tl.float32) - - for i0 in range(start, end, BLOCK_E): - idx = i0 + offs_e - m = idx < end - entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) - vals = tl.load( - Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], - mask=m[:, None], - other=0.0, - ) - acc += tl.sum(vals.to(tl.float32), axis=0) - - tl.store(dKV_acc_ptr + acc_base + offs_d, acc) - - -def build_inverted_topk(topk_indices_slice, num_kv): - """CSR inverted index over ``num_kv`` KV rows. - - One stable sort yields both the permutation (``inv_data``) and the sorted keys; - ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k. - Invalid (-1) entries sort to the front, so ``inv_ptr[0]`` starts past them and they are - never visited. - - The sort key is narrowed to int16 when ``num_kv`` fits, which is what keeps the radix sort - to two byte-passes. - - Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. - """ - # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. - # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by - # accident. - flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid - if num_kv < 32767: # int16 range, -1 included - keys = flat_kv.to(torch.int16) - row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) - else: - keys = flat_kv.to(torch.int32) - row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) - sorted_vals, inv_data = torch.sort(keys, stable=True) - inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) - return inv_ptr, inv_data.to(torch.int32) - - -def dkv_gather_acc( - interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True -): - """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. - - Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. - """ - _, _, D = interm.shape - num_kv = dkv_acc.shape[0] - _bwd_dkv_gather_acc_v4[(num_kv,)]( - interm, - inv_ptr, - inv_data, - dkv_acc, - interm.stride(1), - dkv_acc.stride(0), - D=D, - BLOCK_E=BLOCK_E, - ACCUMULATE=accumulate, - num_warps=num_warps, - ) diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py new file mode 100644 index 0000000000..b2295c8066 --- /dev/null +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Triton kernels for the DeepSeek-V4 sparse-MLA training BACKWARD (gfx950 / CDNA4). + +The two memory-bound phases; the MFMA phases are Gluon and live in +``aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd``. + +``_delta_v4_kernel`` + ``delta = rowsum(O * dO)`` -- the standard flash-attention "o_dot_do" preamble. Streams the + bf16 inputs and accumulates in fp32, so it moves exactly the working set. + +``_bwd_dkv_gather_acc_v4`` + Reduces ``interm[t, slot, :]`` into ``dkv[kv_row, :]`` over the top-k mapping. The scatter is + inverted into a CSR gather (each output KV row collects its own contributors), so no atomics + are needed. + +Launchers live in ``aiter.ops.triton.attention.sparse_attention_dsv4_bwd``; this module stays +free of torch so the kernels can be called without it. +""" + +import triton +import triton.language as tl + +from aiter.ops.triton.utils._triton.kernel_repr import make_kernel_repr + +_delta_v4_kernel_repr = make_kernel_repr( + "_delta_v4_kernel", + [ + "D", + "BLOCK_R", + ], +) + + +@triton.jit(repr=_delta_v4_kernel_repr) +def _delta_v4_kernel( + O_ptr, # [n_rows, D] bf16 (rows = T*H, contiguous) + dO_ptr, # [n_rows, D] bf16 + Delta_ptr, # [n_rows] fp32 + n_rows, + D: tl.constexpr, + BLOCK_R: tl.constexpr, +): + """Grid (cdiv(n_rows, BLOCK_R),) — each program reduces BLOCK_R rows of width D.""" + pid = tl.program_id(0) + rows = pid * BLOCK_R + tl.arange(0, BLOCK_R) + mask = rows < n_rows + offs = rows.to(tl.int64)[:, None] * D + tl.arange(0, D)[None, :] + o = tl.load(O_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + d = tl.load(dO_ptr + offs, mask=mask[:, None], other=0.0).to(tl.float32) + tl.store(Delta_ptr + rows, tl.sum(o * d, axis=1), mask=mask) + + +_bwd_dkv_gather_acc_v4_repr = make_kernel_repr( + "_bwd_dkv_gather_acc_v4", + [ + "D", + "BLOCK_E", + "ACCUMULATE", + ], +) + + +@triton.jit(repr=_bwd_dkv_gather_acc_v4_repr) +def _bwd_dkv_gather_acc_v4( + Interm_ptr, # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D] + InvPtr_ptr, # [num_kv+1] int32 — CSR row pointers + InvData_ptr, # [valid] int32 — encoded q*R_CHUNK+local_r, sorted by KV token + dKV_acc_ptr, # [num_kv, D] fp32 — accumulator + stride_interm_r: tl.int64, + stride_acc_t: tl.int64, + D: tl.constexpr, + BLOCK_E: tl.constexpr, + ACCUMULATE: tl.constexpr, +): + """Grid (num_kv,) — one CTA per KV token, BLOCK_E CSR entries in flight. + + ``BLOCK_E`` entries are carried per iteration, which the gather needs for two reasons: + + * **load width.** A bare ``tl.arange(0, D)`` block over 256 threads is 2 bf16 = 4 B per + lane -- a dword. The [BLOCK_E, D] block gives ``BLOCK_E*D/threads`` elements per lane, + so the loads become dwordx4. The gather is issue-bound, so this dominates. + * **trip count.** ``tl.sum`` folds the entry axis, so the run is consumed BLOCK_E at a + time. A realistic top-k gives run lengths up to ~3000 on the pool rows. + + ``ACCUMULATE=False`` writes the destination instead of reading it back first. The caller + uses it for the first chunk, where the accumulator is still zero. + """ + k = tl.program_id(0) + offs_d = tl.arange(0, D) + offs_e = tl.arange(0, BLOCK_E) + start = tl.load(InvPtr_ptr + k) + end = tl.load(InvPtr_ptr + k + 1) + acc_base = k.to(tl.int64) * stride_acc_t + + if ACCUMULATE: + acc = tl.load(dKV_acc_ptr + acc_base + offs_d).to(tl.float32) + else: + acc = tl.zeros([D], dtype=tl.float32) + + for i0 in range(start, end, BLOCK_E): + idx = i0 + offs_e + m = idx < end + entry = tl.load(InvData_ptr + idx, mask=m, other=0).to(tl.int64) + vals = tl.load( + Interm_ptr + entry[:, None] * stride_interm_r + offs_d[None, :], + mask=m[:, None], + other=0.0, + ) + acc += tl.sum(vals.to(tl.float32), axis=0) + + tl.store(dKV_acc_ptr + acc_base + offs_d, acc) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index f2e77125d8..9ba1042330 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -37,18 +37,18 @@ 3.380 ms total = 407 TFLOPS. """ +from dataclasses import dataclass + import torch +import triton from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - build_inverted_topk, - delta_v4, - dkv_gather_acc, -) -from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, + _dkv_interm_v4_kernel, + _dq_v4_kernel, ) -from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - sparse_mla_bwd_dq as _dq_gluon, +from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4_bwd import ( + _bwd_dkv_gather_acc_v4, + _delta_v4_kernel, ) from aiter.ops.triton.utils._triton import arch_info @@ -58,39 +58,240 @@ _TILE_K_DKV = 128 -def sparse_mla_bwd_dsv4( +def sparse_mla_bwd_dq( q, kv, do, - o, + topk, lse, - topk_indices, - attn_sink=None, - scale=None, - R_CHUNK=None, + delta, + dq, + chunk_dS, + chunk_P, + scale, + r_start, + R_CHUNK, + BLOCK_H=64, + TILE_K=32, + is_first_chunk=True, ): - """Backward for the DSv4 sparse-MLA prefill attention. gfx950 (CDNA4) only. + """Launch the dQ kernel for one rank chunk. Writes ``dq`` (RMW when not the first chunk) + plus this chunk's ``chunk_dS`` / ``chunk_P``.""" + T, H, D = q.shape + _dq_v4_kernel[(T, triton.cdiv(H, BLOCK_H))]( + q, + kv, + do, + topk, + lse, + delta, + dq, + chunk_dS, + chunk_P, + q.stride(0), + q.stride(1), + kv.stride(0), + do.stride(0), + do.stride(1), + dq.stride(0), + dq.stride(1), + topk.stride(0), + chunk_dS.stride(0), + chunk_dS.stride(1), + scale, + H, + r_start, + R_CHUNK=R_CHUNK, + BLOCK_H=BLOCK_H, + TILE_K=TILE_K, + D=D, + IS_FIRST_CHUNK=is_first_chunk, + num_warps=4, + waves_per_eu=1, + ) - Args: - q: [T, H, 512] bf16 - kv: [num_kv, 512] bf16, K == V. ``num_kv >= T``; rows ``T..num_kv-1`` are - the compressed pool, which ``topk_indices`` may reference. - do: [T, H, 512] bf16, gradient of the attention output - o: [T, H, 512] bf16, the forward output - lse: [T, H] fp32, sink-inclusive log-sum-exp from the forward - topk_indices: [T, TOPK] int32, -1 marks an invalid slot - attn_sink: [H] fp32 per-head sink bias, or None - scale: softmax scale, defaults to ``1/sqrt(512)`` - R_CHUNK: split the rank dimension into chunks of this width. ``None`` (default) - runs unchunked, which is what you want. Chunking exists only to bound the - ``interm`` intermediate, which is ``T*TOPK*512`` bf16 (2.0 GiB at - T=4096, TOPK=512); it costs a dQ read-modify-write between chunks and one - CSR build per chunk. Must be a multiple of 32 (the mfma tile width) and - must divide ``TOPK`` -- a partial tail chunk is rejected rather than - handled, since the chunk width is a kernel constexpr. - Returns: - dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) +def sparse_mla_bwd_dkv_interm_v4( + q, + do, + chunk_dS, + chunk_P, + R_CHUNK, + BD=256, + TILE_K=128, + MFMA_K=32, + DUAL_STAGE=1, + H_POW2=None, + num_warps=4, + interm=None, +): + """V4 dKV-interm, Q/dO read once. Returns interm [T, R_CHUNK, D] bf16. + + ``BD`` splits D across ``grid.y``; dS/P are re-read once per D block, so a larger BD moves + less of them. ``MFMA_K=32`` is the CDNA4 16x16x32 depth. + """ + T, H, D = q.shape + assert R_CHUNK % TILE_K == 0 + assert D % BD == 0 + h_pow2 = H_POW2 or triton.next_power_of_2(H) + if interm is None: + interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) + _dkv_interm_v4_kernel[(T, D // BD)]( + q, + do, + chunk_dS, + chunk_P, + interm, + q.stride(0), + q.stride(1), + do.stride(0), + do.stride(1), + chunk_dS.stride(0), + chunk_dS.stride(1), + interm.stride(0), + interm.stride(1), + H, + R_CHUNK=R_CHUNK, + TILE_K=TILE_K, + NH=h_pow2, + BD=BD, + D=D, + MFMA_K=MFMA_K, + DUAL_STAGE=DUAL_STAGE, + num_warps=num_warps, + ) + return interm + + +def delta_v4(o, do, out=None, BLOCK_R=8, num_warps=8): + # BLOCK_R=8 keeps each lane loading >= 8 bf16, i.e. a dwordx4; narrower blocks drop to a + # dword and the kernel loses most of its bandwidth. + """o[T,H,D] bf16, do[T,H,D] bf16 -> delta[T,H] fp32 = sum_d o*do. + + ``do`` must already be the D-wide (lora) slice, contiguous — same contract as the dQ kernel. + """ + assert o.shape == do.shape and o.is_contiguous() and do.is_contiguous() + T, H, D = o.shape + n_rows = T * H + if out is None: + out = torch.empty(T, H, dtype=torch.float32, device=o.device) + _delta_v4_kernel[(triton.cdiv(n_rows, BLOCK_R),)]( + o, + do, + out, + n_rows, + D=D, + BLOCK_R=BLOCK_R, + num_warps=num_warps, + ) + return out + + +def build_inverted_topk(topk_indices_slice, num_kv): + """CSR inverted index over ``num_kv`` KV rows. + + One stable sort yields both the permutation (``inv_data``) and the sorted keys; + ``inv_ptr[k] = searchsorted(sorted, k, 'left')`` = the number of entries with value < k. + Invalid (-1) entries sort to the front, so ``inv_ptr[0]`` starts past them and they are + never visited. + + The sort key is narrowed to int16 when ``num_kv`` fits, which is what keeps the radix sort + to two byte-passes. + + Returns ``inv_ptr[num_kv+1]`` int32, ``inv_data[T*R]`` int32. + """ + # row_ids is the searchsorted query: [0 .. num_kv], one per KV row plus the end sentinel. + # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by + # accident. + flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid + if num_kv < 32767: # int16 range, -1 included + keys = flat_kv.to(torch.int16) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) + else: + keys = flat_kv.to(torch.int32) + row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int32) + sorted_vals, inv_data = torch.sort(keys, stable=True) + inv_ptr = torch.searchsorted(sorted_vals, row_ids).to(torch.int32) + return inv_ptr, inv_data.to(torch.int32) + + +def dkv_gather_acc( + interm, inv_ptr, inv_data, dkv_acc, BLOCK_E=64, num_warps=8, accumulate=True +): + """interm[T,R,D] bf16 -> dkv_acc[num_kv,D] fp32 via the entry-blocked CSR gather. + + Grid is ``num_kv`` (from ``dkv_acc``), not ``T``, so a compressed-pool KV works. + """ + _, _, D = interm.shape + num_kv = dkv_acc.shape[0] + _bwd_dkv_gather_acc_v4[(num_kv,)]( + interm, + inv_ptr, + inv_data, + dkv_acc, + interm.stride(1), + dkv_acc.stride(0), + D=D, + BLOCK_E=BLOCK_E, + ACCUMULATE=accumulate, + num_warps=num_warps, + ) + + +@dataclass +class _BwdPlan: + """Validated inputs, tile choices and workspace for one backward call. + + Holds the values that flow between phases (``delta``, the CSR pair, the accumulator) so + ``bwd_phases`` can hand out independent thunks that still compose into the real pipeline. + """ + + q: torch.Tensor + kv: torch.Tensor + do: torch.Tensor + o: torch.Tensor + lse: torch.Tensor + topk_indices: torch.Tensor + attn_sink: object + scale: float + T: int + H: int + D: int + TOPK: int + num_kv: int + R_CHUNK: int + tk_dq: int + tk_dkv: int + delta: torch.Tensor + dq: torch.Tensor + dkv_acc: torch.Tensor + chunk_dS: torch.Tensor + chunk_P: torch.Tensor + interm: torch.Tensor + inv_ptr: object = None + inv_data: object = None + d_sink: object = None + + def build_csr(self, r): + self.inv_ptr, self.inv_data = build_inverted_topk( + self.topk_indices[:, r : r + self.R_CHUNK], self.num_kv + ) + + def compute_d_sink(self): + # d_sink[h] = -sum_t exp(sink[h] - lse[t,h]) * delta[t,h] + self.d_sink = -( + torch.exp(self.attn_sink[None, :].float() - self.lse) * self.delta + ).sum(dim=0) + + def result(self): + return self.dq, self.dkv_acc.to(self.kv.dtype), self.d_sink + + +def plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None): + """Validate the inputs, pick the tile widths and allocate the workspace for one call. + + Split out from ``sparse_mla_bwd_dsv4`` so the op benchmark can build the same plan and then + time ``bwd_phases`` against it. Argument meanings are documented on the public entry. """ assert ( arch_info.get_arch() == "gfx950" @@ -163,51 +364,132 @@ def sparse_mla_bwd_dsv4( "of each top-k row and desynchronizes the CSR index space from interm" ) - delta = delta_v4(o, do) - - dq = torch.empty_like(q) - dkv_acc = torch.zeros(num_kv, D, dtype=torch.float32, device=q.device) - chunk_dS = torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device) - chunk_P = torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device) - interm = torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device) - - for r in range(0, TOPK, R_CHUNK): - _dq_gluon( - q, - kv, - do, - topk_indices, - lse, - delta, - dq, - chunk_dS, - chunk_P, - scale, - r, - R_CHUNK, - BLOCK_H=_BLOCK_H_DQ, - TILE_K=tk_dq, - is_first_chunk=(r == 0), - ) - _dkv_interm_gluon( - q, do, chunk_dS, chunk_P, R_CHUNK, BD=_BD_DKV, TILE_K=tk_dkv, interm=interm + return _BwdPlan( + q=q, + kv=kv, + do=do, + lse=lse, + topk_indices=topk_indices, + attn_sink=attn_sink, + scale=scale, + T=T, + H=H, + D=D, + TOPK=TOPK, + num_kv=num_kv, + R_CHUNK=R_CHUNK, + tk_dq=tk_dq, + tk_dkv=tk_dkv, + delta=torch.empty(T, H, dtype=torch.float32, device=q.device), + dq=torch.empty_like(q), + dkv_acc=torch.zeros(num_kv, D, dtype=torch.float32, device=q.device), + chunk_dS=torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device), + chunk_P=torch.empty(T, H, R_CHUNK, dtype=torch.bfloat16, device=q.device), + interm=torch.empty(T, R_CHUNK, D, dtype=torch.bfloat16, device=q.device), + o=o, + ) + + +def bwd_phases(plan): + """``(name, thunk)`` for every phase of one backward call, in execution order. + + ``sparse_mla_bwd_dsv4`` runs these; the op benchmark times them one at a time. Sharing the + sequence is the point -- a per-phase measurement is only meaningful if it is measuring the + phases the op actually runs, with the same tile configuration and the same workspace. + """ + yield "delta", lambda: delta_v4(plan.o, plan.do, out=plan.delta) + + for r in range(0, plan.TOPK, plan.R_CHUNK): + yield ( + "dq", + lambda r=r: sparse_mla_bwd_dq( + plan.q, + plan.kv, + plan.do, + plan.topk_indices, + plan.lse, + plan.delta, + plan.dq, + plan.chunk_dS, + plan.chunk_P, + plan.scale, + r, + plan.R_CHUNK, + BLOCK_H=_BLOCK_H_DQ, + TILE_K=plan.tk_dq, + is_first_chunk=(r == 0), + ), ) - inv_ptr, inv_data = build_inverted_topk( - topk_indices[:, r : r + R_CHUNK], num_kv + yield ( + "interm", + lambda: sparse_mla_bwd_dkv_interm_v4( + plan.q, + plan.do, + plan.chunk_dS, + plan.chunk_P, + plan.R_CHUNK, + BD=_BD_DKV, + TILE_K=plan.tk_dkv, + interm=plan.interm, + ), ) + yield "csr_build", lambda r=r: plan.build_csr(r) # dkv_acc was just zeroed, so the first chunk can write instead of read-modify-write. # Unchunked -- the default -- that is the only chunk, and it saves reading back the # whole [num_kv, 512] fp32 accumulator. - dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc, accumulate=(r > 0)) + yield ( + "gather", + lambda r=r: dkv_gather_acc( + plan.interm, + plan.inv_ptr, + plan.inv_data, + plan.dkv_acc, + accumulate=(r > 0), + ), + ) - dkv = dkv_acc.to(kv.dtype) + if plan.attn_sink is not None: + yield "d_sink", plan.compute_d_sink - d_sink = None - if attn_sink is not None: - # d_sink[h] = -sum_t exp(sink[h] - lse[t,h]) * delta[t,h] - d_sink = -(torch.exp(attn_sink[None, :].float() - lse) * delta).sum(dim=0) - return dq, dkv, d_sink +def sparse_mla_bwd_dsv4( + q, + kv, + do, + o, + lse, + topk_indices, + attn_sink=None, + scale=None, + R_CHUNK=None, +): + """Backward for the DSv4 sparse-MLA prefill attention. gfx950 (CDNA4) only. + + Args: + q: [T, H, 512] bf16 + kv: [num_kv, 512] bf16, K == V. ``num_kv >= T``; rows ``T..num_kv-1`` are + the compressed pool, which ``topk_indices`` may reference. + do: [T, H, 512] bf16, gradient of the attention output + o: [T, H, 512] bf16, the forward output + lse: [T, H] fp32, sink-inclusive log-sum-exp from the forward + topk_indices: [T, TOPK] int32, -1 marks an invalid slot + attn_sink: [H] fp32 per-head sink bias, or None + scale: softmax scale, defaults to ``1/sqrt(512)`` + R_CHUNK: split the rank dimension into chunks of this width. ``None`` (default) + runs unchunked, which is what you want. Chunking exists only to bound the + ``interm`` intermediate, which is ``T*TOPK*512`` bf16 (2.0 GiB at + T=4096, TOPK=512); it costs a dQ read-modify-write between chunks and one + CSR build per chunk. Must be a multiple of 32 (the mfma tile width) and + must divide ``TOPK`` -- a partial tail chunk is rejected rather than + handled, since the chunk width is a kernel constexpr. + + Returns: + dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) + """ + plan = plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink, scale, R_CHUNK) + for _, run in bwd_phases(plan): + run() + return plan.result() __all__ = ["sparse_mla_bwd_dsv4"] diff --git a/aiter/test_mha_common.py b/aiter/test_mha_common.py index d4b4cdda89..20e5675de4 100644 --- a/aiter/test_mha_common.py +++ b/aiter/test_mha_common.py @@ -678,3 +678,45 @@ def _tol(ref_val, pt_val, is_forward=False): bwd_tols = [_tol(dq, dq_pt), _tol(dk, dk_pt), _tol(dv, dv_pt)] return out, (dq, dk, dv), fwd_tol, bwd_tols + + +def sparse_mla_dsv4_ref(q, kv, topk_indices, attn_sink=None, scale=None): + """fp32 reference forward for the DeepSeek-V4 sparse MLA attention. + + The official V4 form: shared-KV GQA where ``K == V == kv`` is one dense ``head_dim``-wide + tensor, RoPE already applied in place caller-side, ``attn_sink`` folded into the softmax + denominator only, and ``topk_indices == -1`` masked out. + + Differentiable in ``q`` / ``kv`` / ``attn_sink``, so a backward reference is just autograd + through this. It is a per-token python loop, so keep the shapes small. + + Args: + q: [T, H, D] float32, requires_grad for a backward reference + kv: [num_kv, D] float32, ``num_kv >= T``; rows ``T..`` are the compressed pool + topk_indices: [T, TOPK] int, -1 marks an invalid slot + attn_sink: [H] float32 or None + scale: softmax scale, defaults to ``1/sqrt(D)`` + + Returns: + ``(o, lse)`` -- ``o`` [T, H, D] float32 carrying grad, and ``lse`` [T, H] float32 + detached and sink-inclusive, which is the convention the backward kernel expects. + """ + if scale is None: + scale = 1.0 / (q.shape[-1] ** 0.5) + outs, lses = [], [] + for t in range(q.shape[0]): + idx = topk_indices[t].long() + valid = idx != -1 + k = kv[idx.clamp(min=0)] + k = torch.where(valid[:, None], k, torch.zeros_like(k)) + s = (q[t] @ k.t()) * scale + s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) + row_max = s.max(dim=1).values + m = row_max if attn_sink is None else torch.maximum(row_max, attn_sink) + p = torch.where(valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s)) + denom = p.sum(dim=1) + if attn_sink is not None: + denom = denom + torch.exp(attn_sink - m) + outs.append((p @ k) / denom[:, None]) + lses.append((m + torch.log(denom)).detach()) + return torch.stack(outs, dim=0), torch.stack(lses, dim=0) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py index 002c507a77..3c141b5792 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -24,19 +24,13 @@ import torch import triton -from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - build_inverted_topk, - delta_v4, - dkv_gather_acc, +from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import ( + bwd_phases, + plan_bwd, + sparse_mla_bwd_dsv4, ) -from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - sparse_mla_bwd_dkv_interm_v4 as _dkv_interm_gluon, -) -from aiter.ops.triton._gluon_kernels.gfx950.attention.sparse_attention_dsv4_bwd import ( - sparse_mla_bwd_dq as _dq_gluon, -) -from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 from aiter.ops.triton.utils._triton import arch_info +from aiter.test_mha_common import sparse_mla_dsv4_ref D = 512 SWA = 128 # sliding-window width @@ -108,21 +102,6 @@ def _bf16(*shape): # --------------------------------------------------------------------------- # Timing # --------------------------------------------------------------------------- -def _bench(fn, *, warmup=5, reps=20): - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - ev0 = torch.cuda.Event(enable_timing=True) - ev1 = torch.cuda.Event(enable_timing=True) - ev0.record() - for _ in range(reps): - fn() - ev1.record() - torch.cuda.synchronize() - return ev0.elapsed_time(ev1) / reps - - def _flops(T, H, topk): # dQ contributes S, dP and dS@kv; dKV-interm contributes dS*Q and P*dO. Five 2*D-flop # products per (token, head, top-k slot). @@ -130,100 +109,29 @@ def _flops(T, H, topk): def _time_phases(case, scale): - """Per-phase timings, in the order the public entry runs them.""" - q, kv, do, o = case["q"], case["kv"], case["do"], case["o"] - lse, sink, indices = case["lse"], case["sink"], case["indices"] - T, H, _ = q.shape - topk = indices.shape[1] - num_kv = case["num_kv"] - - delta = delta_v4(o, do) - dq = torch.empty_like(q) - dkv_acc = torch.zeros(num_kv, D, dtype=torch.float32, device=q.device) - chunk_dS = torch.empty(T, H, topk, dtype=torch.bfloat16, device=q.device) - chunk_P = torch.empty(T, H, topk, dtype=torch.bfloat16, device=q.device) - interm = torch.empty(T, topk, D, dtype=torch.bfloat16, device=q.device) - inv_ptr, inv_data = build_inverted_topk(indices, num_kv) - - phases = [ - ("delta", lambda: delta_v4(o, do)), - ( - "dq", - lambda: _dq_gluon( - q, - kv, - do, - indices, - lse, - delta, - dq, - chunk_dS, - chunk_P, - scale, - 0, - topk, - BLOCK_H=64, - TILE_K=32, - is_first_chunk=True, - ), - ), - ( - "interm", - lambda: _dkv_interm_gluon( - q, do, chunk_dS, chunk_P, topk, BD=256, TILE_K=128, interm=interm - ), - ), - ("csr_build", lambda: build_inverted_topk(indices, num_kv)), - ("gather", lambda: dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc)), - ( - "d_sink", - lambda: -(torch.exp(sink[None, :].float() - lse) * delta).sum(dim=0), - ), - ] - return [(name, _bench(fn)) for name, fn in phases] + """Time each phase of the real pipeline. + + `plan_bwd` and `bwd_phases` are the same helpers `sparse_mla_bwd_dsv4` runs, so this cannot + drift from the op: the tile widths, the workspace and the phase order all come from the + wrapper rather than being restated here. Timing a phase repeatedly re-runs its side effects, + which is harmless -- only the duration is read. + """ + plan = plan_bwd( + case["q"], + case["kv"], + case["do"], + case["o"], + case["lse"], + case["indices"], + case["sink"], + scale, + ) + return [(name, triton.testing.do_bench(run)) for name, run in bwd_phases(plan)] # --------------------------------------------------------------------------- # Correctness gate # --------------------------------------------------------------------------- -def _ref_scores(q_t, kv_f, idx, scale): - """Masked fp32 scores for one token: [H, topk], invalid slots -inf. Also returns the rows.""" - valid = idx != -1 - k = kv_f[idx.clamp(min=0)] - k = torch.where(valid[:, None], k, torch.zeros_like(k)) - s = (q_t @ k.t()) * scale - return torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))), k, valid - - -def _ref_forward(q, kv, sink, topk, scale): - """fp32 reference forward. Per-token python loop, so small shapes only. - - Returns `o` and the leaf tensors, so the caller can drive autograd through it. - """ - qg = q.float().clone().requires_grad_(True) - kvg = kv.float().clone().requires_grad_(True) - sg = sink.clone().requires_grad_(True) - outs = [] - for t in range(q.shape[0]): - s, k, valid = _ref_scores(qg[t], kvg, topk[t].long(), scale) - m = torch.maximum(s.max(dim=1).values, sg) - p = torch.where(valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s)) - denom = p.sum(dim=1) + torch.exp(sg - m) - outs.append((p @ k) / denom[:, None]) - return torch.stack(outs, dim=0), (qg, kvg, sg) - - -def _ref_lse(q, kv, sink, topk, scale): - """Sink-inclusive log-sum-exp matching the forward's convention.""" - with torch.no_grad(): - kv_f = kv.float() - rows = [] - for t in range(q.shape[0]): - s, _, _ = _ref_scores(q[t].float(), kv_f, topk[t].long(), scale) - rows.append(torch.logsumexp(torch.cat([s, sink[:, None]], dim=1), dim=1)) - return torch.stack(rows) - - def check_correctness(device): """Small-shape gate against autograd, so a broken kernel fails before it is timed.""" print("\n========== CORRECTNESS ==========") @@ -238,9 +146,11 @@ def check_correctness(device): case["indices"], ) - ref_o, (qg, kvg, sg) = _ref_forward(q, kv, sink, indices, scale) + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) + ref_o, lse = sparse_mla_dsv4_ref(qg, kvg, indices, sg, scale) o = ref_o.detach().to(torch.bfloat16).contiguous() - lse = _ref_lse(q, kv, sink, indices, scale) dq, dkv, d_sink = sparse_mla_bwd_dsv4( q, kv, do, o, lse, indices, attn_sink=sink, scale=scale @@ -287,7 +197,7 @@ def run_bwd_bench(args, device): scale = 1.0 / (D**0.5) tflops = _flops(T, H, topk) / 1e12 - ms = _bench( + ms = triton.testing.do_bench( lambda c=case, s=scale: sparse_mla_bwd_dsv4( c["q"], c["kv"], @@ -325,13 +235,19 @@ def _parse_args(): "4096,128,1024", "8192,128,512", ], + help="one or more shapes as T,H,topk (e.g. 4096,128,512). topk must be >= the sliding " + "window (128); the remainder is drawn from the compressed pool.", ) p.add_argument( "--breakdown", action="store_true", help="also time each phase separately (the per-kernel table in the PR description)", ) - p.add_argument("--skip-correctness", action="store_true") + p.add_argument( + "--skip-correctness", + action="store_true", + help="skip the small-shape autograd check that otherwise runs before any timing.", + ) args = p.parse_args() args.cfgs = [tuple(int(x) for x in s.split(",")) for s in args.cfgs] return args diff --git a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py index b9a5abc12b..4493048e9b 100644 --- a/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py +++ b/op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py @@ -3,10 +3,11 @@ """Correctness for the DSv4 sparse-MLA training backward against torch autograd. -The reference is a differentiable fp32 re-implementation of the V4 forward, differentiated by -autograd -- an independent path from the kernels, not a re-expression of them. It is a per-token -python loop, so the shapes here are deliberately small; the kernels are exercised at production -shapes (T=4096, H=128, topk=512) out of tree. +The reference is ``aiter.test_mha_common.sparse_mla_dsv4_ref``, a differentiable fp32 +re-implementation of the V4 forward -- an independent path from the kernels, not a re-expression +of them. Backward values come from autograd through it. It is a per-token python loop, so the +shapes here are deliberately small; the op benchmark exercises the kernels at production shapes +and shares the same reference for its correctness gate. """ import pytest @@ -14,6 +15,7 @@ from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import sparse_mla_bwd_dsv4 from aiter.ops.triton.utils._triton import arch_info +from aiter.test_mha_common import sparse_mla_dsv4_ref D = 512 COS_TOL = 0.999 @@ -24,55 +26,6 @@ ) -def _ref_fwd_diff(q, kv, attn_sink, topk, scale): - """Differentiable fp32 V4 forward. q/kv/attn_sink require grad. Returns O [T, H, D].""" - outs = [] - for t in range(q.shape[0]): - idx = topk[t].long() - valid = idx != -1 - k = kv[idx.clamp(min=0)] - k = torch.where(valid[:, None], k, torch.zeros_like(k)) - s = (q[t] @ k.t()) * scale - s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) - if attn_sink is None: - m = s.max(dim=1).values - p = torch.where( - valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) - ) - denom = p.sum(dim=1) - else: - m = torch.maximum(s.max(dim=1).values, attn_sink) - p = torch.where( - valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) - ) - denom = p.sum(dim=1) + torch.exp(attn_sink - m) - outs.append((p @ k) / denom[:, None]) - return torch.stack(outs, dim=0) - - -def _ref_fwd(q, kv, attn_sink, topk, scale): - """Non-differentiable forward giving the kernel inputs O (bf16) and sink-inclusive lse.""" - with torch.no_grad(): - o = _ref_fwd_diff(q.float(), kv.float(), attn_sink, topk, scale) - lse = torch.empty(q.shape[0], q.shape[1], device=q.device, dtype=torch.float32) - for t in range(q.shape[0]): - idx = topk[t].long() - valid = idx != -1 - k = kv.float()[idx.clamp(min=0)] - k = torch.where(valid[:, None], k, torch.zeros_like(k)) - s = (q.float()[t] @ k.t()) * scale - s = torch.where(valid[None, :], s, torch.full_like(s, float("-inf"))) - m = s.max(dim=1).values - p = torch.where( - valid[None, :], torch.exp(s - m[:, None]), torch.zeros_like(s) - ) - denom = p.sum(dim=1) - if attn_sink is not None: - denom = denom + torch.exp(attn_sink - m) - lse[t] = m + torch.log(denom) - return o.to(torch.bfloat16).contiguous(), lse - - def _cos(a, b): return torch.nn.functional.cosine_similarity( a.float().reshape(-1), b.float().reshape(-1), dim=0 @@ -104,16 +57,17 @@ def test_sparse_mla_bwd_dsv4(T, H, topk, npool, has_sink, r_chunk): invalid = torch.rand(T, topk, device=dev) < 0.1 indices = torch.where(invalid, torch.full_like(indices, -1), indices).contiguous() - o, lse = _ref_fwd(q, kv, sink, indices, scale) + qg = q.float().clone().requires_grad_(True) + kvg = kv.float().clone().requires_grad_(True) + sg = sink.clone().requires_grad_(True) if has_sink else None + ref_o, lse = sparse_mla_dsv4_ref(qg, kvg, indices, sg, scale) + o = ref_o.detach().to(torch.bfloat16).contiguous() dq, dkv, d_sink = sparse_mla_bwd_dsv4( q, kv, do, o, lse, indices, attn_sink=sink, scale=scale, R_CHUNK=r_chunk ) - qg = q.float().clone().requires_grad_(True) - kvg = kv.float().clone().requires_grad_(True) - sg = sink.clone().requires_grad_(True) if has_sink else None - _ref_fwd_diff(qg, kvg, sg, indices, scale).backward(do.float()) + ref_o.backward(do.float()) assert _cos(dq, qg.grad) > COS_TOL, f"dq cos {_cos(dq, qg.grad)}" assert _cos(dkv, kvg.grad) > COS_TOL, f"dkv cos {_cos(dkv, kvg.grad)}" From 5abadaff5e82eabb0b9ba3afb687b69e1ba501b1 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Mon, 24 Aug 2026 15:20:35 -0500 Subject: [PATCH 12/15] Widen the int16 CSR cutoff and drop an unused benchmark parameter Both from review. `build_inverted_topk` gated the int16 sort key on `num_kv < 32767`. The largest value that has to fit is `num_kv` itself (`row_ids` runs to `num_kv` inclusive), and that is exactly `torch.iinfo(torch.int16).max`, so the boundary case was falling back to an int32 key -- an 8-pass radix sort instead of 2 -- for no reason. Now `<= torch.iinfo(torch.int16).max`, written against the type rather than a literal. `_print_table` took a `title` it never used, carried over from the forward's benchmark. Dropped; the section header is already printed by the caller. 8/8 tests pass on gfx950 (MI355X). --- aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py | 4 +++- .../op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 9ba1042330..0baf505d95 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -204,7 +204,9 @@ def build_inverted_topk(topk_indices_slice, num_kv): # Its dtype must match `keys` -- searchsorted is built per branch for that reason, not by # accident. flat_kv = topk_indices_slice.reshape(-1) # [T*R] int32; -1 = invalid - if num_kv < 32767: # int16 range, -1 included + # row_ids reaches num_kv, so int16 holds it as long as num_kv is within the type; + # the -1 sentinel is fine either way. + if num_kv <= torch.iinfo(torch.int16).max: keys = flat_kv.to(torch.int16) row_ids = torch.arange(num_kv + 1, device=flat_kv.device, dtype=torch.int16) else: diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py index 3c141b5792..3225ff34e2 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -175,7 +175,7 @@ def _cos(a, b): # --------------------------------------------------------------------------- # Reporting # --------------------------------------------------------------------------- -def _print_table(title, headers, rows): +def _print_table(headers, rows): def _fmt(x): if isinstance(x, float): return f"{x:.3f}" if x >= 1 or x == 0 else f"{x:.4f}" @@ -221,7 +221,7 @@ def run_bwd_bench(args, device): f" {'SUM':10s} {total:7.3f} ms -> {tflops / (total * 1e-3):.0f} TFLOPS" ) - _print_table("BACKWARD", ["T", "H", "Kv", "topk", "ms", "TFLOPS"], rows) + _print_table(["T", "H", "Kv", "topk", "ms", "TFLOPS"], rows) def _parse_args(): From 6993fcc11b4122a16169e9bef08bd5da42131aba Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Mon, 24 Aug 2026 21:01:58 -0500 Subject: [PATCH 13/15] Reject host tensors, and give a fully-masked row a defined output Both from review. `plan_bwd` checked the arch but never that the tensors were on device, so a host tensor got as far as the kernel launch and failed there. It now raises the same way the forward's entry does, naming the offending argument and its device. `sparse_mla_dsv4_ref` produced NaN for a row whose top-k is entirely -1 with no sink: every contributor is masked, the denominator is 0, and the output was 0/0. Such a row now yields 0, with the lse left at -inf, which is the honest value for a row with nothing in it. The V4 top-k always keeps at least the token itself so the in-tree callers never reach this, but the reference is shared from `test_mha_common` now and should not hand a caller NaN for a case it can legitimately construct. 8/8 tests pass on gfx950 (MI355X); benchmark unchanged at 3.229 ms / 426 TFLOPS on an idle GPU. --- .../triton/attention/sparse_attention_dsv4_bwd.py | 13 +++++++++++++ aiter/test_mha_common.py | 8 +++++++- 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 0baf505d95..73f3e29bf8 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -299,6 +299,19 @@ def plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUN arch_info.get_arch() == "gfx950" ), f"sparse_mla_bwd_dsv4 requires gfx950 (CDNA4), got {arch_info.get_arch()}" + for name, t in ( + ("q", q), + ("kv", kv), + ("do", do), + ("o", o), + ("lse", lse), + ("topk_indices", topk_indices), + ): + if not t.is_cuda: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 requires CUDA/HIP tensors, {name} is on {t.device}" + ) + if q.dtype != torch.bfloat16: raise RuntimeError(f"sparse_mla_bwd_dsv4 expects bf16 q, got {q.dtype}") for name, t in (("kv", kv), ("do", do), ("o", o)): diff --git a/aiter/test_mha_common.py b/aiter/test_mha_common.py index 3bda903d8d..97ac286ff2 100644 --- a/aiter/test_mha_common.py +++ b/aiter/test_mha_common.py @@ -746,6 +746,7 @@ def sparse_mla_dsv4_ref(q, kv, topk_indices, attn_sink=None, scale=None): Returns: ``(o, lse)`` -- ``o`` [T, H, D] float32 carrying grad, and ``lse`` [T, H] float32 detached and sink-inclusive, which is the convention the backward kernel expects. + A row whose top-k is entirely -1 gets ``o = 0`` and ``lse = -inf``. """ if scale is None: scale = 1.0 / (q.shape[-1] ** 0.5) @@ -763,6 +764,11 @@ def sparse_mla_dsv4_ref(q, kv, topk_indices, attn_sink=None, scale=None): denom = p.sum(dim=1) if attn_sink is not None: denom = denom + torch.exp(attn_sink - m) - outs.append((p @ k) / denom[:, None]) + # A row whose top-k is entirely -1 has no contributors at all, and without a sink the + # denominator is then 0. Define that row's output as zero rather than letting 0/0 make + # it NaN; the lse stays -inf, which is the honest value for an empty row. + outs.append( + (p @ k) / torch.where(denom > 0, denom, torch.ones_like(denom))[:, None] + ) lses.append((m + torch.log(denom)).detach()) return torch.stack(outs, dim=0), torch.stack(lses, dim=0) From 77778c44716c6e314efc986f16fa5d6f74df50c6 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Mon, 24 Aug 2026 21:17:21 -0500 Subject: [PATCH 14/15] Reject a benchmark topk narrower than the sliding window From review. `_build_topk_swa_pool` documents a `[T, topk]` result but returned the full `[T, SWA]` window whenever `topk <= SWA`, so `--cfgs 4096,128,64` would have reported topk=64 in the results table while timing 128 -- and computed its TFLOPS from the number it printed. The window is the floor of a V4 top-k, so a narrower request has no meaning; it is now an error rather than a silently different shape. The `--cfgs` help text already stated the constraint, which is exactly the sort of thing that should be enforced rather than documented. The `n_pool <= 0` branch becomes `n_pool == 0`, the only case it can still see. The default configurations are unaffected: all three have topk >= 128. --- .../triton/bench_sparse_attention_dsv4_bwd.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py index 3225ff34e2..60daa1571b 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -47,14 +47,20 @@ def _build_topk_swa_pool(T, topk, num_pool, device, generator): KV row layout is ``[ per-token 0..T-1 | pool T..T+num_pool-1 ]``. Early tokens see fewer than `topk - SWA` pool blocks, so their trailing slots stay -1, exactly as in production. + + ``topk`` must be at least the window width: the window is the floor of a V4 top-k, and a + narrower request has no meaning here. Returning the full window anyway would report one + ``topk`` in the results table while timing another. """ + assert topk >= SWA, f"topk={topk} is narrower than the SWA({SWA}) window" + idx = torch.arange(T, device=device) off = torch.arange(SWA, device=device) swa = idx[:, None] - (SWA - 1) + off[None, :] swa = torch.where(swa >= 0, swa, torch.full_like(swa, -1)) n_pool = topk - SWA - if n_pool <= 0 or num_pool == 0: + if n_pool == 0 or num_pool == 0: return swa.to(torch.int32).contiguous() assert n_pool <= num_pool, f"need {n_pool} pool ranks but only {num_pool} blocks" From 564274e2fe09044c649de308b4906df84beb8284 Mon Sep 17 00:00:00 2001 From: Ye Wang Date: Mon, 24 Aug 2026 21:54:42 -0500 Subject: [PATCH 15/15] Validate attn_sink, and mark the plan helpers private Two suppressed review comments on the last round, both fair. `attn_sink` is optional, so it sat outside the device sweep added for the required tensors and was only shape-checked -- yet `compute_d_sink` does GPU math with it, so a host or non-fp32 sink failed late and unhelpfully. It is now checked the same way as the rest. `plan_bwd` / `bwd_phases` exist so the benchmark can time the phases the op actually runs, not as API. They read as public in a public wrapper module, so they are `_plan_bwd` / `_bwd_phases` now, leaving `sparse_mla_bwd_dsv4` as the only supported entry -- which is what `__all__` already said. The benchmark importing the private names matches what the other benchmarks here do with kernel-layer symbols. 8/8 tests pass on gfx950 (MI355X); benchmark 3.231 ms / 425 TFLOPS on an idle GPU. --- .../attention/sparse_attention_dsv4_bwd.py | 24 ++++++++++++++----- .../triton/bench_sparse_attention_dsv4_bwd.py | 10 ++++---- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py index 73f3e29bf8..bdec379db7 100644 --- a/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py +++ b/aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py @@ -245,7 +245,7 @@ class _BwdPlan: """Validated inputs, tile choices and workspace for one backward call. Holds the values that flow between phases (``delta``, the CSR pair, the accumulator) so - ``bwd_phases`` can hand out independent thunks that still compose into the real pipeline. + ``_bwd_phases`` can hand out independent thunks that still compose into the real pipeline. """ q: torch.Tensor @@ -289,11 +289,13 @@ def result(self): return self.dq, self.dkv_acc.to(self.kv.dtype), self.d_sink -def plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None): +def _plan_bwd( + q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUNK=None +): """Validate the inputs, pick the tile widths and allocate the workspace for one call. Split out from ``sparse_mla_bwd_dsv4`` so the op benchmark can build the same plan and then - time ``bwd_phases`` against it. Argument meanings are documented on the public entry. + time ``_bwd_phases`` against it. Argument meanings are documented on the public entry. """ assert ( arch_info.get_arch() == "gfx950" @@ -334,6 +336,16 @@ def plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUN assert lse.shape == (T, H), f"lse must be [{T}, {H}], got {tuple(lse.shape)}" assert num_kv >= T, f"num_kv ({num_kv}) must be >= T ({T})" if attn_sink is not None: + # Optional, so it misses the device sweep above; d_sink does GPU math with it. + if not attn_sink.is_cuda: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 requires CUDA/HIP tensors, attn_sink is on " + f"{attn_sink.device}" + ) + if attn_sink.dtype != torch.float32: + raise RuntimeError( + f"sparse_mla_bwd_dsv4 expects fp32 attn_sink, got {attn_sink.dtype}" + ) assert attn_sink.shape == ( H, ), f"attn_sink must be [{H}], got {tuple(attn_sink.shape)}" @@ -405,7 +417,7 @@ def plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink=None, scale=None, R_CHUN ) -def bwd_phases(plan): +def _bwd_phases(plan): """``(name, thunk)`` for every phase of one backward call, in execution order. ``sparse_mla_bwd_dsv4`` runs these; the op benchmark times them one at a time. Sharing the @@ -501,8 +513,8 @@ def sparse_mla_bwd_dsv4( Returns: dq [T, H, 512] bf16, dkv [num_kv, 512] bf16, d_sink [H] fp32 (None if no ``attn_sink``) """ - plan = plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink, scale, R_CHUNK) - for _, run in bwd_phases(plan): + plan = _plan_bwd(q, kv, do, o, lse, topk_indices, attn_sink, scale, R_CHUNK) + for _, run in _bwd_phases(plan): run() return plan.result() diff --git a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py index 60daa1571b..efb38e5581 100644 --- a/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py +++ b/op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py @@ -25,8 +25,8 @@ import triton from aiter.ops.triton.attention.sparse_attention_dsv4_bwd import ( - bwd_phases, - plan_bwd, + _bwd_phases, + _plan_bwd, sparse_mla_bwd_dsv4, ) from aiter.ops.triton.utils._triton import arch_info @@ -117,12 +117,12 @@ def _flops(T, H, topk): def _time_phases(case, scale): """Time each phase of the real pipeline. - `plan_bwd` and `bwd_phases` are the same helpers `sparse_mla_bwd_dsv4` runs, so this cannot + `_plan_bwd` and `_bwd_phases` are the same helpers `sparse_mla_bwd_dsv4` runs, so this cannot drift from the op: the tile widths, the workspace and the phase order all come from the wrapper rather than being restated here. Timing a phase repeatedly re-runs its side effects, which is harmless -- only the duration is read. """ - plan = plan_bwd( + plan = _plan_bwd( case["q"], case["kv"], case["do"], @@ -132,7 +132,7 @@ def _time_phases(case, scale): case["sink"], scale, ) - return [(name, triton.testing.do_bench(run)) for name, run in bwd_phases(plan)] + return [(name, triton.testing.do_bench(run)) for name, run in _bwd_phases(plan)] # ---------------------------------------------------------------------------