feat(deepseek-v4): cudnn_flashmla sparse attention backend (FlashMLA forward, cuDNN backward)

configs/advanced/deepseek-v4-flash/sft.toml

@@ -19,6 +19,7 @@ max_steps = 20
 [model]
 name = "deepseek-ai/DeepSeek-V4-Flash-0731"
 ep = 8
+dsa_backend = "cudnn_flashmla"
 
 # Torch all-to-all, not DeepEP. At full depth DeepEP dies on the first backward with
 # "RuntimeError: Autograd not support dtype: Byte", raised from `Tensor.float()` in

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -288,7 +288,7 @@ 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)."""
+    """Kernels for the sparse attention of ``glm_moe_dsa`` (DSA sparse MLA, e.g. GLM-5) and ``deepseek_v4``. ``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, SM100 or SM103 GPU."""
 
     freeze_moe_router: bool | Literal["auto"] = "auto"
     """Freeze MoE router gate parameters. ``auto`` (default): ``True`` for RL, ``False`` for SFT; a no-op on non-MoE models. Frozen routers get no optimizer state, so resuming a checkpoint saved with a frozen router with ``freeze_moe_router=false`` fails on missing DCP keys."""

src/prime_rl/trainer/models/deepseek_v4/attention.py

@@ -140,19 +140,31 @@
 
 # Guarded because tilelang ships in the linux-gated `gpu` extra, so some installs lack it.
 try:
-    from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_sparse_attn import dsv4_sparse_attn, sparse_attn_shape_error
+    from prime_rl.trainer.models.kernels.deepseek_v4.dsv4_sparse_attn import (
+        FORWARD_BACKENDS,
+        dsv4_sparse_attn,
+        sparse_attn_shape_error,
+    )
 except ImportError:
+    FORWARD_BACKENDS = {}
     dsv4_sparse_attn = None  # type: ignore
     sparse_attn_shape_error = None  # type: ignore
 
 
-def _kernel_blocker(num_heads: int, head_dim: int) -> str | None:
-    """Why the fused kernel cannot run at this shape, or ``None`` if it can."""
+def _kernel_blocker(num_heads: int, head_dim: int, backend: str) -> str | None:
+    """Why `backend`'s fused kernel cannot run at this shape, or ``None`` if it can."""
     if dsv4_sparse_attn is None:
         return "the tilelang sparse-attention kernel failed to import; install the `gpu` extra"
     # CSA gives every query head the same single KV head, so the kernel's `kv_group` is 1. The
     # shape constraints themselves are stated once, next to the kernels they come from.
-    return sparse_attn_shape_error(num_heads, 1, head_dim)
+    return sparse_attn_shape_error(num_heads, 1, head_dim, backend)
+
+
+def max_sparse_attn_slots(config: DeepseekV4Config) -> int:
+    """The widest slot axis any layer of `config` can hand the kernel, at the longest document it allows."""
+    hca_rate = config.compress_rates.get("heavily_compressed_attention")
+    hca_picks = config.max_position_embeddings // hca_rate if hca_rate is not None else 0
+    return config.sliding_window + max(config.index_topk, hca_picks)
 
 
 class DeepseekV4GroupedLinear(nn.Linear):
