[Triton/Gluon] [GFX950][DSV4] Sparse MLA training backward by wangye805 · Pull Request #4766 · ROCm/aiter

wangye805 · 2026-08-14T19:44:02Z

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.

Contract

Same op 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, 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.

Full signature and the shapes it enforces:

sparse_mla_bwd_dsv4(q, kv, do, o, lse, topk_indices,
                    attn_sink=None, scale=None, R_CHUNK=None)  ->  (dq, dkv, d_sink)

arg	shape	dtype	notes
`q`, `do`, `o`	`[T, H, 512]`	bf16	contiguous
`kv`	`[num_kv, 512]`	bf16	`K == V`; `num_kv >= T`, rows `T..num_kv-1` 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; `d_sink` is `None` without it
`scale`			defaults to `1/sqrt(512)`
`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.

Structure

Five kernels plus one torch reduction, split across the gluon / triton / public-API trees as usual:

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 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

Reproduce with the op benchmark added in this PR:

python op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py --breakdown

It runs a small-shape autograd check first, then times the public entry end-to-end; --breakdown
additionally times each phase on its own. On gfx950 (MI355X):

T	H	Kv	topk	ms	TFLOPS
4096	128	5120	512	3.271	420
4096	128	5120	1024	5.261	522
8192	128	9216	512	6.177	445

Per-kernel at T=4096 H=128 topk=512, summing to 3.368 ms / 408 TFLOPS:

phase	ms	share
dQ	1.452	43.1%
dKV-interm	1.149	34.1%
dKV gather	0.373	11.1%
delta	0.207	6.1%
CSR build	0.160	4.7%
`d_sink`	0.027	0.8%

The breakdown runs the same plan_bwd / bwd_phases helpers the op itself runs, so it measures
the phases as configured and ordered in production rather than a restatement of them. Timings
come from triton.testing.do_bench, which flushes L2 between repetitions -- the per-phase
numbers are therefore cold-cache and sum slightly above the end-to-end figure.

The top-k the benchmark builds is a sliding window plus a causally-visible compressed pool, not
uniform random. That is not cosmetic: a uniform top-k spreads contributors evenly across the KV
rows, while the real distribution gives the pool rows runs of a few thousand, and the dKV gather
is sized for those runs. On uniform indices the gather looks about twice as fast as it is.

For reference, the Primus-Turbo FlyDSL backward
measured 3.223 ms / 426 TFLOPS on identical tensors in one process at the first shape above, so
the two are at parity. That comparison is not in the benchmark script: the FlyDSL backward has
no single entry point (it is driven phase by phase through compile() calls), so including it
would mean carrying third-party launch code that cannot run here. It was measured out of tree.

Environment

Absolute times above are MI355X; the kernels themselves are gfx950 and run unchanged on any CDNA4 part, but a lower-TDP SKU (e.g. MI350X) will report different wall-clock.


GPU	AMD Instinct MI355 OAM — `gfx950:sramecc+:xnack-` (CDNA4)
ROCm	7.2.1 (HIP `7.2.26054-7ab7bb831c`)
Python	3.11.14
PyTorch	`2.8.0+rocm7.2.1.git08d38866`
Triton	3.7.0, built from source — not the Triton shipped in the image
Container	`te-ci:rocm-7.2_ubuntu22.04_py3.11_pytorch_release-2.8_08d38866_jax_0.8.0_fa_2.8.1_aiter_77455e3ecf` (AMD-internal `registry-sc-harbor` CI image)
aiter	this branch, built from source

The source-built Triton is the one part worth calling out: these kernels use Gluon (buffer_load_to_shared, CDNA4 MFMA layouts) and want a Triton newer than the one baked into the container image. Reproducing on the stock image Triton has not been tried and may fail to compile rather than run slow.

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.

op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py  ......  6 passed

Please 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. The kernels use CDNA4 MFMA layouts and buffer_load_to_shared and are not portable to gfx942 as written; the public wrapper asserts the arch so other targets get a clear error rather than a Gluon compile failure.

github-actions · 2026-08-14T19:44:23Z

🏷️ CI Guide

Runs automatically on every PR:

✅ Pre-checks (submodule verification, code formatting)
✅ Aiter op tests (gfx942 + gfx950)
✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
`ci:gfx1250-ffm-triton`	Run the five-shard gfx1250 FFM Triton test suite
`ci:triton-300x`	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
`ci:sglang`	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
`ci:atom`	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
`ci:atom_full`	ATOM accuracy suite for PR and main models from ATOM `models_accuracy.json`
`ci:vllm`	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
`ci:all`	All standard extended tests (excludes `ci:atom_full`)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4766 --add-label <label>

Copilot

Pull request overview

Adds the DeepSeek‑V4 sparse MLA training backward path for gfx950/CDNA4, complementing the existing sparse prefill forward by introducing the full backward pipeline (delta precompute, dQ, dKV interm, inverted top‑k CSR, dKV gather, and optional sink gradient) with Gluon + Triton kernels, plus correctness tests and Gluon docs.

Changes:

Introduce public wrapper sparse_mla_bwd_dsv4(...) under aiter/ops/triton/attention/ that orchestrates the backward pipeline and chunking via R_CHUNK.
Add new Gluon kernels for dQ and dKV-intermediate computation, and Triton kernels for delta and CSR gather accumulation.
Add gfx950-only correctness tests and document the new Gluon module in aiter/ops/triton/gluon/README.md.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py	New gfx950-gated correctness tests vs autograd reference (including chunking + sink cases).
aiter/ops/triton/gluon/sparse_attention_dsv4_bwd_gluon.py	New Gluon kernels for dQ and dKV-intermediate for DSv4 sparse MLA backward.
aiter/ops/triton/gluon/README.md	Document the new DSv4 sparse MLA backward entry and performance notes.
aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py	New public API wrapper implementing the DSv4 sparse MLA backward pipeline (chunked/unchunked).
aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py	New Triton kernels/utilities for `delta`, inverted-topk CSR build, and atomic-free dKV gather accumulation.

Suppressed comments (3)

aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py:110

R_CHUNK can be set to a value that does not evenly divide TOPK (or even exceeds it). In that case, the Gluon kernels still iterate over R_CHUNK entries starting at R_START=r, which can read past the end of each topk_indices row and produce incorrect results / OOB reads. Please enforce that R_CHUNK is a positive divisor of TOPK (or pad topk_indices to a multiple).

    if scale is None:
        scale = 1.0 / (D**0.5)
    if R_CHUNK is None:
        R_CHUNK = TOPK
    lse = lse.float().contiguous()

aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py:104

The wrapper documents strict dtypes/shapes (bf16 Q/KV/dO/O, fp32 lse/sink, int32 indices) but currently only checks contiguity/shapes partially. Adding explicit dtype + shape validation (including lse.shape == (T, H) and attn_sink shape/dtype when provided) will prevent hard-to-debug miscompiles or silent correctness issues.

    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()

aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4_bwd.py:70

This new Triton kernel is also missing a config-aware repr=make_kernel_repr(...) (rule: "Kernel conventions"). Adding a repr here improves profiling/trace readability and aligns with other attention kernels that define repr objects.

@triton.jit
def _bwd_dkv_gather_acc_v4_be(
    Interm_ptr,  # [T, R_CHUNK, D] bf16, flat [T*R_CHUNK, D]

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

leonling-ll

Overall good to me, only some minor comments, please take a look.

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

op_tests/triton_tests/attention/test_sparse_attention_dsv4_bwd.py:72

_ref_fwd() claims to produce a sink-inclusive lse, but when attn_sink is not None it still sets m = s.max(dim=1).values instead of folding the sink into the max (m = max(max(s), attn_sink)). This is less numerically stable and diverges from the stated convention (sink folded into e_max/e_sum), and can overflow exp(attn_sink - m) if attn_sink is much larger than the score max. Compute m with torch.maximum(..., attn_sink) and use that m consistently for p, denom, and lse.

            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)

leonling-ll

LGTM

Copilot

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

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.

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).

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).

`_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).

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).

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py:756

