From 91874c3deb9e5c4a2803d76bfaca4367a0211556 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sat, 3 Oct 2026 23:28:16 +0000
Subject: [PATCH 01/11] perf(kernels): use cuDNN CuTe-DSL DSA kernel for sparse
 MLA backward on SM90

The vendored TileLang sparse MLA backward reaches ~93-130 TFLOP/s on H200
(no pipelining, fp32 atomics for every top-k slot including sentinels).
On Hopper, dispatch to cuDNN frontend's SM90 DSA backward instead
(warp-specialized, TMA Q/dO, cp.async KV gather, skips dKV atomics for
masked slots): 3.5-4.7x faster, ~410-440 TFLOP/s at S=8k/16k, topk=2048.
Sentinel indices are remapped to -1 and the base-2 LSE is converted to
natural log. TileLang remains the fallback on other architectures.

Bumps nvidia-cudnn-frontend to >=1.30 (1.25 is incompatible with the
pinned CUTLASS DSL 4.7.1) and adds a standalone microbenchmark with an
fp32 reference check.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/scripts/bench_sparse_mla.py        | 170 ++++++++++++++++++
 pyproject.toml                                |   2 +
 .../trainer/models/kernels/sparse_mla_bwd.py  |  83 +++++++--
 uv.lock                                       |  41 +++--
 4 files changed, 266 insertions(+), 30 deletions(-)
 create mode 100644 benchmarks/scripts/bench_sparse_mla.py

