perf(dsa): cuDNN + FlashMLA kernels for GLM-5 sparse attention (model.dsa_backend)

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.

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

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

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,6 +326,9 @@ class ModelConfig(BaseModelConfig):
 
     quantization: QuantizationConfig | None = None
 
+    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
     """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."""
 

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'",
@@ -94,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 = [
@@ -107,6 +112,7 @@ all = [
     "prime-rl[sonic-moe]",
     "prime-rl[flash-attn-3]",
     "prime-rl[flash-attn-cute]",
+    "prime-rl[flash-mla]",
 ]
 
 [dependency-groups]
@@ -246,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 }
@@ -286,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" }

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

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),
+        dsa_backend=getattr(config, "dsa_backend", "tilelang"),
     )
 
 

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,6 +31,7 @@ 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
 
@@ -230,7 +232,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.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)

src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py

@@ -13,10 +13,14 @@
     # This is expected on CPU-only machines
     pass
 
+import math
+
 import tilelang
 import torch
 from tilelang import language as T
 
+from prime_rl.configs.trainer import DSABackend
+
 
 @tilelang.jit(out_idx=[-1])
 def preprocess(
@@ -276,8 +280,32 @@ 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(
+_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,
     out: torch.Tensor,
@@ -286,30 +314,96 @@ 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()
+    """TileLang backward; takes a natural-log LSE."""
     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)
     postprocess_kernel = postprocess(D, D_tail, kv_group)
 
     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
+
+
+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,
+    arch: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """cuDNN-frontend CuTe-DSL DSA backward (single KV head with K=V).
+
+    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.
+    """
+    B, S, H, D_qk = q.shape
+    S_kv = kv.shape[1]
+    topk = indices.shape[-1]
+    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.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)
+
+
+@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,
+    out: torch.Tensor,
+    grad_out: torch.Tensor,
+    indices: torch.Tensor,
+    lse: torch.Tensor,
+    backend: DSABackend,
+    sm_scale: float | None = None,
+) -> 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()
+    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 backend == "cudnn_flashmla":
+        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
 
 
@@ -321,6 +415,7 @@ def _sparse_mla_backward_fake(
     grad_out: torch.Tensor,
     indices: torch.Tensor,
     lse: torch.Tensor,
+    backend: DSABackend,
     sm_scale: float | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), torch.empty_like(kv)