@@ -696,7 +708,14 @@ def __init__(self, config: DeepseekV4Config, layer_idx: int, rotary_emb: Deepsee
         self.rotary_emb = rotary_emb
         # Raised here rather than from the first forward, where it would surface as an ImportError
         # or a tilelang compile failure a long way from the config that caused it.
-        blocker = _kernel_blocker(self.num_heads, self.head_dim)
+        self.dsa_backend = getattr(config, "dsa_backend", "tilelang")
+        if dsv4_sparse_attn is not None and self.dsa_backend not in FORWARD_BACKENDS:
+            raise ValueError(
+                f"DeepSeek V4 has no {self.dsa_backend!r} sparse-attention backend, "
+                f"expected one of {sorted(FORWARD_BACKENDS)}"
+            )
+        self.max_slots = max_sparse_attn_slots(config)
+        blocker = _kernel_blocker(self.num_heads, self.head_dim, self.dsa_backend)
         if blocker is not None:
             raise ValueError(f"DeepSeek V4 cannot run the fused sparse-attention kernel: {blocker}")
         assert config.attention_dropout == 0.0, "the fused sparse attention kernel implements no dropout"
@@ -769,6 +788,8 @@ def forward(self, hidden_states: torch.Tensor, packed: PackedContext) -> tuple[t
             inputs.indices,
             self.sinks,
             self.scaling,
+            backend=self.dsa_backend,
+            max_slots=self.max_slots,
         )  # (b, t, h, d)
 
         # The value stream is the key stream, so it arrived rotated. Rotating the output

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

@@ -0,0 +1,162 @@
+"""Compile the cuDNN DSA backward kernel once, instead of again whenever sequence length or packing changes.
+
+Out of the box, cuDNN compiles a new kernel for every number of key slots per query (`topk_idxs.shape[-1]`), and
+that number changes with sequence length and document packing. The kernel only uses it as an upper limit, so
+when a call passes `topk_length` we compile once with a fixed limit that covers every call. Calls without
+`topk_length` are left alone. On SM100 this applies to the `generic_m64` kernel only, and the number of key slots
+must be a multiple of 64. This patches private cuDNN frontend internals and raises if they change.
+"""
+
+import math
+import time
+from collections.abc import Callable
+
+import cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 as interface_sm90
+import cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 as interface_sm100
+import cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm90 as kernel_sm90
+import cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100 as kernel_sm100
+import cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100_deterministic as kernel_sm100_deterministic
+import torch
+
+from prime_rl.utils.logger import get_logger
+
+_topk_clamp: int | None = None
+_stock_init = kernel_sm90.FlashAttentionDSABackwardSm90.__init__
+_sm100_topk_clamp: int | None = None
+_sm100_stock_init = kernel_sm100.FlashAttentionDSABackwardSm100.__init__
+_SM100_GENERIC_KERNELS = (
+    kernel_sm100.FlashAttentionDSABackwardSm100,
+    kernel_sm100_deterministic.FlashAttentionDSABackwardSm100Deterministic,
+)
+
+
+def topk_clamp(max_slots: int) -> int:
+    """The `max_topk` the main kernel is built with: the smallest power of two that is at least `max_slots`.
+
+    A power of two compiles to nearly the stock kernel's SASS, where `2**31 - 1` costs 1% of forward+backward.
+    """
+    exponent = math.ceil(math.log2(max_slots))
+    return 2**exponent
+
+
+def main_kernel_key(key: tuple) -> tuple:
+    *static_key, have_topk_length, max_topk = key
+    if not isinstance(have_topk_length, bool) or not isinstance(max_topk, int):
+        raise RuntimeError(f"the cuDNN DSA main-kernel key no longer ends in (have_topk_length, max_topk): {key}")
+    if not have_topk_length:
+        return key
+    if _topk_clamp is None:
+        raise RuntimeError("a cuDNN DSA backward call with topk_length must go through cudnn_dsa_patch")
+    if max_topk > _topk_clamp:
+        raise ValueError(f"slot width {max_topk} exceeds the cuDNN DSA backward's compiled clamp {_topk_clamp}")
+    return (*static_key, have_topk_length, _topk_clamp)
+
+
+def _init_with_clamped_max_topk(self, *args, have_topk_length: bool, max_topk: int, **kwargs) -> None:
+    _stock_init(
+        self, *args, have_topk_length=have_topk_length, max_topk=_topk_clamp if have_topk_length else max_topk, **kwargs
+    )
+
+
+class ExecutorCache(dict):
+    """An executor cache that stores entries under `normalize_key(key)` and logs each compile.
+
+    The frontend checks `key in cache`, compiles on a miss, then assigns the executor, so the time between
+    the missed lookup and the assignment is the compile.
+    """
+
+    def __init__(
+        self,
+        kernel: str,
+        normalize_key: Callable[[tuple], tuple] = lambda key: key,
+        max_topk_position: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.kernel = kernel
+        self.normalize_key = normalize_key
+        self.max_topk_position = max_topk_position
+        self.miss_time = 0.0
+
+    def __contains__(self, key) -> bool:
+        found = dict.__contains__(self, self.normalize_key(key))
+        if not found:
+            self.miss_time = time.perf_counter()
+        return found
+
+    def __getitem__(self, key):
+        return dict.__getitem__(self, self.normalize_key(key))
+
+    def __setitem__(self, key, compiled) -> None:
+        normalized_key = self.normalize_key(key)
+        dict.__setitem__(self, normalized_key, compiled)
+        seconds = time.perf_counter() - self.miss_time
+        max_topk = "" if self.max_topk_position is None else f" with max_topk {normalized_key[self.max_topk_position]}"
+        get_logger().info(f"Compiled the cuDNN DSA backward {self.kernel} kernel{max_topk} in {seconds:.2f}s")
+
+
+def patch_sm90_compile_cache() -> None:
+    if kernel_sm90.FlashAttentionDSABackwardSm90.__init__ is _init_with_clamped_max_topk:
+        return
+    kernel_sm90.FlashAttentionDSABackwardSm90.__init__ = _init_with_clamped_max_topk
+    bwd = interface_sm90.flash_attn_bwd_sm90
+    bwd.compile_cache_pre = ExecutorCache("preprocess")
+    bwd.compile_cache = ExecutorCache("main", main_kernel_key, max_topk_position=-1)
+    bwd.compile_cache_post = ExecutorCache("postprocess")
+
+
+def flash_attn_bwd_sm90(*args, max_slots: int | None = None, **kwargs) -> tuple[torch.Tensor, ...]:
+    """The frontend's `flash_attn_bwd_sm90`; a call with `topk_length` serves widths up to `max_slots`."""
+    global _topk_clamp
+    patch_sm90_compile_cache()
+    if kwargs.get("topk_length") is not None:
+        if max_slots is None:
+            raise ValueError("a cuDNN DSA backward call with topk_length needs max_slots, its widest slot width")
+        _topk_clamp = topk_clamp(max_slots)
+    return interface_sm90.flash_attn_bwd_sm90(*args, **kwargs)
+
+
+def sm100_main_kernel_key(key: tuple) -> tuple:
+    backend, *static_key, max_topk, has_topk_length = key
+    if not isinstance(has_topk_length, bool) or not isinstance(max_topk, int):
+        raise RuntimeError(f"the cuDNN SM100 DSA key no longer ends in (max_topk, has_topk_length): {key}")
+    if backend != "generic_m64" or not has_topk_length:
+        return key
+    if _sm100_topk_clamp is None:
+        raise RuntimeError("a cuDNN SM100 DSA backward call with topk_length must go through cudnn_dsa_patch")
+    if max_topk > _sm100_topk_clamp:
+        raise ValueError(f"slot width {max_topk} exceeds the cuDNN DSA backward's compiled clamp {_sm100_topk_clamp}")
+    if max_topk % 64 != 0:
+        raise ValueError(
+            f"the clamped cuDNN SM100 DSA backward needs a slot width that is a multiple of 64, got {max_topk}"
+        )
+    return (backend, *static_key, _sm100_topk_clamp, has_topk_length)
+
+
+def _sm100_init_with_clamped_max_topk(self, element_dtype, head_dim, head_dim_v, block_tile, max_topk=0) -> None:
+    if _sm100_topk_clamp is not None and type(self) in _SM100_GENERIC_KERNELS:
+        max_topk = _sm100_topk_clamp
+    _sm100_stock_init(self, element_dtype, head_dim, head_dim_v, block_tile, max_topk)
+
+
+def patch_sm100_compile_cache() -> None:
+    if kernel_sm100.FlashAttentionDSABackwardSm100.__init__ is _sm100_init_with_clamped_max_topk:
+        return
+    kernel_sm100.FlashAttentionDSABackwardSm100.__init__ = _sm100_init_with_clamped_max_topk
+    interface_sm100.flash_attn_bwd_sm100.compile_cache = ExecutorCache(
+        "main", sm100_main_kernel_key, max_topk_position=-2
+    )
+
+
+def flash_attn_bwd_sm100(*args, max_slots: int | None = None, **kwargs) -> tuple[torch.Tensor, ...]:
+    """The frontend's `flash_attn_bwd_sm100`; a call with `topk_length` serves widths up to `max_slots`."""
+    global _sm100_topk_clamp
+    patch_sm100_compile_cache()
+    if kwargs.get("topk_length") is None:
+        return interface_sm100.flash_attn_bwd_sm100(*args, **kwargs)
+    if max_slots is None:
+        raise ValueError("a cuDNN DSA backward call with topk_length needs max_slots, its widest slot width")
+    _sm100_topk_clamp = topk_clamp(max_slots)
+    try:
+        return interface_sm100.flash_attn_bwd_sm100(*args, **kwargs)
+    finally:
+        _sm100_topk_clamp = None