In _bwd_dkv_gather_acc_v4, the loop for i0 in range(start, end, BLOCK_E): uses start/end loaded from device memory (tl.load). Python’s range() can’t take runtime tl.tensor values, so this will not compile as written. Use a Triton-supported dynamic loop (e.g., i0 = start; while i0 < end: ...; i0 += BLOCK_E).

    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

op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py:178

_time_phases() times the gather phase with dkv_gather_acc(interm, inv_ptr, inv_data, dkv_acc) using the default accumulate=True and a persistent dkv_acc. Because _bench() repeats the lambda, dkv_acc will keep accumulating across warmup/reps and the kernel will also include the extra read of the fp32 accumulator, so the reported gather time won’t match the public entry’s first-chunk accumulate=False behavior. Consider zeroing dkv_acc inside the lambda and/or passing accumulate=False for the breakdown case intended to mirror the end-to-end path.

    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)),
        (

brunomazzottiamd

Hi @wangye805. I have some comments regarding code organization and code duplication.

… 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).

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

brunomazzottiamd

LGTM!

brunomazzottiamd · 2026-08-24T11:30:00Z

I think we have a conflict in aiter/test_mha_common.py.

Only conflict is `aiter/test_mha_common.py`: main added `opus_ref_lse` / `opus_check_lse` (#4877) at the end of the file and this branch added `sparse_mla_dsv4_ref` there. Unrelated additions, both kept.

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).

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

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.

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py:340

