[WIP][TRITON][GLUON][GFX950][DSV4] Sparse MLA training backward

[wangye805]

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.

num_kv >= T is supported, so a compressed pool (kv rows T..num_kv-1) works.

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

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 / 407 TFLOPS	3.223 ms / 426 TFLOPS
end-to-end	3.238 ms / 424 TFLOPS	3.471 ms / 396 TFLOPS

Read that as parity — the two differ mainly in how much each allocates per call, which is why the ordering flips between the two rows.

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]

🏷️ 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.