diff --git a/benchmarks/scripts/bench_sparse_mla.py b/benchmarks/scripts/bench_sparse_mla.py
new file mode 100644
index 00000000000..f3185127cf0
--- /dev/null
+++ b/benchmarks/scripts/bench_sparse_mla.py
@@ -0,0 +1,170 @@
+"""Microbenchmark + correctness check for the DSA sparse MLA kernels (GLM-5 / DeepSeek-V3.2 shapes).
+
+Usage:
+    uv run python benchmarks/scripts/bench_sparse_mla.py
+    uv run python benchmarks/scripts/bench_sparse_mla.py --seqlens 8192 16384 --check-seqlen 4096
+"""
+
+import argparse
+import math
+
+import torch
+
+from prime_rl.trainer.models.kernels import sparse_mla_bwd
+from prime_rl.trainer.models.kernels.sparse_mla_fwd import sparse_mla
+
+H = 64
+D_QK = 576
+D_V = 512
+
+
+def make_indices(seqlen: int, topk: int, doc_len: int, device: str, generator: torch.Generator) -> torch.Tensor:
+    """Causal, document-masked top-k indices in the layout produced by the indexer.
+
+    Tokens with fewer than `topk` visible keys get every visible key plus the sentinel
+    (`seqlen`, i.e. the extra zero KV row) for the remaining slots; later tokens get a
+    random subset of their visible keys. Order inside a row is random.
+    """
+    sentinel = seqlen
+    positions = torch.arange(seqlen, device=device)
+    doc_start = (positions // doc_len) * doc_len
+    indices = torch.empty(seqlen, topk, dtype=torch.int32, device=device)
+    chunk = 512
+    for start in range(0, seqlen, chunk):
+        rows = positions[start : start + chunk]
+        scores = torch.rand(len(rows), seqlen, device=device, generator=generator)
+        visible = (positions[None, :] <= rows[:, None]) & (positions[None, :] >= doc_start[rows][:, None])
+        scores = torch.where(visible, scores, -1.0)
+        top_scores, top_idx = scores.topk(topk, dim=-1)
+        top_idx = torch.where(top_scores >= 0, top_idx, sentinel)
+        perm = torch.rand(len(rows), topk, device=device, generator=generator).argsort(dim=-1)
+        indices[start : start + chunk] = top_idx.gather(1, perm).int()
+    return indices.view(1, seqlen, 1, topk)
+
+
+def make_inputs(seqlen: int, topk: int, doc_len: int, seed: int = 0):
+    device = "cuda"
+    g = torch.Generator(device=device).manual_seed(seed)
+    q = torch.randn(1, seqlen, H, D_QK, device=device, dtype=torch.bfloat16, generator=g)
+    kv = torch.randn(1, seqlen + 1, 1, D_QK, device=device, dtype=torch.bfloat16, generator=g)
+    kv[:, -1] = 0
+    grad_out = torch.randn(1, seqlen, H, D_V, device=device, dtype=torch.bfloat16, generator=g)
+    indices = make_indices(seqlen, topk, doc_len, device, g)
+    return q, kv, grad_out, indices
+
+
+def reference_grads(q, kv, grad_out, indices, sm_scale, chunk: int = 128):
+    """fp32 autograd reference: gather KV by index, mask sentinel, softmax, matmuls."""
+    _, seqlen, _, _ = q.shape
+    sentinel = kv.shape[1] - 1
+    kv32 = kv[0, :, 0].float().requires_grad_(True)
+    dq = torch.empty(seqlen, H, D_QK, device=q.device, dtype=torch.float32)
+    for start in range(0, seqlen, chunk):
+        idx = indices[0, start : start + chunk, 0].long()
+        q32 = q[0, start : start + chunk].float().requires_grad_(True)
+        kv_g = kv32[idx]  # [c, topk, 576]
+        scores = torch.einsum("chd,ckd->chk", q32, kv_g) * sm_scale
+        scores = scores.masked_fill((idx == sentinel)[:, None, :], float("-inf"))
+        probs = scores.softmax(dim=-1)
+        out = torch.einsum("chk,ckd->chd", probs, kv_g[..., :D_V])
+        out.backward(grad_out[0, start : start + chunk].float())
+        dq[start : start + chunk] = q32.grad
+    dkv = kv32.grad.clone()
+    dkv[sentinel] = 0
+    return dq.unsqueeze(0), dkv.view(1, -1, 1, D_QK)
+
+
+def compare(name: str, got: torch.Tensor, ref: torch.Tensor) -> None:
+    got, ref = got.float(), ref.float()
+    abs_err = (got - ref).abs().max().item()
+    rel_err = ((got - ref).norm() / ref.norm()).item()
+    print(f"  {name}: max_abs={abs_err:.3e} ref_max={ref.abs().max().item():.3e} rel_l2={rel_err:.3e}")
+    assert rel_err < 1e-2, f"{name} rel_l2 {rel_err} too large"
+
+
+def time_ms(fn, warmup: int = 3, iters: int = 10) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def backends():
+    return {"tilelang": sparse_mla_bwd.tilelang_sparse_mla_backward, "default": sparse_mla_bwd.sparse_mla_backward}
+
+
+def print_device() -> None:
+    print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())
+
+
+def check(seqlen: int, topk: int, doc_len: int) -> None:
+    sm_scale = D_QK**-0.5
+    q, kv, grad_out, indices = make_inputs(seqlen, topk, doc_len, seed=1)
+    out, lse = sparse_mla(q, kv, indices, sm_scale)
+    dq_ref, dkv_ref = reference_grads(q, kv, grad_out, indices, sm_scale)
+    for name, fn in backends().items():
+        print(f"[check S={seqlen} topk={topk} doc_len={doc_len}] {name}")
+        dq, dkv = fn(q, kv, out, grad_out, indices, lse, sm_scale)
+        compare("dQ", dq, dq_ref)
+        compare("dKV", dkv, dkv_ref)
+
+
+def check_batched_cp_shard(seqlen: int, topk: int, doc_len: int) -> None:
+    """B=2 with a CP-style query shard (second half of the queries, full KV): every backend vs tilelang."""
+    sm_scale = D_QK**-0.5
+    inputs = [make_inputs(seqlen, topk, doc_len, seed=seed) for seed in (2, 3)]
+    half = seqlen // 2
+    q, kv, grad_out, indices = (torch.cat(ts).contiguous() for ts in zip(*inputs))
+    q, grad_out, indices = (t[:, half:].contiguous() for t in (q, grad_out, indices))
+    out, lse = sparse_mla(q, kv, indices, sm_scale)
+    dq_ref, dkv_ref = sparse_mla_bwd.tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+    for name, fn in backends().items():
+        print(f"[check B=2 S_q={seqlen - half} S_kv={seqlen + 1}] {name} vs tilelang")
+        dq, dkv = fn(q, kv, out, grad_out, indices, lse, sm_scale)
+        compare("dQ", dq, dq_ref)
+        compare("dKV", dkv, dkv_ref)
+
+
+def bench(seqlen: int, topk: int, doc_len: int) -> None:
+    sm_scale = D_QK**-0.5
+    q, kv, grad_out, indices = make_inputs(seqlen, topk, doc_len)
+    fwd_flops = 2 * seqlen * H * topk * (D_QK + D_V)
+    bwd_flops = 2.5 * fwd_flops
+
+    fwd_ms = time_ms(lambda: sparse_mla(q, kv, indices, sm_scale))
+    print(
+        f"[bench S={seqlen} topk={topk} doc_len={doc_len}] fwd: {fwd_ms:.3f} ms {fwd_flops / fwd_ms / 1e9:.1f} TFLOP/s"
+    )
+    out, lse = sparse_mla(q, kv, indices, sm_scale)
+    for name, fn in backends().items():
+        bwd_ms = time_ms(lambda: fn(q, kv, out, grad_out, indices, lse, sm_scale))
+        print(f"  bwd {name}: {bwd_ms:.3f} ms {bwd_flops / bwd_ms / 1e9:.1f} TFLOP/s")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seqlens", type=int, nargs="+", default=[8192, 16384])
+    parser.add_argument("--topk", type=int, default=2048)
+    parser.add_argument("--doc-len", type=int, default=None, help="packed document length (default: one document)")
+    parser.add_argument("--check-seqlen", type=int, default=4096)
+    parser.add_argument("--check-doc-len", type=int, default=1536)
+    args = parser.parse_args()
+    print_device()
+
+    if args.check_seqlen > 0:
+        check(args.check_seqlen, args.topk, args.check_doc_len)
+        check_batched_cp_shard(args.check_seqlen, args.topk, args.check_doc_len)
+    for seqlen in args.seqlens:
+        bench(seqlen, args.topk, args.doc_len or seqlen)
+        if args.doc_len is None:
+            bench(seqlen, args.topk, math.gcd(seqlen, 4096))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index fdb2c2dbc6f..f7773daa081 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,6 +59,8 @@ gpu = [
     "torchtitan ; sys_platform == 'linux'",
     "dion ; sys_platform == 'linux'",
     "tilelang>=0.1.8 ; sys_platform == 'linux'",
+    # DSA sparse MLA backward kernel (CuTe DSL); >=1.29 matches the pinned CUTLASS DSL 4.7.
+    "nvidia-cudnn-frontend>=1.30.0 ; sys_platform == 'linux'",
     "flash-linear-attention ; sys_platform == 'linux'",
     "mamba-ssm>=2.3.0 ; sys_platform == 'linux'",
     "nvidia-ml-py>=12.575.51 ; sys_platform == 'linux'",
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
index 93a6ffd7f94..58317d5e465 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -13,6 +13,8 @@
     # This is expected on CPU-only machines
     pass
 
+import math
+
 import tilelang
 import torch
 from tilelang import language as T
@@ -276,8 +278,7 @@ def sparse_mla_bwd_kernel(
     return sparse_mla_bwd_kernel
 
 
-@torch.library.custom_op("prime_rl::sparse_mla_backward", mutates_args=())
-def sparse_mla_backward(
+def tilelang_sparse_mla_backward(
     q: torch.Tensor,
     kv: torch.Tensor,
     out: torch.Tensor,
@@ -286,20 +287,12 @@ def sparse_mla_backward(
     lse: torch.Tensor,
     sm_scale: float | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    assert q.is_contiguous()
-    assert kv.is_contiguous()
-    assert indices.is_contiguous()
-    assert lse.is_contiguous()
     grad_out = grad_out.contiguous()
-    B, S, H, dim_plus_tail_dim = q.shape
-    _, S_kv, kv_group, _ = kv.shape
-    assert kv.shape[-1] == dim_plus_tail_dim
-    assert kv.shape[0] == B
+    H, dim_plus_tail_dim = q.shape[2:]
+    kv_group = kv.shape[2]
     D = 512
     D_tail = dim_plus_tail_dim - D
     topk = indices.shape[-1]
-    assert indices.shape == (B, S, kv_group, topk)
-    assert lse.shape == (B, S, H)
 
     preprocess_kernel = preprocess(H, D)
     bwd_kernel = bwd(H, D, D_tail, topk, kv_group, sm_scale, True)
@@ -309,10 +302,74 @@ def sparse_mla_backward(
     dkv = torch.zeros_like(kv, dtype=torch.float32)
     dq = bwd_kernel(q, kv, grad_out, indices, lse, delta, dkv)
     dkv = postprocess_kernel(dkv)
-
     return dq, dkv
 
 
+def cudnn_sparse_mla_backward(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    out: torch.Tensor,
+    grad_out: torch.Tensor,
+    indices: torch.Tensor,
+    lse: torch.Tensor,
+    sm_scale: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """cuDNN-frontend CuTe-DSL DSA backward (SM90, single KV head with K=V).
+
+    The cuDNN kernel takes flat (tokens, ...) tensors with global indices, skips
+    negative indices (no KV load, no dKV atomics) and expects a natural-log LSE.
+    """
+    from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90
+
+    B, S, H, D_qk = q.shape
+    S_kv = kv.shape[1]
+    topk = indices.shape[-1]
+    # Sentinel / out-of-range indices (see sparse_mla_fwd) become -1; valid ones are
+    # offset into the batch-flattened KV.
+    batch_offset = torch.arange(B, device=indices.device, dtype=indices.dtype).view(B, 1, 1, 1) * S_kv
+    flat_indices = torch.where(indices <= S_kv - 2, indices + batch_offset, -1).view(B * S, topk)
+    # The forward kernel returns a base-2 LSE.
+    lse_ln = lse * math.log(2.0)
+
+    dq, dkv = flash_attn_bwd_sm90(
+        q.view(B * S, H, D_qk),
+        kv.view(B * S_kv, D_qk),
+        out.view(B * S, H, -1),
+        grad_out.contiguous().view(B * S, H, -1),
+        lse_ln.view(B * S, H),
+        softmax_scale=sm_scale,
+        topk_idxs=flat_indices,
+    )
+    return dq.view_as(q), dkv.view_as(kv)
+
+
+@torch.library.custom_op("prime_rl::sparse_mla_backward", mutates_args=())
+def sparse_mla_backward(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    out: torch.Tensor,
+    grad_out: torch.Tensor,
+    indices: torch.Tensor,
+    lse: torch.Tensor,
+    sm_scale: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert q.is_contiguous()
+    assert kv.is_contiguous()
+    assert indices.is_contiguous()
+    assert lse.is_contiguous()
+    B, S, H, dim_plus_tail_dim = q.shape
+    _, S_kv, kv_group, _ = kv.shape
+    assert kv.shape[-1] == dim_plus_tail_dim
+    assert kv.shape[0] == B
+    topk = indices.shape[-1]
+    assert indices.shape == (B, S, kv_group, topk)
+    assert lse.shape == (B, S, H)
+
+    if torch.cuda.get_device_capability(q.device)[0] == 9 and kv_group == 1:
+        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+    return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+
+
 @sparse_mla_backward.register_fake
 def _sparse_mla_backward_fake(
     q: torch.Tensor,
diff --git a/uv.lock b/uv.lock
index 0ccfe2a481c..d1897a0ed44 100644
--- a/uv.lock
+++ b/uv.lock
@@ -17,26 +17,26 @@ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for
 exclude-newer-span = "P7D"
 
 [options.exclude-newer-package]
-verifiers = false
-prime-pydantic-config = false
-harbor = "2026-08-11T00:00:00Z"
-prime-traces = false
-vllm-router = false
+deep-ep = false
+deep-gemm = false
 dion = false
 fastokens = false
-prime-runs = false
 flash-attn-3 = false
-prime-tunnel = false
-deep-gemm = false
-prime-evals = false
-torchao = false
-tokenspeed-mla = false
+harbor = "2026-08-11T00:00:00Z"
 nixl-cu13 = false
-deep-ep = false
+prime = false
+prime-evals = false
+prime-kernels = false
+prime-pydantic-config = false
+prime-runs = false
 prime-sandboxes = false
+prime-traces = false
+prime-tunnel = false
 renderers = false
-prime-kernels = false
-prime = false
+tokenspeed-mla = false
+torchao = false
+verifiers = false
+vllm-router = false
 
 [manifest]
 members = [
@@ -4045,11 +4045,15 @@ wheels = [
 
 [[package]]
 name = "nvidia-cudnn-frontend"
-version = "1.25.0"
+version = "1.30.0"
 source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "apache-tvm-ffi" },
+    { name = "nvidia-cutlass-dsl" },
+]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/28/0f/df39a194f2529093db737d43cc4cbf594c6a79712a09aa104b999e4d95d4/nvidia_cudnn_frontend-1.25.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:09e6e1bc48ce1235743f89d8ea699c52b3008fd6dae7f2ecadb744bebf272a2b", size = 3263306, upload-time = "2026-06-10T21:07:48.093Z" },
-    { url = "https://files.pythonhosted.org/packages/03/65/3b45941d8a22128b971e910f2e9af6bf5ef453e92cc329c56b6eb53c53de/nvidia_cudnn_frontend-1.25.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9a94a72d736bd79eb35f451aaf26d9493778e02ecabccc92c05425508c9e7a83", size = 3414884, upload-time = "2026-06-10T21:08:08.603Z" },
+    { url = "https://files.pythonhosted.org/packages/e4/53/07836d9f74f1416a8768530d428c155f67d7728f283bce3a9b3096b7e18f/nvidia_cudnn_frontend-1.30.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:279d227a226429288ae9e4e88db744ce144ae69605f63a286dc84ae180869a6e", size = 8631629, upload-time = "2026-09-23T03:44:15.265Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/95/c92076af7561e6e43da1ee3ac4155202dcd3876e5b999ff0dea5754eb3d2/nvidia_cudnn_frontend-1.30.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5dda72db38ea0c39f158ebffa9596eefab21ce7cc6a12824490ab8b6cb60e76d", size = 8797507, upload-time = "2026-09-23T03:44:35.051Z" },
 ]
 
 [[package]]
@@ -4988,6 +4992,7 @@ all = [
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux'" },
     { name = "nixl", marker = "sys_platform == 'linux'" },
     { name = "nixl-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux'" },
     { name = "nvidia-ml-py", marker = "sys_platform == 'linux'" },
     { name = "quack-kernels", marker = "sys_platform == 'linux'" },
     { name = "ring-flash-attn", marker = "sys_platform == 'linux'" },
@@ -5039,6 +5044,7 @@ gpu = [
     { name = "modelexpress", marker = "sys_platform == 'linux'" },
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux'" },
     { name = "nixl-cu13", marker = "sys_platform == 'linux'" },
+    { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux'" },
     { name = "nvidia-ml-py", marker = "sys_platform == 'linux'" },
     { name = "ring-flash-attn", marker = "sys_platform == 'linux'" },
     { name = "tilelang" },
@@ -5103,6 +5109,7 @@ requires-dist = [
     { name = "nixl", marker = "sys_platform == 'linux' and extra == 'disagg'" },
     { name = "nixl-cu13", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==0.10.1" },
     { name = "numpy", specifier = ">=2.2.6" },
+    { name = "nvidia-cudnn-frontend", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=1.30.0" },
     { name = "nvidia-ml-py", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = ">=12.575.51" },
     { name = "openai", specifier = ">=1.106.1" },
     { name = "orjson", specifier = ">=3.11.0" },

From 30cd8c9ed71317b700483795a4b8ddd8778d9010 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 00:29:53 +0000
Subject: [PATCH 02/11] chore: drop sparse MLA microbenchmark script

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 benchmarks/scripts/bench_sparse_mla.py | 170 -------------------------
 1 file changed, 170 deletions(-)
 delete mode 100644 benchmarks/scripts/bench_sparse_mla.py

diff --git a/benchmarks/scripts/bench_sparse_mla.py b/benchmarks/scripts/bench_sparse_mla.py
deleted file mode 100644
index f3185127cf0..00000000000
--- a/benchmarks/scripts/bench_sparse_mla.py
+++ /dev/null
@@ -1,170 +0,0 @@
-"""Microbenchmark + correctness check for the DSA sparse MLA kernels (GLM-5 / DeepSeek-V3.2 shapes).
-
-Usage:
-    uv run python benchmarks/scripts/bench_sparse_mla.py
-    uv run python benchmarks/scripts/bench_sparse_mla.py --seqlens 8192 16384 --check-seqlen 4096
-"""
-
-import argparse
-import math
-
-import torch
-
-from prime_rl.trainer.models.kernels import sparse_mla_bwd
-from prime_rl.trainer.models.kernels.sparse_mla_fwd import sparse_mla
-
-H = 64
-D_QK = 576
-D_V = 512
-
-
-def make_indices(seqlen: int, topk: int, doc_len: int, device: str, generator: torch.Generator) -> torch.Tensor:
-    """Causal, document-masked top-k indices in the layout produced by the indexer.
-
-    Tokens with fewer than `topk` visible keys get every visible key plus the sentinel
-    (`seqlen`, i.e. the extra zero KV row) for the remaining slots; later tokens get a
-    random subset of their visible keys. Order inside a row is random.
-    """
-    sentinel = seqlen
-    positions = torch.arange(seqlen, device=device)
-    doc_start = (positions // doc_len) * doc_len
-    indices = torch.empty(seqlen, topk, dtype=torch.int32, device=device)
-    chunk = 512
-    for start in range(0, seqlen, chunk):
-        rows = positions[start : start + chunk]
-        scores = torch.rand(len(rows), seqlen, device=device, generator=generator)
-        visible = (positions[None, :] <= rows[:, None]) & (positions[None, :] >= doc_start[rows][:, None])
-        scores = torch.where(visible, scores, -1.0)
-        top_scores, top_idx = scores.topk(topk, dim=-1)
-        top_idx = torch.where(top_scores >= 0, top_idx, sentinel)
-        perm = torch.rand(len(rows), topk, device=device, generator=generator).argsort(dim=-1)
-        indices[start : start + chunk] = top_idx.gather(1, perm).int()
-    return indices.view(1, seqlen, 1, topk)
-
-
-def make_inputs(seqlen: int, topk: int, doc_len: int, seed: int = 0):
-    device = "cuda"
-    g = torch.Generator(device=device).manual_seed(seed)
-    q = torch.randn(1, seqlen, H, D_QK, device=device, dtype=torch.bfloat16, generator=g)
-    kv = torch.randn(1, seqlen + 1, 1, D_QK, device=device, dtype=torch.bfloat16, generator=g)
-    kv[:, -1] = 0
-    grad_out = torch.randn(1, seqlen, H, D_V, device=device, dtype=torch.bfloat16, generator=g)
-    indices = make_indices(seqlen, topk, doc_len, device, g)
-    return q, kv, grad_out, indices
-
-
-def reference_grads(q, kv, grad_out, indices, sm_scale, chunk: int = 128):
-    """fp32 autograd reference: gather KV by index, mask sentinel, softmax, matmuls."""
-    _, seqlen, _, _ = q.shape
-    sentinel = kv.shape[1] - 1
-    kv32 = kv[0, :, 0].float().requires_grad_(True)
-    dq = torch.empty(seqlen, H, D_QK, device=q.device, dtype=torch.float32)
-    for start in range(0, seqlen, chunk):
-        idx = indices[0, start : start + chunk, 0].long()
-        q32 = q[0, start : start + chunk].float().requires_grad_(True)
-        kv_g = kv32[idx]  # [c, topk, 576]
-        scores = torch.einsum("chd,ckd->chk", q32, kv_g) * sm_scale
-        scores = scores.masked_fill((idx == sentinel)[:, None, :], float("-inf"))
-        probs = scores.softmax(dim=-1)
-        out = torch.einsum("chk,ckd->chd", probs, kv_g[..., :D_V])
-        out.backward(grad_out[0, start : start + chunk].float())
-        dq[start : start + chunk] = q32.grad
-    dkv = kv32.grad.clone()
-    dkv[sentinel] = 0
-    return dq.unsqueeze(0), dkv.view(1, -1, 1, D_QK)
-
-
-def compare(name: str, got: torch.Tensor, ref: torch.Tensor) -> None:
-    got, ref = got.float(), ref.float()
-    abs_err = (got - ref).abs().max().item()
-    rel_err = ((got - ref).norm() / ref.norm()).item()
-    print(f"  {name}: max_abs={abs_err:.3e} ref_max={ref.abs().max().item():.3e} rel_l2={rel_err:.3e}")
-    assert rel_err < 1e-2, f"{name} rel_l2 {rel_err} too large"
-
-
-def time_ms(fn, warmup: int = 3, iters: int = 10) -> float:
-    for _ in range(warmup):
-        fn()
-    torch.cuda.synchronize()
-    start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
-    start.record()
-    for _ in range(iters):
-        fn()
-    end.record()
-    torch.cuda.synchronize()
-    return start.elapsed_time(end) / iters
-
-
-def backends():
-    return {"tilelang": sparse_mla_bwd.tilelang_sparse_mla_backward, "default": sparse_mla_bwd.sparse_mla_backward}
-
-
-def print_device() -> None:
-    print(torch.cuda.get_device_name(), torch.cuda.get_device_capability())
-
-
-def check(seqlen: int, topk: int, doc_len: int) -> None:
-    sm_scale = D_QK**-0.5
-    q, kv, grad_out, indices = make_inputs(seqlen, topk, doc_len, seed=1)
-    out, lse = sparse_mla(q, kv, indices, sm_scale)
-    dq_ref, dkv_ref = reference_grads(q, kv, grad_out, indices, sm_scale)
-    for name, fn in backends().items():
-        print(f"[check S={seqlen} topk={topk} doc_len={doc_len}] {name}")
-        dq, dkv = fn(q, kv, out, grad_out, indices, lse, sm_scale)
-        compare("dQ", dq, dq_ref)
-        compare("dKV", dkv, dkv_ref)
-
-
-def check_batched_cp_shard(seqlen: int, topk: int, doc_len: int) -> None:
-    """B=2 with a CP-style query shard (second half of the queries, full KV): every backend vs tilelang."""
-    sm_scale = D_QK**-0.5
-    inputs = [make_inputs(seqlen, topk, doc_len, seed=seed) for seed in (2, 3)]
-    half = seqlen // 2
-    q, kv, grad_out, indices = (torch.cat(ts).contiguous() for ts in zip(*inputs))
-    q, grad_out, indices = (t[:, half:].contiguous() for t in (q, grad_out, indices))
-    out, lse = sparse_mla(q, kv, indices, sm_scale)
-    dq_ref, dkv_ref = sparse_mla_bwd.tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
-    for name, fn in backends().items():
-        print(f"[check B=2 S_q={seqlen - half} S_kv={seqlen + 1}] {name} vs tilelang")
-        dq, dkv = fn(q, kv, out, grad_out, indices, lse, sm_scale)
-        compare("dQ", dq, dq_ref)
-        compare("dKV", dkv, dkv_ref)
-
-
-def bench(seqlen: int, topk: int, doc_len: int) -> None:
-    sm_scale = D_QK**-0.5
-    q, kv, grad_out, indices = make_inputs(seqlen, topk, doc_len)
-    fwd_flops = 2 * seqlen * H * topk * (D_QK + D_V)
-    bwd_flops = 2.5 * fwd_flops
-
-    fwd_ms = time_ms(lambda: sparse_mla(q, kv, indices, sm_scale))
-    print(
-        f"[bench S={seqlen} topk={topk} doc_len={doc_len}] fwd: {fwd_ms:.3f} ms {fwd_flops / fwd_ms / 1e9:.1f} TFLOP/s"
-    )
-    out, lse = sparse_mla(q, kv, indices, sm_scale)
-    for name, fn in backends().items():
-        bwd_ms = time_ms(lambda: fn(q, kv, out, grad_out, indices, lse, sm_scale))
-        print(f"  bwd {name}: {bwd_ms:.3f} ms {bwd_flops / bwd_ms / 1e9:.1f} TFLOP/s")
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--seqlens", type=int, nargs="+", default=[8192, 16384])
-    parser.add_argument("--topk", type=int, default=2048)
-    parser.add_argument("--doc-len", type=int, default=None, help="packed document length (default: one document)")
-    parser.add_argument("--check-seqlen", type=int, default=4096)
-    parser.add_argument("--check-doc-len", type=int, default=1536)
-    args = parser.parse_args()
-    print_device()
-
-    if args.check_seqlen > 0:
-        check(args.check_seqlen, args.topk, args.check_doc_len)
-        check_batched_cp_shard(args.check_seqlen, args.topk, args.check_doc_len)
-    for seqlen in args.seqlens:
-        bench(seqlen, args.topk, args.doc_len or seqlen)
-        if args.doc_len is None:
-            bench(seqlen, args.topk, math.gcd(seqlen, 4096))
-
-
-if __name__ == "__main__":
-    main()

From 528f731ed0fbf04c742e1e766df465ff64bc6d72 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 01:04:20 +0000
Subject: [PATCH 03/11] feat(config): model.sparse_mla_backward to choose the
 DSA backward kernel

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../prime-rl-configs/src/prime_rl/configs/trainer.py     | 3 +++
 src/prime_rl/trainer/model.py                            | 1 +
 .../trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py   | 1 +
 .../trainer/models/glm_moe_dsa/sparse_mla_attention.py   | 3 ++-
 src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py    | 9 ++++++++-
 src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py    | 8 ++++++--
 6 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index f31d24f6397..64f01ac3a65 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -325,6 +325,9 @@ class ModelConfig(BaseModelConfig):
 
     quantization: QuantizationConfig | None = None
 
+    sparse_mla_backward: Literal["auto", "cudnn", "tilelang"] = "auto"
+    """Kernel for the DSA sparse MLA attention backward (``glm_moe_dsa``, e.g. GLM-5). ``cudnn`` uses the cuDNN frontend CuTe-DSL DSA backward (Hopper/SM90, single KV head only). ``tilelang`` uses the vendored TileLang kernel. ``auto`` picks ``cudnn`` where supported and ``tilelang`` otherwise."""
+
     index_cache: IndexCacheConfig | None = None
     """DSA IndexCache sub-configuration. If set, sparse-attention top-k indices are reused across decoder layers per the configured schedule (mirrors vLLM's IndexCache HF overrides). If None, every layer recomputes its own indices."""
 
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index b78e12936c1..6fe768b18b9 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -365,6 +365,7 @@ def get_model(
         subconfig = getattr(model_config, subconfig_key, None)
         if subconfig is not None and hasattr(subconfig, "use_cache"):
             subconfig.use_cache = False
+    model_config.sparse_mla_backward = config.sparse_mla_backward
     if config.index_cache is not None:
         model_config.use_index_cache = True
         model_config.index_topk_freq = config.index_cache.topk_freq
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
index 1278f73b369..91d76079f38 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
@@ -41,6 +41,7 @@ def _sparse_mla_attention_args(config: GlmMoeDsaConfig, layer_idx: int) -> Spars
         index_topk=config.index_topk,
         use_index_cache=getattr(config, "use_index_cache", False),
         skip_topk=_index_cache_skip_topk(config, layer_idx),
+        sparse_mla_backward=getattr(config, "sparse_mla_backward", "auto"),
     )
 
 
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 02673f6d3f0..567568d5856 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -32,6 +32,7 @@ class SparseMlaAttentionArgs:
     index_topk: int
     use_index_cache: bool = False
     skip_topk: bool = False
+    sparse_mla_backward: str = "auto"
 
 
 def apply_rope_interleave_single(
@@ -230,7 +231,7 @@ def forward(
             position_embeddings_full=position_embeddings,
         )
 
-        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling)
+        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling, backward_backend=self.args.sparse_mla_backward)
         out = torch.einsum("bshk,hdk->bshd", out, w_v)
         batch_size, total_tokens = out.shape[:2]
         out = out.reshape(batch_size, total_tokens, -1)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
index 58317d5e465..1698cd5433f 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -352,6 +352,7 @@ def sparse_mla_backward(
     indices: torch.Tensor,
     lse: torch.Tensor,
     sm_scale: float | None = None,
+    backend: str = "auto",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert q.is_contiguous()
     assert kv.is_contiguous()
@@ -365,8 +366,13 @@ def sparse_mla_backward(
     assert indices.shape == (B, S, kv_group, topk)
     assert lse.shape == (B, S, H)
 
-    if torch.cuda.get_device_capability(q.device)[0] == 9 and kv_group == 1:
+    cudnn_supported = torch.cuda.get_device_capability(q.device)[0] == 9 and kv_group == 1
+    if backend == "auto":
+        backend = "cudnn" if cudnn_supported else "tilelang"
+    if backend == "cudnn":
+        assert cudnn_supported, "cuDNN sparse MLA backward requires an SM90 GPU and a single KV head"
         return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+    assert backend == "tilelang", f"Unknown sparse MLA backward backend: {backend}"
     return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
 
 
@@ -379,5 +385,6 @@ def _sparse_mla_backward_fake(
     indices: torch.Tensor,
     lse: torch.Tensor,
     sm_scale: float | None = None,
+    backend: str = "auto",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), torch.empty_like(kv)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
index aace98c36e9..c0cd428e4f3 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -196,6 +196,7 @@ def sparse_mla(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
+    backward_backend: str = "auto",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
     batch, seq_len, heads, dim_plus_tail_dim = q.shape
@@ -235,15 +236,17 @@ def _sparse_mla_fake(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
+    backward_backend: str = "auto",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return q.new_empty((*q.shape[:-1], d_v)), q.new_empty(q.shape[:-1], dtype=torch.float32)
 
 
 def _sparse_mla_setup_context(ctx, inputs, output) -> None:
-    q, kv, indices, sm_scale, _d_v, _block_I, _num_stages, _threads = inputs
+    q, kv, indices, sm_scale, _d_v, _block_I, _num_stages, _threads, backward_backend = inputs
     out, lse = output
     ctx.save_for_backward(q, kv, out, indices, lse)
     ctx.sm_scale = sm_scale
+    ctx.backward_backend = backward_backend
     ctx.mark_non_differentiable(lse)
 
 
@@ -257,8 +260,9 @@ def _sparse_mla_autograd_backward(ctx, grad_out: torch.Tensor, _grad_lse: torch.
         indices,
         lse.detach(),
         ctx.sm_scale,
+        ctx.backward_backend,
     )
-    return dq, dkv, None, None, None, None, None, None
+    return dq, dkv, None, None, None, None, None, None, None
 
 
 sparse_mla.register_autograd(_sparse_mla_autograd_backward, setup_context=_sparse_mla_setup_context)

From f4f1076c8d944b204131ab7c6ac5b5b60cb5a577 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 05:51:32 +0000
Subject: [PATCH 04/11] feat(trainer): model.dsa_backend with a FlashMLA
 forward + cuDNN backward route

Replace model.sparse_mla_backward with model.dsa_backend, selecting the DSA
sparse MLA kernels for both passes through the single prime_rl::sparse_mla op:

- tilelang (default): TileLang forward, cuDNN backward on SM90, TileLang
  backward elsewhere.
- cudnn_flashmla: FlashMLA flash_mla_sparse_fwd forward + cuDNN backward
  (SM90; the SM100/SM103 cuDNN backward is wired up but untested).
- tilelang_legacy: TileLang forward and backward.

The op now always returns a natural-log LSE. The FlashMLA route flattens the
batch, maps sentinel indices to -1 and passes a topk_length covering every
valid index so the sentinel tail is skipped.

FlashMLA (nv_dev @ b7643bd, sm_90a + sm_100f) ships as the new flash-mla
extra (part of `all`) from a prebuilt torch 2.13 / cu130 x86_64 wheel.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/prime_rl/configs/trainer.py           |  4 +-
 pyproject.toml                                |  7 ++
 src/prime_rl/trainer/model.py                 |  2 +-
 .../glm_moe_dsa/modeling_glm_moe_dsa.py       |  2 +-
 .../glm_moe_dsa/sparse_mla_attention.py       |  4 +-
 .../trainer/models/kernels/sparse_mla_bwd.py  | 86 +++++++++++++------
 .../trainer/models/kernels/sparse_mla_fwd.py  | 66 ++++++++++++--
 uv.lock                                       | 17 +++-
 8 files changed, 146 insertions(+), 42 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 64f01ac3a65..024aa79a592 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -325,8 +325,8 @@ class ModelConfig(BaseModelConfig):
 
     quantization: QuantizationConfig | None = None
 
-    sparse_mla_backward: Literal["auto", "cudnn", "tilelang"] = "auto"
-    """Kernel for the DSA sparse MLA attention backward (``glm_moe_dsa``, e.g. GLM-5). ``cudnn`` uses the cuDNN frontend CuTe-DSL DSA backward (Hopper/SM90, single KV head only). ``tilelang`` uses the vendored TileLang kernel. ``auto`` picks ``cudnn`` where supported and ``tilelang`` otherwise."""
+    dsa_backend: Literal["tilelang", "cudnn_flashmla", "tilelang_legacy"] = "tilelang"
+    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward with the cuDNN frontend CuTe-DSL backward on Hopper (SM90) and the TileLang backward elsewhere. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested). ``tilelang_legacy`` runs TileLang for both."""
 
     index_cache: IndexCacheConfig | None = None
     """DSA IndexCache sub-configuration. If set, sparse-attention top-k indices are reused across decoder layers per the configured schedule (mirrors vLLM's IndexCache HF overrides). If None, every layer recomputes its own indices."""
diff --git a/pyproject.toml b/pyproject.toml
index f7773daa081..2280978c199 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -96,6 +96,9 @@ quack = [
     "quack-kernels>=0.4.1 ; sys_platform == 'linux'",
 ]
 sonic-moe = ["sonic-moe ; sys_platform == 'linux'"]
+# FlashMLA sparse prefill forward for `model.dsa_backend = "cudnn_flashmla"` (sm_90a + sm_100f).
+# FlashMLA is not on PyPI and only an x86_64 wheel is published.
+flash-mla = ["flash-mla ; sys_platform == 'linux' and platform_machine == 'x86_64'"]
 # Prebuilt `prime-kernels` wheels from a release (see the `kernels` skill). Compiled for
 # sm_100a; on other GPUs the package installs fine and reports its kernels unavailable.
 kernels = [
@@ -109,6 +112,7 @@ all = [
     "prime-rl[sonic-moe]",
     "prime-rl[flash-attn-3]",
     "prime-rl[flash-attn-cute]",
+    "prime-rl[flash-mla]",
 ]
 
 [dependency-groups]
@@ -248,6 +252,7 @@ deep-gemm = false
 nixl-cu13 = false
 torchao = false
 prime-kernels = false
+flash-mla = false
 
 [tool.uv.sources]
 prime-rl-configs = { path = "packages/prime-rl-configs", editable = true }
@@ -288,6 +293,8 @@ prime-kernels = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-6f48176/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.0-6f48176/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
+# FlashMLA nv_dev @ b7643bd built against torch 2.13.0+cu130 (see the release notes for the recipe).
+flash-mla = { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" }
 flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" }
 mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" }
 flash_attn_3 = { index = "pytorch-cu130-test" }
diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py
index 6fe768b18b9..0d073174928 100644
--- a/src/prime_rl/trainer/model.py
+++ b/src/prime_rl/trainer/model.py
@@ -365,7 +365,7 @@ def get_model(
         subconfig = getattr(model_config, subconfig_key, None)
         if subconfig is not None and hasattr(subconfig, "use_cache"):
             subconfig.use_cache = False
-    model_config.sparse_mla_backward = config.sparse_mla_backward
+    model_config.dsa_backend = config.dsa_backend
     if config.index_cache is not None:
         model_config.use_index_cache = True
         model_config.index_topk_freq = config.index_cache.topk_freq
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
index 91d76079f38..feeedc83418 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/modeling_glm_moe_dsa.py
@@ -41,7 +41,7 @@ def _sparse_mla_attention_args(config: GlmMoeDsaConfig, layer_idx: int) -> Spars
         index_topk=config.index_topk,
         use_index_cache=getattr(config, "use_index_cache", False),
         skip_topk=_index_cache_skip_topk(config, layer_idx),
-        sparse_mla_backward=getattr(config, "sparse_mla_backward", "auto"),
+        dsa_backend=getattr(config, "dsa_backend", "tilelang"),
     )
 
 
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 567568d5856..9e16c7c59dd 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -32,7 +32,7 @@ class SparseMlaAttentionArgs:
     index_topk: int
     use_index_cache: bool = False
     skip_topk: bool = False
-    sparse_mla_backward: str = "auto"
+    dsa_backend: str = "tilelang"
 
 
 def apply_rope_interleave_single(
@@ -231,7 +231,7 @@ def forward(
             position_embeddings_full=position_embeddings,
         )
 
-        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling, backward_backend=self.args.sparse_mla_backward)
+        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling, backend=self.args.dsa_backend)
         out = torch.einsum("bshk,hdk->bshd", out, w_v)
         batch_size, total_tokens = out.shape[:2]
         out = out.reshape(batch_size, total_tokens, -1)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
index 1698cd5433f..4c70bf95794 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -278,6 +278,33 @@ def sparse_mla_bwd_kernel(
     return sparse_mla_bwd_kernel
 
 
+DSA_BACKENDS = ("tilelang", "cudnn_flashmla", "tilelang_legacy")
+
+_CUDNN_SM100_CAPABILITIES = ((10, 0), (10, 3))
+
+
+def cudnn_backward_arch(q: torch.Tensor, kv: torch.Tensor) -> str | None:
+    """The cuDNN DSA backward variant for this device ("sm90" / "sm100"), or None if unsupported."""
+    if kv.shape[2] != 1:
+        return None
+    capability = torch.cuda.get_device_capability(q.device)
+    if capability[0] == 9:
+        return "sm90"
+    if capability in _CUDNN_SM100_CAPABILITIES:
+        return "sm100"
+    return None
+
+
+def flat_kv_indices(indices: torch.Tensor, S_kv: int) -> torch.Tensor:
+    """Global indices into the batch-flattened KV, with sentinel / out-of-range indices as -1.
+
+    `indices` is [B, S, kv_group, topk]; valid entries live in [0, S_kv - 1) (see sparse_mla_fwd).
+    """
+    B = indices.shape[0]
+    batch_offset = torch.arange(B, device=indices.device, dtype=indices.dtype).view(B, 1, 1, 1) * S_kv
+    return torch.where(indices <= S_kv - 2, indices + batch_offset, -1)
+
+
 def tilelang_sparse_mla_backward(
     q: torch.Tensor,
     kv: torch.Tensor,
@@ -287,6 +314,7 @@ def tilelang_sparse_mla_backward(
     lse: torch.Tensor,
     sm_scale: float | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    """TileLang backward; takes a natural-log LSE."""
     grad_out = grad_out.contiguous()
     H, dim_plus_tail_dim = q.shape[2:]
     kv_group = kv.shape[2]
@@ -300,7 +328,7 @@ def tilelang_sparse_mla_backward(
 
     delta = preprocess_kernel(out, grad_out)
     dkv = torch.zeros_like(kv, dtype=torch.float32)
-    dq = bwd_kernel(q, kv, grad_out, indices, lse, delta, dkv)
+    dq = bwd_kernel(q, kv, grad_out, indices, lse * math.log2(math.e), delta, dkv)
     dkv = postprocess_kernel(dkv)
     return dq, dkv
 
@@ -312,34 +340,35 @@ def cudnn_sparse_mla_backward(
     grad_out: torch.Tensor,
     indices: torch.Tensor,
     lse: torch.Tensor,
-    sm_scale: float | None = None,
+    sm_scale: float | None,
+    arch: str,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """cuDNN-frontend CuTe-DSL DSA backward (SM90, single KV head with K=V).
+    """cuDNN-frontend CuTe-DSL DSA backward (single KV head with K=V).
 
-    The cuDNN kernel takes flat (tokens, ...) tensors with global indices, skips
-    negative indices (no KV load, no dKV atomics) and expects a natural-log LSE.
+    The cuDNN kernels take flat (tokens, ...) tensors with global indices, skip
+    negative indices (no KV load, no dKV atomics) and expect a natural-log LSE.
     """
-    from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90
-
     B, S, H, D_qk = q.shape
     S_kv = kv.shape[1]
     topk = indices.shape[-1]
-    # Sentinel / out-of-range indices (see sparse_mla_fwd) become -1; valid ones are
-    # offset into the batch-flattened KV.
-    batch_offset = torch.arange(B, device=indices.device, dtype=indices.dtype).view(B, 1, 1, 1) * S_kv
-    flat_indices = torch.where(indices <= S_kv - 2, indices + batch_offset, -1).view(B * S, topk)
-    # The forward kernel returns a base-2 LSE.
-    lse_ln = lse * math.log(2.0)
-
-    dq, dkv = flash_attn_bwd_sm90(
+    args = (
         q.view(B * S, H, D_qk),
         kv.view(B * S_kv, D_qk),
         out.view(B * S, H, -1),
         grad_out.contiguous().view(B * S, H, -1),
-        lse_ln.view(B * S, H),
-        softmax_scale=sm_scale,
-        topk_idxs=flat_indices,
+        lse.view(B * S, H),
     )
+    flat_indices = flat_kv_indices(indices, S_kv).view(B * S, topk)
+    if arch == "sm90":
+        from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90
+
+        dq, dkv = flash_attn_bwd_sm90(*args, softmax_scale=sm_scale, topk_idxs=flat_indices)
+    else:
+        # Untested: no SM100 hardware has run this path. A -inf sink leaves the softmax unchanged.
+        from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 import flash_attn_bwd_sm100
+
+        no_sink = torch.full((H,), float("-inf"), dtype=torch.float32, device=q.device)
+        dq, dkv, _ = flash_attn_bwd_sm100(*args, no_sink, flat_indices, softmax_scale=sm_scale)
     return dq.view_as(q), dkv.view_as(kv)
 
 
@@ -352,8 +381,9 @@ def sparse_mla_backward(
     indices: torch.Tensor,
     lse: torch.Tensor,
     sm_scale: float | None = None,
-    backend: str = "auto",
+    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sparse MLA backward for the `backend` route (see `sparse_mla`); `lse` is natural-log."""
     assert q.is_contiguous()
     assert kv.is_contiguous()
     assert indices.is_contiguous()
@@ -365,14 +395,14 @@ def sparse_mla_backward(
     topk = indices.shape[-1]
     assert indices.shape == (B, S, kv_group, topk)
     assert lse.shape == (B, S, H)
-
-    cudnn_supported = torch.cuda.get_device_capability(q.device)[0] == 9 and kv_group == 1
-    if backend == "auto":
-        backend = "cudnn" if cudnn_supported else "tilelang"
-    if backend == "cudnn":
-        assert cudnn_supported, "cuDNN sparse MLA backward requires an SM90 GPU and a single KV head"
-        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
-    assert backend == "tilelang", f"Unknown sparse MLA backward backend: {backend}"
+    assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"
+
+    arch = cudnn_backward_arch(q, kv)
+    if backend == "cudnn_flashmla":
+        assert arch is not None, "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
+        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
+    if backend == "tilelang" and arch == "sm90":
+        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
     return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
 
 
@@ -385,6 +415,6 @@ def _sparse_mla_backward_fake(
     indices: torch.Tensor,
     lse: torch.Tensor,
     sm_scale: float | None = None,
-    backend: str = "auto",
+    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), torch.empty_like(kv)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
index c0cd428e4f3..4d5f1162bb4 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -13,11 +13,19 @@
     # This is expected on CPU-only machines
     pass
 
+import importlib.util
+import math
+
 import tilelang
 import torch
 from tilelang import language as T
 
-from prime_rl.trainer.models.kernels.sparse_mla_bwd import sparse_mla_backward
+from prime_rl.trainer.models.kernels.sparse_mla_bwd import (
+    DSA_BACKENDS,
+    cudnn_backward_arch,
+    flat_kv_indices,
+    sparse_mla_backward,
+)
 
 
 @tilelang.jit(
@@ -186,6 +194,30 @@ def main(
     return main
 
 
+def flashmla_sparse_mla_forward(
+    q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sm_scale: float, d_v: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """FlashMLA sparse prefill forward; returns a natural-log LSE."""
+    from flash_mla import flash_mla_sparse_fwd
+
+    B, S, H, D_qk = q.shape
+    S_kv, kv_group = kv.shape[1:3]
+    topk = indices.shape[-1]
+    # Shortest prefix holding every valid index, so the kernel skips the sentinel tail
+    # (the indexer sorts valid indices first); invalid entries inside it are -1.
+    positions = torch.arange(1, topk + 1, device=indices.device, dtype=torch.int32)
+    topk_length = torch.where(indices <= S_kv - 2, positions, 0).amax(-1).view(B * S)
+    out, _max_logits, lse = flash_mla_sparse_fwd(
+        q.view(B * S, H, D_qk),
+        kv.view(B * S_kv, kv_group, D_qk),
+        flat_kv_indices(indices, S_kv).view(B * S, kv_group, topk),
+        sm_scale,
+        d_v,
+        topk_length=topk_length,
+    )
+    return out.view(B, S, H, d_v), lse.view(B, S, H)
+
+
 @torch.library.custom_op("prime_rl::sparse_mla", mutates_args=())
 def sparse_mla(
     q: torch.Tensor,
@@ -196,9 +228,17 @@ def sparse_mla(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
-    backward_backend: str = "auto",
+    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Sparse MLA attention over the top-k `indices`; returns the output and its natural-log LSE.
+
+    `backend` picks the kernels for the forward and (via autograd) the backward:
+    ``tilelang`` is TileLang forward + cuDNN backward on SM90 with a single KV head, TileLang
+    backward otherwise; ``cudnn_flashmla`` is FlashMLA forward + cuDNN backward (SM90; SM100 /
+    SM103 untested); ``tilelang_legacy`` is TileLang for both.
+    """
     assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
+    assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"
     batch, seq_len, heads, dim_plus_tail_dim = q.shape
     _, seq_len_kv, kv_group, _ = kv.shape
 
@@ -210,6 +250,17 @@ def sparse_mla(
     _, _, _, topk = indices.shape
     assert indices.shape == (batch, seq_len, kv_group, topk)
 
+    if backend == "cudnn_flashmla":
+        assert cudnn_backward_arch(q, kv) is not None, (
+            "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
+        )
+        assert importlib.util.find_spec("flash_mla") is not None, (
+            "dsa_backend='cudnn_flashmla' requires FlashMLA; install it with `uv sync --extra flash-mla`"
+        )
+        if sm_scale is None:
+            sm_scale = dim_plus_tail_dim**-0.5
+        return flashmla_sparse_mla_forward(q, kv, indices, sm_scale, d_v)
+
     kernel = sparse_mla_fwd(
         heads,
         dim,
@@ -223,7 +274,8 @@ def sparse_mla(
         threads=threads,
     )
     out, lse = kernel(q, kv, indices)
-    return out, lse
+    # The TileLang kernel returns a base-2 LSE.
+    return out, lse * math.log(2.0)
 
 
 @sparse_mla.register_fake
@@ -236,17 +288,17 @@ def _sparse_mla_fake(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
-    backward_backend: str = "auto",
+    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return q.new_empty((*q.shape[:-1], d_v)), q.new_empty(q.shape[:-1], dtype=torch.float32)
 
 
 def _sparse_mla_setup_context(ctx, inputs, output) -> None:
-    q, kv, indices, sm_scale, _d_v, _block_I, _num_stages, _threads, backward_backend = inputs
+    q, kv, indices, sm_scale, _d_v, _block_I, _num_stages, _threads, backend = inputs
     out, lse = output
     ctx.save_for_backward(q, kv, out, indices, lse)
     ctx.sm_scale = sm_scale
-    ctx.backward_backend = backward_backend
+    ctx.backend = backend
     ctx.mark_non_differentiable(lse)
 
 
@@ -260,7 +312,7 @@ def _sparse_mla_autograd_backward(ctx, grad_out: torch.Tensor, _grad_lse: torch.
         indices,
         lse.detach(),
         ctx.sm_scale,
-        ctx.backward_backend,
+        ctx.backend,
     )
     return dq, dkv, None, None, None, None, None, None, None
 
diff --git a/uv.lock b/uv.lock
index d1897a0ed44..ddac7c69fc7 100644
--- a/uv.lock
+++ b/uv.lock
@@ -22,6 +22,7 @@ deep-gemm = false
 dion = false
 fastokens = false
 flash-attn-3 = false
+flash-mla = false
 harbor = "2026-08-11T00:00:00Z"
 nixl-cu13 = false
 prime = false
@@ -1860,6 +1861,14 @@ dependencies = [
     { name = "transformers" },
 ]
 
+[[package]]
+name = "flash-mla"
+version = "1.0.0+cu.13.0.torch.2.13.b7643bd"
+source = { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl" }
+wheels = [
+    { url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl", hash = "sha256:0bdb32340979c739659b305537dbe0445ba9f3148a4ce2f97f4fd170887c01e4" },
+]
+
 [[package]]
 name = "flashinfer-python"
 version = "0.6.18.post1"
@@ -4987,6 +4996,7 @@ all = [
     { name = "flash-attn-3", marker = "sys_platform == 'linux'" },
     { name = "flash-attn-4", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
+    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
     { name = "mamba-ssm", marker = "sys_platform == 'linux'" },
     { name = "modelexpress", marker = "sys_platform == 'linux'" },
     { name = "mooncake-transfer-engine", marker = "sys_platform == 'linux'" },
@@ -5037,6 +5047,9 @@ flash-attn-3 = [
 flash-attn-cute = [
     { name = "flash-attn-4", marker = "sys_platform == 'linux'" },
 ]
+flash-mla = [
+    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" },
+]
 gpu = [
     { name = "dion", marker = "sys_platform == 'linux'" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux'" },
@@ -5100,6 +5113,7 @@ requires-dist = [
     { name = "flash-attn-3", marker = "sys_platform == 'linux' and extra == 'flash-attn-3'", index = "https://download.pytorch.org/whl/test/cu130" },
     { name = "flash-attn-4", marker = "sys_platform == 'linux' and extra == 'flash-attn-cute'", git = "https://github.com/Dao-AILab/flash-attention.git?subdirectory=flash_attn%2Fcute&rev=0251105" },
     { name = "flash-linear-attention", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/fla-org/flash-linear-attention" },
+    { name = "flash-mla", marker = "platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'flash-mla'", url = "https://github.com/NovaSky-AI/skyrl-wheels/releases/download/cu13torch2.13-flashmla-b7643bd/flash_mla-1.0.0%2Bcu.13.0.torch.2.13.b7643bd-cp312-cp312-linux_x86_64.whl" },
     { name = "jaxtyping", specifier = ">=0.3.2" },
     { name = "loguru", specifier = ">=0.7.3" },
     { name = "mamba-ssm", marker = "sys_platform == 'linux' and extra == 'gpu'", git = "https://github.com/state-spaces/mamba?rev=f577286" },
@@ -5122,6 +5136,7 @@ requires-dist = [
     { name = "prime-rl", extras = ["flash-attn"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["flash-attn-3"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["flash-attn-cute"], marker = "extra == 'all'" },
+    { name = "prime-rl", extras = ["flash-mla"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["gpu"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["quack"], marker = "extra == 'all'" },
     { name = "prime-rl", extras = ["sonic-moe"], marker = "extra == 'all'" },
@@ -5161,7 +5176,7 @@ requires-dist = [
     { name = "wandb", specifier = ">=0.26.1" },
     { name = "wandb-workspaces", specifier = ">=0.4.3" },
 ]
-provides-extras = ["gpu", "dashboard", "dynamo", "flash-attn", "flash-attn-3", "flash-attn-cute", "disagg", "quack", "sonic-moe", "kernels", "all"]
+provides-extras = ["gpu", "dashboard", "dynamo", "flash-attn", "flash-attn-3", "flash-attn-cute", "disagg", "quack", "sonic-moe", "flash-mla", "kernels", "all"]
 
 [package.metadata.requires-dev]
 dev = [

From 1b935b6dbcb085aeaae26caffc7d9e03ad43450c Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 06:48:36 +0000
Subject: [PATCH 05/11] examples(glm-5.3): SFT on cudnn_flashmla DSA backend,
 add 12-node 16k overlay

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 examples/advanced/glm-5.3/README.md           | 12 ++++++++++++
 .../glm-5.3/sft/h200/16k-12-node.toml         | 19 +++++++++++++++++++
 examples/advanced/glm-5.3/sft/h200/base.toml  |  1 +
 3 files changed, 32 insertions(+)
 create mode 100644 examples/advanced/glm-5.3/sft/h200/16k-12-node.toml

diff --git a/examples/advanced/glm-5.3/README.md b/examples/advanced/glm-5.3/README.md
index 8b5fa9c67ea..e525ee479bd 100644
--- a/examples/advanced/glm-5.3/README.md
+++ b/examples/advanced/glm-5.3/README.md
@@ -113,5 +113,17 @@ This will start a SFT run with the following configuration:
 
 - The model is `zai-org/GLM-5.3-BF16`
 - The data is `PrimeIntellect/INTELLECT-3-SFT-10K` (math split)
+- 8 nodes at 131k context: CP 8, EP 8, full activation checkpointing with activation offloading, optimizer state offloaded to CPU
+- DSA attention runs on the FlashMLA sparse forward and the cuDNN sparse backward (`model.dsa_backend = "cudnn_flashmla"`)
 
 For a fake-data dry run, append [`fake.toml`](sft/h200/fake.toml) instead. You can use the same dashboard to monitor the SFT run.
+
+### 12 nodes at 16k
+
+For 16k-token samples, append [`16k-12-node.toml`](sft/h200/16k-12-node.toml) after the data overlay:
+
+```bash
+uv run sft @ examples/advanced/glm-5.3/sft/h200/base.toml @ examples/advanced/glm-5.3/sft/h200/math-10k.toml @ examples/advanced/glm-5.3/sft/h200/16k-12-node.toml
+```
+
+On 12 nodes the sharded model state fits in GPU memory without CP or any offloading. Each GPU runs one 16k sequence per step (batch 96), and transformer blocks compile with `fullgraph = true`. Full activation checkpointing stays on; peak memory is about 115 GiB per GPU.
diff --git a/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml b/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml
new file mode 100644
index 00000000000..f72b0b16352
--- /dev/null
+++ b/examples/advanced/glm-5.3/sft/h200/16k-12-node.toml
@@ -0,0 +1,19 @@
+# 12-node, 16k-token overlay: enough GPUs that the model state fits without CP or offloading,
+# so each GPU runs one 16k sequence per micro-batch and blocks compile with fullgraph.
+# Compose after the data overlay, e.g. base.toml @ math-10k.toml @ 16k-12-node.toml.
+
+[model]
+cp = 1
+optim_cpu_offload = false
+ac_offloading = "None"
+
+[model.compile]
+fullgraph = true
+
+[data]
+seq_len = 16384
+batch_size = 96
+micro_batch_size = 1
+
+[deployment]
+num_train_nodes = 12
diff --git a/examples/advanced/glm-5.3/sft/h200/base.toml b/examples/advanced/glm-5.3/sft/h200/base.toml
index 2d402be3dd9..db3fd6e969c 100644
--- a/examples/advanced/glm-5.3/sft/h200/base.toml
+++ b/examples/advanced/glm-5.3/sft/h200/base.toml
@@ -10,6 +10,7 @@ attn = "flash_attention_2"
 cp = 8
 ep = 8
 optim_cpu_offload = true
+dsa_backend = "cudnn_flashmla"
 
 [model.ac]
 freq = 1

From f2372e9aedc11ea1e69db85e219ac0123bcc9203 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 07:23:40 +0000
Subject: [PATCH 06/11] refactor(trainer): dsa_backend is tilelang (TileLang
 fwd+bwd) or cudnn_flashmla

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 packages/prime-rl-configs/src/prime_rl/configs/trainer.py | 4 ++--
 src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py     | 6 ++----
 src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py     | 5 ++---
 3 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 024aa79a592..3067560151a 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -325,8 +325,8 @@ class ModelConfig(BaseModelConfig):
 
     quantization: QuantizationConfig | None = None
 
-    dsa_backend: Literal["tilelang", "cudnn_flashmla", "tilelang_legacy"] = "tilelang"
-    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward with the cuDNN frontend CuTe-DSL backward on Hopper (SM90) and the TileLang backward elsewhere. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested). ``tilelang_legacy`` runs TileLang for both."""
+    dsa_backend: Literal["tilelang", "cudnn_flashmla"] = "tilelang"
+    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
 
     index_cache: IndexCacheConfig | None = None
     """DSA IndexCache sub-configuration. If set, sparse-attention top-k indices are reused across decoder layers per the configured schedule (mirrors vLLM's IndexCache HF overrides). If None, every layer recomputes its own indices."""
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
index 4c70bf95794..08e51ef663c 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -278,7 +278,7 @@ def sparse_mla_bwd_kernel(
     return sparse_mla_bwd_kernel
 
 
-DSA_BACKENDS = ("tilelang", "cudnn_flashmla", "tilelang_legacy")
+DSA_BACKENDS = ("tilelang", "cudnn_flashmla")
 
 _CUDNN_SM100_CAPABILITIES = ((10, 0), (10, 3))
 
@@ -397,12 +397,10 @@ def sparse_mla_backward(
     assert lse.shape == (B, S, H)
     assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"
 
-    arch = cudnn_backward_arch(q, kv)
     if backend == "cudnn_flashmla":
+        arch = cudnn_backward_arch(q, kv)
         assert arch is not None, "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
         return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
-    if backend == "tilelang" and arch == "sm90":
-        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
     return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
 
 
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
index 4d5f1162bb4..cfb58852be5 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -233,9 +233,8 @@ def sparse_mla(
     """Sparse MLA attention over the top-k `indices`; returns the output and its natural-log LSE.
 
     `backend` picks the kernels for the forward and (via autograd) the backward:
-    ``tilelang`` is TileLang forward + cuDNN backward on SM90 with a single KV head, TileLang
-    backward otherwise; ``cudnn_flashmla`` is FlashMLA forward + cuDNN backward (SM90; SM100 /
-    SM103 untested); ``tilelang_legacy`` is TileLang for both.
+    ``tilelang`` is TileLang for both; ``cudnn_flashmla`` is FlashMLA forward + cuDNN backward
+    (SM90; SM100 / SM103 untested).
     """
     assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
     assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"

From 8b5e3b65d68ee548ed85b61c3613eac5aadaa5e7 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 23:16:18 +0000
Subject: [PATCH 07/11] perf(trainer): cuDNN frontend FP8 indexer scoring +
 radix top-k under dsa_backend=cudnn_flashmla

On SM90 the cudnn_flashmla route now selects DSA indices with the cuDNN
frontend SM90 FP8 indexer forward (ratio=1 token-level causal mask, packed
documents as THD segments with q_causal_offsets for the CP shard) and the
cuDNN radix top-k (per-row causal lengths, chunked to bound its scratch
buffer), behind the prime_rl::cudnn_fp8_indexer custom op (AC mandatory
save). Same UE8M0 FP8 quantization as the Triton indexer; ~2x faster per
layer on H200. SM100 and the tilelang backend keep the Triton indexer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/prime_rl/configs/trainer.py           |   2 +-
 .../trainer/activation_checkpointing.py       |   1 +
 .../glm_moe_dsa/sparse_mla_attention.py       |   8 +-
 .../trainer/models/kernels/cudnn_indexer.py   | 114 ++++++++++++++++++
 tests/unit/train/models/test_fp8_indexer.py   |  37 ++++++
 5 files changed, 160 insertions(+), 2 deletions(-)
 create mode 100644 src/prime_rl/trainer/models/kernels/cudnn_indexer.py

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 3067560151a..ae08afa21f4 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -326,7 +326,7 @@ class ModelConfig(BaseModelConfig):
     quantization: QuantizationConfig | None = None
 
     dsa_backend: Literal["tilelang", "cudnn_flashmla"] = "tilelang"
-    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
+    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward, and on SM90 also the cuDNN frontend FP8 indexer scoring and top-k (SM100/SM103 keep the Triton indexer); it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
 
     index_cache: IndexCacheConfig | None = None
     """DSA IndexCache sub-configuration. If set, sparse-attention top-k indices are reused across decoder layers per the configured schedule (mirrors vLLM's IndexCache HF overrides). If None, every layer recomputes its own indices."""
diff --git a/src/prime_rl/trainer/activation_checkpointing.py b/src/prime_rl/trainer/activation_checkpointing.py
index 035bfe119bd..bbdfbd9bc2b 100644
--- a/src/prime_rl/trainer/activation_checkpointing.py
+++ b/src/prime_rl/trainer/activation_checkpointing.py
@@ -33,6 +33,7 @@
     {
         "aten::topk",
         "prime_rl::fp8_indexer",
+        "prime_rl::cudnn_fp8_indexer",
         "prime_kernels::select_indexed_blocks",
         "prime_rl::record_moe_routing_statistics",
     }
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 9e16c7c59dd..ef5be210c7d 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -4,6 +4,7 @@
 import torch.distributed as dist
 from torch import nn
 
+from prime_rl.trainer.models.kernels.cudnn_indexer import cudnn_fp8_indexer
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
@@ -56,6 +57,8 @@ def __init__(self, args: SparseMlaAttentionArgs):
         self.k_norm = LayerNorm(dim=self.head_dim, eps=1e-6)
         self.weights_proj = nn.Linear(args.hidden_size, self.n_head, bias=False)
         self.weight_scale = (self.head_dim**-0.5) * (self.n_head**-0.5)
+        # The cuDNN indexer kernels are SM90 ports; SM100 keeps the Triton indexer.
+        self.use_cudnn_indexer = args.dsa_backend == "cudnn_flashmla" and torch.cuda.get_device_capability()[0] == 9
 
     @torch.no_grad()
     def compute_sparse_indices(
@@ -106,7 +109,10 @@ def compute_sparse_indices(
         q_idx = torch.cat([q_pe, q_nope], dim=-1)
         k_idx = torch.cat([k_pe, k_nope], dim=-1)
 
-        indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale)
+        if self.use_cudnn_indexer:
+            indices = cudnn_fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk)
+        else:
+            indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale)
         # indices shape: [S_local, topk] in K's coordinate space (sentinel = s_full)
         # KV passed to sparse MLA has length s_full + 1 (sentinel zeros at index s_full).
         return indices.view(1, s_local, 1, index_topk)
diff --git a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
new file mode 100644
index 00000000000..adc25762b25
--- /dev/null
+++ b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
@@ -0,0 +1,114 @@
+"""cuDNN frontend FP8 indexer for DSA sparse token selection (SM90).
+
+Same contract and scoring formula as `fp8_indexer`, I_{t,s} = Σ_j w_{t,j} · ReLU(q_{t,j} · k_s),
+with the same UE8M0 per-token FP8 quantization, but scored by the cuDNN SM90 indexer forward and
+selected by the cuDNN radix top-k. Packed documents become THD segments, so each document is scored
+only against its own keys.
+"""
+
+import torch
+
+from prime_rl.trainer.models.kernels.fp8_indexer import per_token_group_quant_fp8
+
+# Upper bound on fp32 logits per top-k launch: the kernel allocates a scratch buffer of twice the
+# logits' size, which would be 16 GiB for a 16k x 131k CP shard in one launch.
+TOPK_CHUNK_ELEMS = 1 << 27
+
+
+def _document_segments(ks: torch.Tensor, ke: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+    """THD segments of the queries, one per document (run of equal `ks`).
+
+    Within a document `ke` advances by one per query, so the segment's keys are `[ks, ke[-1])` and
+    its first query sits at doc-local position `ke[0] - ks - 1`. Consecutive documents' key ranges
+    abut, so the key segments tile the gathered K from the first document's start.
+    """
+    S_q = ks.shape[0]
+    is_start = torch.ones(S_q, dtype=torch.bool, device=ks.device)
+    is_start[1:] = ks[1:] != ks[:-1]
+    starts = is_start.nonzero().squeeze(1).to(torch.int32)
+    cu_seqlens_q = torch.cat([starts, starts.new_tensor([S_q])])
+    cu_seqlens_k = torch.cat([ks[starts], ke[-1:]])
+    q_causal_offsets = ke[starts] - ks[starts] - 1
+    max_seqlen_q, max_seqlen_k = torch.stack([cu_seqlens_q.diff().max(), (ke - ks).max()]).tolist()
+    return cu_seqlens_q, cu_seqlens_k, q_causal_offsets, max_seqlen_q, max_seqlen_k
+
+
+def _chunked_top_k(logits: torch.Tensor, lengths: torch.Tensor, topk: int) -> torch.Tensor:
+    from cudnn.deepseek_sparse_attention import indexer_top_k_wrapper
+
+    num_rows, num_cols = logits.shape
+    chunk_rows = max(256, TOPK_CHUNK_ELEMS // num_cols)
+    return torch.cat(
+        [
+            indexer_top_k_wrapper(logits[lo : lo + chunk_rows], lengths[lo : lo + chunk_rows], topk, return_val=False)[
+                "indices"
+            ]
+            for lo in range(0, num_rows, chunk_rows)
+        ]
+    )
+
+
+@torch.library.custom_op("prime_rl::cudnn_fp8_indexer", mutates_args=())
+def cudnn_fp8_indexer(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    w: torch.Tensor,
+    ks: torch.Tensor,
+    ke: torch.Tensor,
+    topk: int,
+) -> torch.Tensor:
+    """cuDNN FP8 indexer: UE8M0 quantization + cuDNN SM90 scoring + cuDNN radix top-k.
+
+    Args:
+        q: [S_q, H, D] bf16 query vectors per head (H in {32, 64}, D = 128)
+        k: [S_k, D] bf16 key vectors (shared across heads)
+        w: [S_q, H] bf16 per-head weights
+        ks: [S_q] int32 document start per query (in K's coordinate system)
+        ke: [S_q] int32 causal end per query (= global position + 1); must advance by one per
+            query within a document
+        topk: number of top indices to return (<= 2048)
+
+    Returns:
+        [S_q, topk] int32 selected token indices per query (sentinel = S_k), in no particular order
+    """
+    from cudnn.deepseek_sparse_attention.indexer_forward._interface_sm90 import indexer_fwd
+
+    S_q, H, D = q.shape
+    S_k = k.shape[0]
+
+    q_fp8, q_scales = per_token_group_quant_fp8(q.reshape(S_q * H, D).contiguous(), group_size=D, use_ue8m0=True)
+    k_fp8, k_scales = per_token_group_quant_fp8(k.contiguous(), group_size=D, use_ue8m0=True)
+    q_scales = q_scales.view(S_q, H)
+    # Fold q's descale into the weights (fp32), as fp8_indexer does; K's descale goes to the kernel.
+    w = w * q_scales
+
+    cu_seqlens_q, cu_seqlens_k, q_causal_offsets, max_seqlen_q, max_seqlen_k = _document_segments(ks, ke)
+    # Rows 32-byte aligned for the top-k kernel's vector loads.
+    num_cols = (max_seqlen_k + 7) // 8 * 8
+
+    # ratio=1 is the plain token-level causal mask; scores are doc-local: column j is key ks + j.
+    logits = indexer_fwd(
+        q_fp8.view(S_q, H, D),
+        k_fp8.view(S_k, 1, D),
+        w,
+        ratio=1,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=num_cols,
+        q_causal_offsets=q_causal_offsets,
+        precision="fp8",
+        q_scale=q_scales,
+        k_scale=k_scales,
+    )
+
+    lengths = ke - ks
+    local = _chunked_top_k(logits, lengths, topk)
+    # Slots past a short row's length come back as -1.
+    valid = (local >= 0) & (local < lengths[:, None])
+    return torch.where(valid, local + ks[:, None], S_k).to(torch.int32)
+
+
+@cudnn_fp8_indexer.register_fake
+def _cudnn_fp8_indexer_fake(q, k, w, ks, ke, topk):
+    return q.new_empty((q.shape[0], topk), dtype=torch.int32)
diff --git a/tests/unit/train/models/test_fp8_indexer.py b/tests/unit/train/models/test_fp8_indexer.py
index a1f210b113f..c08ba8f87bb 100644
--- a/tests/unit/train/models/test_fp8_indexer.py
+++ b/tests/unit/train/models/test_fp8_indexer.py
@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from prime_rl.trainer.models.kernels.cudnn_indexer import cudnn_fp8_indexer
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 
 pytestmark = [
@@ -109,3 +110,39 @@ def test_selection_agreement(segments):
     mean, p1 = agreement.mean().item(), agreement.quantile(0.01).item()
     assert mean > AGREEMENT_MEAN, f"mean set agreement {mean} below {AGREEMENT_MEAN}"
     assert p1 > AGREEMENT_P1, f"p1 set agreement {p1} below {AGREEMENT_P1}"
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9,
+    reason="the cuDNN indexer kernels are SM90-only",
+)
+def test_cudnn_matches_triton_on_causal_documents():
+    """The cuDNN indexer against the Triton one on a CP-style shard of packed causal documents.
+
+    The queries are the last 2560 tokens of a 4096-token row, so the shard starts mid-document and
+    `ks` / `ke` live in the gathered K's coordinates. Both kernels quantize identically and their
+    scores agree to fp32 rounding, so the selections only differ by ties (e.g. keys every head's
+    ReLU zeroes) at the top-k boundary.
+    """
+    num_q, num_k, doc_lens = 2560, 4096, (1500, 1300, 200, 1096)
+    doc_start = torch.cat([torch.full((n,), sum(doc_lens[:i])) for i, n in enumerate(doc_lens)]).cuda()
+    ks = doc_start[num_k - num_q :].int().contiguous()
+    ke = torch.arange(num_k - num_q + 1, num_k + 1, device="cuda", dtype=torch.int32)
+
+    q = torch.randn(num_q, HEADS, DIM, device="cuda").bfloat16()
+    k = torch.randn(num_k, DIM, device="cuda").bfloat16()
+    w = torch.randn(num_q, HEADS, device="cuda").bfloat16()
+
+    def selected(picks):
+        mask = torch.zeros(num_q, num_k + 1, dtype=torch.bool, device="cuda")
+        return mask.scatter_(1, picks.long(), True)[:, :num_k]
+
+    triton_selected = selected(fp8_indexer(q, k, w, ks, ke, TOPK))
+    cudnn_selected = selected(cudnn_fp8_indexer(q, k, w, ks, ke, TOPK))
+
+    assert torch.equal(cudnn_selected.sum(-1), (ke - ks).clamp(max=TOPK)), "wrong number of readable picks"
+    saturated = (ke - ks) > TOPK
+    assert torch.equal(cudnn_selected[~saturated], triton_selected[~saturated])
+    assert saturated.any(), "vacuous probe: no query has more readable keys than it can pick"
+    agreement = (cudnn_selected & triton_selected).sum(-1)[saturated] / TOPK
+    assert agreement.mean().item() > 0.999 and agreement.min().item() > 0.99

From 513ccfcd30920109206e77940e56e7614d09a2d6 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Sun, 4 Oct 2026 23:55:59 +0000
Subject: [PATCH 08/11] docs(glm-5.3): note the cuDNN indexer on
 cudnn_flashmla; comment top-k tie behaviour

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 examples/advanced/glm-5.3/README.md                  | 2 +-
 src/prime_rl/trainer/models/kernels/cudnn_indexer.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/advanced/glm-5.3/README.md b/examples/advanced/glm-5.3/README.md
index e525ee479bd..4a010c20225 100644
--- a/examples/advanced/glm-5.3/README.md
+++ b/examples/advanced/glm-5.3/README.md
@@ -114,7 +114,7 @@ This will start a SFT run with the following configuration:
 - The model is `zai-org/GLM-5.3-BF16`
 - The data is `PrimeIntellect/INTELLECT-3-SFT-10K` (math split)
 - 8 nodes at 131k context: CP 8, EP 8, full activation checkpointing with activation offloading, optimizer state offloaded to CPU
-- DSA attention runs on the FlashMLA sparse forward and the cuDNN sparse backward (`model.dsa_backend = "cudnn_flashmla"`)
+- DSA attention runs on the FlashMLA sparse forward and the cuDNN sparse backward, and the DSA indexer on the cuDNN FP8 scoring and top-k kernels (`model.dsa_backend = "cudnn_flashmla"`)
 
 For a fake-data dry run, append [`fake.toml`](sft/h200/fake.toml) instead. You can use the same dashboard to monitor the SFT run.
 
diff --git a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
index adc25762b25..2801d2a4934 100644
--- a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
+++ b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
@@ -36,6 +36,8 @@ def _document_segments(ks: torch.Tensor, ke: torch.Tensor) -> tuple[torch.Tensor
 def _chunked_top_k(logits: torch.Tensor, lengths: torch.Tensor, topk: int) -> torch.Tensor:
     from cudnn.deepseek_sparse_attention import indexer_top_k_wrapper
 
+    # Radix select is ~5x faster than torch.topk on real indexer scores, but ~1.4x slower when the
+    # k-th value is tied across most of a row (e.g. the all-zero scores of a degenerate random init).
     num_rows, num_cols = logits.shape
     chunk_rows = max(256, TOPK_CHUNK_ELEMS // num_cols)
     return torch.cat(

From 03887532f925dfff287b379769ed349548694300 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Mon, 5 Oct 2026 03:56:21 +0000
Subject: [PATCH 09/11] Revert "docs(glm-5.3): note the cuDNN indexer on
 cudnn_flashmla; comment top-k tie behaviour"

This reverts commit 06f5c6086cdde24012b5df3cdf0a4ff20f020d49.
---
 examples/advanced/glm-5.3/README.md                  | 2 +-
 src/prime_rl/trainer/models/kernels/cudnn_indexer.py | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/examples/advanced/glm-5.3/README.md b/examples/advanced/glm-5.3/README.md
index 4a010c20225..e525ee479bd 100644
--- a/examples/advanced/glm-5.3/README.md
+++ b/examples/advanced/glm-5.3/README.md
@@ -114,7 +114,7 @@ This will start a SFT run with the following configuration:
 - The model is `zai-org/GLM-5.3-BF16`
 - The data is `PrimeIntellect/INTELLECT-3-SFT-10K` (math split)
 - 8 nodes at 131k context: CP 8, EP 8, full activation checkpointing with activation offloading, optimizer state offloaded to CPU
-- DSA attention runs on the FlashMLA sparse forward and the cuDNN sparse backward, and the DSA indexer on the cuDNN FP8 scoring and top-k kernels (`model.dsa_backend = "cudnn_flashmla"`)
+- DSA attention runs on the FlashMLA sparse forward and the cuDNN sparse backward (`model.dsa_backend = "cudnn_flashmla"`)
 
 For a fake-data dry run, append [`fake.toml`](sft/h200/fake.toml) instead. You can use the same dashboard to monitor the SFT run.
 
diff --git a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
index 2801d2a4934..adc25762b25 100644
--- a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
+++ b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
@@ -36,8 +36,6 @@ def _document_segments(ks: torch.Tensor, ke: torch.Tensor) -> tuple[torch.Tensor
 def _chunked_top_k(logits: torch.Tensor, lengths: torch.Tensor, topk: int) -> torch.Tensor:
     from cudnn.deepseek_sparse_attention import indexer_top_k_wrapper
 
-    # Radix select is ~5x faster than torch.topk on real indexer scores, but ~1.4x slower when the
-    # k-th value is tied across most of a row (e.g. the all-zero scores of a degenerate random init).
     num_rows, num_cols = logits.shape
     chunk_rows = max(256, TOPK_CHUNK_ELEMS // num_cols)
     return torch.cat(

From fb5da5bd29dc4885d7aa7230f2eeee5f7eba2781 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Mon, 5 Oct 2026 03:56:21 +0000
Subject: [PATCH 10/11] Revert "perf(trainer): cuDNN frontend FP8 indexer
 scoring + radix top-k under dsa_backend=cudnn_flashmla"

This reverts commit 5724fef94f6b3ef8d6ead9663d63cd45f37ccb87.
---
 .../src/prime_rl/configs/trainer.py           |   2 +-
 .../trainer/activation_checkpointing.py       |   1 -
 .../glm_moe_dsa/sparse_mla_attention.py       |   8 +-
 .../trainer/models/kernels/cudnn_indexer.py   | 114 ------------------
 tests/unit/train/models/test_fp8_indexer.py   |  37 ------
 5 files changed, 2 insertions(+), 160 deletions(-)
 delete mode 100644 src/prime_rl/trainer/models/kernels/cudnn_indexer.py

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index ae08afa21f4..3067560151a 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -326,7 +326,7 @@ class ModelConfig(BaseModelConfig):
     quantization: QuantizationConfig | None = None
 
     dsa_backend: Literal["tilelang", "cudnn_flashmla"] = "tilelang"
-    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward, and on SM90 also the cuDNN frontend FP8 indexer scoring and top-k (SM100/SM103 keep the Triton indexer); it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
+    """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
 
     index_cache: IndexCacheConfig | None = None
     """DSA IndexCache sub-configuration. If set, sparse-attention top-k indices are reused across decoder layers per the configured schedule (mirrors vLLM's IndexCache HF overrides). If None, every layer recomputes its own indices."""
diff --git a/src/prime_rl/trainer/activation_checkpointing.py b/src/prime_rl/trainer/activation_checkpointing.py
index bbdfbd9bc2b..035bfe119bd 100644
--- a/src/prime_rl/trainer/activation_checkpointing.py
+++ b/src/prime_rl/trainer/activation_checkpointing.py
@@ -33,7 +33,6 @@
     {
         "aten::topk",
         "prime_rl::fp8_indexer",
-        "prime_rl::cudnn_fp8_indexer",
         "prime_kernels::select_indexed_blocks",
         "prime_rl::record_moe_routing_statistics",
     }
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index ef5be210c7d..9e16c7c59dd 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -4,7 +4,6 @@
 import torch.distributed as dist
 from torch import nn
 
-from prime_rl.trainer.models.kernels.cudnn_indexer import cudnn_fp8_indexer
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
@@ -57,8 +56,6 @@ def __init__(self, args: SparseMlaAttentionArgs):
         self.k_norm = LayerNorm(dim=self.head_dim, eps=1e-6)
         self.weights_proj = nn.Linear(args.hidden_size, self.n_head, bias=False)
         self.weight_scale = (self.head_dim**-0.5) * (self.n_head**-0.5)
-        # The cuDNN indexer kernels are SM90 ports; SM100 keeps the Triton indexer.
-        self.use_cudnn_indexer = args.dsa_backend == "cudnn_flashmla" and torch.cuda.get_device_capability()[0] == 9
 
     @torch.no_grad()
     def compute_sparse_indices(
@@ -109,10 +106,7 @@ def compute_sparse_indices(
         q_idx = torch.cat([q_pe, q_nope], dim=-1)
         k_idx = torch.cat([k_pe, k_nope], dim=-1)
 
-        if self.use_cudnn_indexer:
-            indices = cudnn_fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk)
-        else:
-            indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale)
+        indices = fp8_indexer(q_idx, k_idx, w, ks, ke, index_topk, self.weight_scale)
         # indices shape: [S_local, topk] in K's coordinate space (sentinel = s_full)
         # KV passed to sparse MLA has length s_full + 1 (sentinel zeros at index s_full).
         return indices.view(1, s_local, 1, index_topk)
diff --git a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py b/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
deleted file mode 100644
index adc25762b25..00000000000
--- a/src/prime_rl/trainer/models/kernels/cudnn_indexer.py
+++ /dev/null
@@ -1,114 +0,0 @@
-"""cuDNN frontend FP8 indexer for DSA sparse token selection (SM90).
-
-Same contract and scoring formula as `fp8_indexer`, I_{t,s} = Σ_j w_{t,j} · ReLU(q_{t,j} · k_s),
-with the same UE8M0 per-token FP8 quantization, but scored by the cuDNN SM90 indexer forward and
-selected by the cuDNN radix top-k. Packed documents become THD segments, so each document is scored
-only against its own keys.
-"""
-
-import torch
-
-from prime_rl.trainer.models.kernels.fp8_indexer import per_token_group_quant_fp8
-
-# Upper bound on fp32 logits per top-k launch: the kernel allocates a scratch buffer of twice the
-# logits' size, which would be 16 GiB for a 16k x 131k CP shard in one launch.
-TOPK_CHUNK_ELEMS = 1 << 27
-
-
-def _document_segments(ks: torch.Tensor, ke: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
-    """THD segments of the queries, one per document (run of equal `ks`).
-
-    Within a document `ke` advances by one per query, so the segment's keys are `[ks, ke[-1])` and
-    its first query sits at doc-local position `ke[0] - ks - 1`. Consecutive documents' key ranges
-    abut, so the key segments tile the gathered K from the first document's start.
-    """
-    S_q = ks.shape[0]
-    is_start = torch.ones(S_q, dtype=torch.bool, device=ks.device)
-    is_start[1:] = ks[1:] != ks[:-1]
-    starts = is_start.nonzero().squeeze(1).to(torch.int32)
-    cu_seqlens_q = torch.cat([starts, starts.new_tensor([S_q])])
-    cu_seqlens_k = torch.cat([ks[starts], ke[-1:]])
-    q_causal_offsets = ke[starts] - ks[starts] - 1
-    max_seqlen_q, max_seqlen_k = torch.stack([cu_seqlens_q.diff().max(), (ke - ks).max()]).tolist()
-    return cu_seqlens_q, cu_seqlens_k, q_causal_offsets, max_seqlen_q, max_seqlen_k
-
-
-def _chunked_top_k(logits: torch.Tensor, lengths: torch.Tensor, topk: int) -> torch.Tensor:
-    from cudnn.deepseek_sparse_attention import indexer_top_k_wrapper
-
-    num_rows, num_cols = logits.shape
-    chunk_rows = max(256, TOPK_CHUNK_ELEMS // num_cols)
-    return torch.cat(
-        [
-            indexer_top_k_wrapper(logits[lo : lo + chunk_rows], lengths[lo : lo + chunk_rows], topk, return_val=False)[
-                "indices"
-            ]
-            for lo in range(0, num_rows, chunk_rows)
-        ]
-    )
-
-
-@torch.library.custom_op("prime_rl::cudnn_fp8_indexer", mutates_args=())
-def cudnn_fp8_indexer(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    w: torch.Tensor,
-    ks: torch.Tensor,
-    ke: torch.Tensor,
-    topk: int,
-) -> torch.Tensor:
-    """cuDNN FP8 indexer: UE8M0 quantization + cuDNN SM90 scoring + cuDNN radix top-k.
-
-    Args:
-        q: [S_q, H, D] bf16 query vectors per head (H in {32, 64}, D = 128)
-        k: [S_k, D] bf16 key vectors (shared across heads)
-        w: [S_q, H] bf16 per-head weights
-        ks: [S_q] int32 document start per query (in K's coordinate system)
-        ke: [S_q] int32 causal end per query (= global position + 1); must advance by one per
-            query within a document
-        topk: number of top indices to return (<= 2048)
-
-    Returns:
-        [S_q, topk] int32 selected token indices per query (sentinel = S_k), in no particular order
-    """
-    from cudnn.deepseek_sparse_attention.indexer_forward._interface_sm90 import indexer_fwd
-
-    S_q, H, D = q.shape
-    S_k = k.shape[0]
-
-    q_fp8, q_scales = per_token_group_quant_fp8(q.reshape(S_q * H, D).contiguous(), group_size=D, use_ue8m0=True)
-    k_fp8, k_scales = per_token_group_quant_fp8(k.contiguous(), group_size=D, use_ue8m0=True)
-    q_scales = q_scales.view(S_q, H)
-    # Fold q's descale into the weights (fp32), as fp8_indexer does; K's descale goes to the kernel.
-    w = w * q_scales
-
-    cu_seqlens_q, cu_seqlens_k, q_causal_offsets, max_seqlen_q, max_seqlen_k = _document_segments(ks, ke)
-    # Rows 32-byte aligned for the top-k kernel's vector loads.
-    num_cols = (max_seqlen_k + 7) // 8 * 8
-
-    # ratio=1 is the plain token-level causal mask; scores are doc-local: column j is key ks + j.
-    logits = indexer_fwd(
-        q_fp8.view(S_q, H, D),
-        k_fp8.view(S_k, 1, D),
-        w,
-        ratio=1,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=num_cols,
-        q_causal_offsets=q_causal_offsets,
-        precision="fp8",
-        q_scale=q_scales,
-        k_scale=k_scales,
-    )
-
-    lengths = ke - ks
-    local = _chunked_top_k(logits, lengths, topk)
-    # Slots past a short row's length come back as -1.
-    valid = (local >= 0) & (local < lengths[:, None])
-    return torch.where(valid, local + ks[:, None], S_k).to(torch.int32)
-
-
-@cudnn_fp8_indexer.register_fake
-def _cudnn_fp8_indexer_fake(q, k, w, ks, ke, topk):
-    return q.new_empty((q.shape[0], topk), dtype=torch.int32)
diff --git a/tests/unit/train/models/test_fp8_indexer.py b/tests/unit/train/models/test_fp8_indexer.py
index c08ba8f87bb..a1f210b113f 100644
--- a/tests/unit/train/models/test_fp8_indexer.py
+++ b/tests/unit/train/models/test_fp8_indexer.py
@@ -1,7 +1,6 @@
 import pytest
 import torch
 
-from prime_rl.trainer.models.kernels.cudnn_indexer import cudnn_fp8_indexer
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 
 pytestmark = [
@@ -110,39 +109,3 @@ def test_selection_agreement(segments):
     mean, p1 = agreement.mean().item(), agreement.quantile(0.01).item()
     assert mean > AGREEMENT_MEAN, f"mean set agreement {mean} below {AGREEMENT_MEAN}"
     assert p1 > AGREEMENT_P1, f"p1 set agreement {p1} below {AGREEMENT_P1}"
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9,
-    reason="the cuDNN indexer kernels are SM90-only",
-)
-def test_cudnn_matches_triton_on_causal_documents():
-    """The cuDNN indexer against the Triton one on a CP-style shard of packed causal documents.
-
-    The queries are the last 2560 tokens of a 4096-token row, so the shard starts mid-document and
-    `ks` / `ke` live in the gathered K's coordinates. Both kernels quantize identically and their
-    scores agree to fp32 rounding, so the selections only differ by ties (e.g. keys every head's
-    ReLU zeroes) at the top-k boundary.
-    """
-    num_q, num_k, doc_lens = 2560, 4096, (1500, 1300, 200, 1096)
-    doc_start = torch.cat([torch.full((n,), sum(doc_lens[:i])) for i, n in enumerate(doc_lens)]).cuda()
-    ks = doc_start[num_k - num_q :].int().contiguous()
-    ke = torch.arange(num_k - num_q + 1, num_k + 1, device="cuda", dtype=torch.int32)
-
-    q = torch.randn(num_q, HEADS, DIM, device="cuda").bfloat16()
-    k = torch.randn(num_k, DIM, device="cuda").bfloat16()
-    w = torch.randn(num_q, HEADS, device="cuda").bfloat16()
-
-    def selected(picks):
-        mask = torch.zeros(num_q, num_k + 1, dtype=torch.bool, device="cuda")
-        return mask.scatter_(1, picks.long(), True)[:, :num_k]
-
-    triton_selected = selected(fp8_indexer(q, k, w, ks, ke, TOPK))
-    cudnn_selected = selected(cudnn_fp8_indexer(q, k, w, ks, ke, TOPK))
-
-    assert torch.equal(cudnn_selected.sum(-1), (ke - ks).clamp(max=TOPK)), "wrong number of readable picks"
-    saturated = (ke - ks) > TOPK
-    assert torch.equal(cudnn_selected[~saturated], triton_selected[~saturated])
-    assert saturated.any(), "vacuous probe: no query has more readable keys than it can pick"
-    agreement = (cudnn_selected & triton_selected).sum(-1)[saturated] / TOPK
-    assert agreement.mean().item() > 0.999 and agreement.min().item() > 0.99

From 0894aa8897544e3c2913b6a4274a3f9f7c6c9527 Mon Sep 17 00:00:00 2001
From: sami jaghouar <sami@primeintellect.ai>
Date: Mon, 5 Oct 2026 04:44:34 +0000
Subject: [PATCH 11/11] refactor(dsa): DSABackend Literal, required backend
 arg, if/else dispatch without asserts

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../src/prime_rl/configs/trainer.py           |  3 +-
 .../glm_moe_dsa/sparse_mla_attention.py       |  5 +-
 .../trainer/models/kernels/sparse_mla_bwd.py  | 23 +++----
 .../trainer/models/kernels/sparse_mla_fwd.py  | 60 +++++++++----------
 4 files changed, 46 insertions(+), 45 deletions(-)

diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
index 3067560151a..b728964d11f 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
@@ -22,6 +22,7 @@
 # -- Shared trainer configs (used by both SFT and RL trainers) --
 
 AttnImplementation: TypeAlias = Literal["flash_attention_2", "flash_attention_3", "flash_attention_4", "auto"]
+DSABackend: TypeAlias = Literal["tilelang", "cudnn_flashmla"]
 
 
 class GCConfig(BaseConfig):
@@ -325,7 +326,7 @@ class ModelConfig(BaseModelConfig):
 
     quantization: QuantizationConfig | None = None
 
-    dsa_backend: Literal["tilelang", "cudnn_flashmla"] = "tilelang"
+    dsa_backend: DSABackend = "tilelang"
     """Kernels for the DSA sparse MLA attention (``glm_moe_dsa``, e.g. GLM-5). ``tilelang`` runs the TileLang forward and backward. ``cudnn_flashmla`` runs the FlashMLA sparse prefill forward with the cuDNN frontend CuTe-DSL backward; it needs the ``flash-mla`` extra and an SM90 GPU (SM100/SM103 are wired up but untested)."""
 
     index_cache: IndexCacheConfig | None = None
diff --git a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
index 9e16c7c59dd..30ceea86b84 100644
--- a/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
+++ b/src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py
@@ -4,6 +4,7 @@
 import torch.distributed as dist
 from torch import nn
 
+from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.fp8_indexer import fp8_indexer
 from prime_rl.trainer.models.layers.norms import LayerNorm, RMSNorm, RMSNormConfig
 from prime_rl.trainer.models.layers.rotary_emb import rotate_half
@@ -30,9 +31,9 @@ class SparseMlaAttentionArgs:
     index_n_heads: int
     index_head_dim: int
     index_topk: int
+    dsa_backend: DSABackend
     use_index_cache: bool = False
     skip_topk: bool = False
-    dsa_backend: str = "tilelang"
 
 
 def apply_rope_interleave_single(
@@ -231,7 +232,7 @@ def forward(
             position_embeddings_full=position_embeddings,
         )
 
-        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.scaling, backend=self.args.dsa_backend)
+        out, _ = sparse_mla(sparse_q, sparse_kv, indices, self.args.dsa_backend, self.scaling)
         out = torch.einsum("bshk,hdk->bshd", out, w_v)
         batch_size, total_tokens = out.shape[:2]
         out = out.reshape(batch_size, total_tokens, -1)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
index 08e51ef663c..00d38c83c4a 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -19,6 +19,8 @@
 import torch
 from tilelang import language as T
 
+from prime_rl.configs.trainer import DSABackend
+
 
 @tilelang.jit(out_idx=[-1])
 def preprocess(
@@ -278,8 +280,6 @@ def sparse_mla_bwd_kernel(
     return sparse_mla_bwd_kernel
 
 
-DSA_BACKENDS = ("tilelang", "cudnn_flashmla")
-
 _CUDNN_SM100_CAPABILITIES = ((10, 0), (10, 3))
 
 
@@ -372,7 +372,11 @@ def cudnn_sparse_mla_backward(
     return dq.view_as(q), dkv.view_as(kv)
 
 
-@torch.library.custom_op("prime_rl::sparse_mla_backward", mutates_args=())
+@torch.library.custom_op(
+    "prime_rl::sparse_mla_backward",
+    mutates_args=(),
+    schema="(Tensor q, Tensor kv, Tensor out, Tensor grad_out, Tensor indices, Tensor lse, str backend, float? sm_scale=None) -> (Tensor, Tensor)",
+)
 def sparse_mla_backward(
     q: torch.Tensor,
     kv: torch.Tensor,
@@ -380,8 +384,8 @@ def sparse_mla_backward(
     grad_out: torch.Tensor,
     indices: torch.Tensor,
     lse: torch.Tensor,
+    backend: DSABackend,
     sm_scale: float | None = None,
-    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Sparse MLA backward for the `backend` route (see `sparse_mla`); `lse` is natural-log."""
     assert q.is_contiguous()
@@ -395,13 +399,12 @@ def sparse_mla_backward(
     topk = indices.shape[-1]
     assert indices.shape == (B, S, kv_group, topk)
     assert lse.shape == (B, S, H)
-    assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"
 
     if backend == "cudnn_flashmla":
-        arch = cudnn_backward_arch(q, kv)
-        assert arch is not None, "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
-        return cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, arch)
-    return tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+        dq, dkv = cudnn_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale, cudnn_backward_arch(q, kv))
+    else:
+        dq, dkv = tilelang_sparse_mla_backward(q, kv, out, grad_out, indices, lse, sm_scale)
+    return dq, dkv
 
 
 @sparse_mla_backward.register_fake
@@ -412,7 +415,7 @@ def _sparse_mla_backward_fake(
     grad_out: torch.Tensor,
     indices: torch.Tensor,
     lse: torch.Tensor,
+    backend: DSABackend,
     sm_scale: float | None = None,
-    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), torch.empty_like(kv)
diff --git a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
index cfb58852be5..3bf6aa70e0c 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -13,16 +13,14 @@
     # This is expected on CPU-only machines
     pass
 
-import importlib.util
 import math
 
 import tilelang
 import torch
 from tilelang import language as T
 
+from prime_rl.configs.trainer import DSABackend
 from prime_rl.trainer.models.kernels.sparse_mla_bwd import (
-    DSA_BACKENDS,
-    cudnn_backward_arch,
     flat_kv_indices,
     sparse_mla_backward,
 )
@@ -218,17 +216,21 @@ def flashmla_sparse_mla_forward(
     return out.view(B, S, H, d_v), lse.view(B, S, H)
 
 
-@torch.library.custom_op("prime_rl::sparse_mla", mutates_args=())
+@torch.library.custom_op(
+    "prime_rl::sparse_mla",
+    mutates_args=(),
+    schema="(Tensor q, Tensor kv, Tensor indices, str backend, float? sm_scale=None, SymInt d_v=512, SymInt block_I=64, SymInt num_stages=2, SymInt threads=256) -> (Tensor, Tensor)",
+)
 def sparse_mla(
     q: torch.Tensor,
     kv: torch.Tensor,
     indices: torch.Tensor,
+    backend: DSABackend,
     sm_scale: float | None = None,
     d_v: int = 512,
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
-    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Sparse MLA attention over the top-k `indices`; returns the output and its natural-log LSE.
 
@@ -237,7 +239,6 @@ def sparse_mla(
     (SM90; SM100 / SM103 untested).
     """
     assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
-    assert backend in DSA_BACKENDS, f"Unknown DSA backend: {backend}"
     batch, seq_len, heads, dim_plus_tail_dim = q.shape
     _, seq_len_kv, kv_group, _ = kv.shape
 
@@ -250,31 +251,26 @@ def sparse_mla(
     assert indices.shape == (batch, seq_len, kv_group, topk)
 
     if backend == "cudnn_flashmla":
-        assert cudnn_backward_arch(q, kv) is not None, (
-            "dsa_backend='cudnn_flashmla' requires an SM90/SM100/SM103 GPU and a single KV head"
-        )
-        assert importlib.util.find_spec("flash_mla") is not None, (
-            "dsa_backend='cudnn_flashmla' requires FlashMLA; install it with `uv sync --extra flash-mla`"
-        )
         if sm_scale is None:
             sm_scale = dim_plus_tail_dim**-0.5
-        return flashmla_sparse_mla_forward(q, kv, indices, sm_scale, d_v)
-
-    kernel = sparse_mla_fwd(
-        heads,
-        dim,
-        tail_dim,
-        topk,
-        kv_group,
-        sm_scale,
-        True,
-        block_I=block_I,
-        num_stages=num_stages,
-        threads=threads,
-    )
-    out, lse = kernel(q, kv, indices)
-    # The TileLang kernel returns a base-2 LSE.
-    return out, lse * math.log(2.0)
+        out, lse = flashmla_sparse_mla_forward(q, kv, indices, sm_scale, d_v)
+    else:
+        kernel = sparse_mla_fwd(
+            heads,
+            dim,
+            tail_dim,
+            topk,
+            kv_group,
+            sm_scale,
+            True,
+            block_I=block_I,
+            num_stages=num_stages,
+            threads=threads,
+        )
+        out, lse = kernel(q, kv, indices)
+        # The TileLang kernel returns a base-2 LSE.
+        lse = lse * math.log(2.0)
+    return out, lse
 
 
 @sparse_mla.register_fake
@@ -282,18 +278,18 @@ def _sparse_mla_fake(
     q: torch.Tensor,
     kv: torch.Tensor,
     indices: torch.Tensor,
+    backend: DSABackend,
     sm_scale: float | None = None,
     d_v: int = 512,
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
-    backend: str = "tilelang",
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return q.new_empty((*q.shape[:-1], d_v)), q.new_empty(q.shape[:-1], dtype=torch.float32)
 
 
 def _sparse_mla_setup_context(ctx, inputs, output) -> None:
-    q, kv, indices, sm_scale, _d_v, _block_I, _num_stages, _threads, backend = inputs
+    q, kv, indices, backend, sm_scale, _d_v, _block_I, _num_stages, _threads = inputs
     out, lse = output
     ctx.save_for_backward(q, kv, out, indices, lse)
     ctx.sm_scale = sm_scale
@@ -310,8 +306,8 @@ def _sparse_mla_autograd_backward(ctx, grad_out: torch.Tensor, _grad_lse: torch.
         grad_out,
         indices,
         lse.detach(),
-        ctx.sm_scale,
         ctx.backend,
+        ctx.sm_scale,
     )
     return dq, dkv, None, None, None, None, None, None, None