attn_sink is only shape-validated, but unlike the other tensors it is not checked for is_cuda, dtype, or contiguity. If a caller passes a CPU tensor (or a non-fp32 sink), compute_d_sink() will error with a device/dtype mismatch even though plan_bwd() accepted the inputs. Add the same device/dtype/contiguity validation for attn_sink as for the other inputs (or explicitly .to(device, dtype) it).

    if attn_sink is not None:
        assert attn_sink.shape == (
            H,
        ), f"attn_sink must be [{H}], got {tuple(attn_sink.shape)}"
    assert (

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.

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py:297

plan_bwd/bwd_phases appear to be internal helpers (used by the benchmark), but their names are public-looking and the module is a public wrapper module. To avoid accidentally expanding the supported API surface, consider making them explicitly private (e.g., _plan_bwd / _bwd_phases) or moving them under a clearly-internal namespace, while keeping sparse_mla_bwd_dsv4 as the only supported entry point.

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.
    """

aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py:340

plan_bwd validates device/dtype/contiguity for the required tensors, but attn_sink (when provided) is only shape-checked. Since compute_d_sink() later performs GPU math with attn_sink and lse, passing a CPU (or wrong-dtype) attn_sink will fail late with a less actionable device/dtype error. Add an explicit validation/conversion for attn_sink (e.g., require .is_cuda and dtype==torch.float32, and/or normalize with .contiguous().to(torch.float32) similarly to other attention wrappers).

    if attn_sink is not None:
        assert attn_sink.shape == (
            H,
        ), f"attn_sink must be [{H}], got {tuple(attn_sink.shape)}"
    assert (

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.

Copilot

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py:47

In _dq_v4_kernel the parameter comment for KV_ptr says # [T, D] bf16, but the backward wrapper explicitly supports kv as [num_kv, 512] with num_kv >= T (compressed pool rows). This mismatch can mislead future maintainers about the valid range of top-k indices and the expected KV layout; please update the comment to [num_kv, D] (or similar) to match the real contract.

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

wangye805 requested review from a team and a lite review from Copilot August 14, 2026 19:44

Copilot started reviewing on behalf of wangye805 August 14, 2026 19:44 View session

Copilot AI reviewed Aug 14, 2026

View reviewed changes

wangye805 marked this pull request as draft August 14, 2026 19:52

wangye805 changed the title ~~[TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward~~ [WIP][TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward Aug 14, 2026

wangye805 changed the title ~~[WIP][TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward~~ [TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward Aug 18, 2026

wangye805 marked this pull request as ready for review August 18, 2026 17:29

github-actions Bot changed the title ~~[TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward~~ [Triton/Gluon] [GFX950][DSV4] Sparse MLA training backward Aug 18, 2026

wangye805 force-pushed the yewang12/dsv4_train_bwd_gluon branch from ba2b7e8 to bc95c96 Compare August 18, 2026 19:00

wangye805 requested review from Dewei-Wang-sh and leonling-ll August 18, 2026 19:04

leonling-ll reviewed Aug 19, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings August 21, 2026 04:27

Copilot started reviewing on behalf of wangye805 August 21, 2026 04:28 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py

wangye805 requested a review from leonling-ll August 21, 2026 04:37

leonling-ll previously approved these changes Aug 21, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings August 21, 2026 14:38

wangye805 dismissed leonling-ll’s stale review via faf960f August 21, 2026 14:38

Copilot started reviewing on behalf of wangye805 August 21, 2026 14:39 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes

Comment thread aiter/ops/triton/_gluon_kernels/gfx950/attention/sparse_attention_dsv4_bwd.py Outdated

vgokhale previously approved these changes Aug 21, 2026

View reviewed changes

wangye805 force-pushed the yewang12/dsv4_train_bwd_gluon branch from faf960f to 5ef9614 Compare August 21, 2026 15:06

wangye805 added 5 commits August 21, 2026 10:08

Copilot started reviewing on behalf of wangye805 August 21, 2026 15:20 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes

Comment thread op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py Outdated

wangye805 requested a review from brunomazzottiamd August 21, 2026 15:26

vgokhale previously approved these changes Aug 21, 2026

View reviewed changes

brunomazzottiamd reviewed Aug 21, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings August 21, 2026 19:29

wangye805 dismissed vgokhale’s stale review via 7e004f1 August 21, 2026 19:29

Copilot started reviewing on behalf of wangye805 August 21, 2026 19:30 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py

Comment thread aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py Outdated

Comment thread op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py Outdated

brunomazzottiamd previously approved these changes Aug 24, 2026

View reviewed changes

wangye805 added 2 commits August 24, 2026 16:57

Merge main into the DSv4 backward branch

063f056

Only conflict is `aiter/test_mha_common.py`: main added `opus_ref_lse` / `opus_check_lse` (#4877) at the end of the file and this branch added `sparse_mla_dsv4_ref` there. Unrelated additions, both kept.

Copilot AI review requested due to automatic review settings August 25, 2026 00:54

wangye805 dismissed brunomazzottiamd’s stale review via 5abadaf August 25, 2026 00:54

Copilot started reviewing on behalf of wangye805 August 25, 2026 00:55 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py

Comment thread aiter/test_mha_common.py

Copilot AI review requested due to automatic review settings August 25, 2026 02:04

Copilot started reviewing on behalf of wangye805 August 25, 2026 02:05 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Comment thread op_tests/op_benchmarks/triton/bench_sparse_attention_dsv4_bwd.py

Copilot AI review requested due to automatic review settings August 25, 2026 02:17

Copilot started reviewing on behalf of wangye805 August 25, 2026 02:18 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings August 25, 2026 02:54

Copilot started reviewing on behalf of wangye805 August 25, 2026 02:56 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Comment thread aiter/ops/triton/attention/sparse_attention_dsv4_bwd.py

Uh oh!

Conversation

wangye805 commented Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Contract

Structure

Performance

Environment

Correctness

Please note

Uh oh!

github-actions Bot commented Aug 14, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

leonling-ll left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

leonling-ll left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

brunomazzottiamd left a comment

Choose a reason for hiding this comment

Uh oh!

brunomazzottiamd commented Aug 24, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

wangye805 commented Aug 14, 2026 •

edited

Loading