diff --git a/configs/advanced/deepseek-v4-flash/sft.toml b/configs/advanced/deepseek-v4-flash/sft.toml
index fc184af924b..c9e56d93e68 100644
--- a/configs/advanced/deepseek-v4-flash/sft.toml
+++ b/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
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 952b85dadf7..24ecd3dd329 100644
--- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py
+++ b/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."""
diff --git a/src/prime_rl/trainer/models/deepseek_v4/attention.py b/src/prime_rl/trainer/models/deepseek_v4/attention.py
index 9dfd91deeb3..f2711062935 100644
--- a/src/prime_rl/trainer/models/deepseek_v4/attention.py
+++ b/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
diff --git a/src/prime_rl/trainer/models/kernels/cudnn_dsa_patch.py b/src/prime_rl/trainer/models/kernels/cudnn_dsa_patch.py
new file mode 100644
index 00000000000..165cf8cd10d
--- /dev/null
+++ b/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
diff --git a/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py b/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
index 4bfbd334c77..4c319740d08 100644
--- a/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
+++ b/src/prime_rl/trainer/models/kernels/deepseek_v4/dsv4_sparse_attn.py
@@ -2,7 +2,8 @@
 # and `prime_rl::dsv4_sparse_attn_backward` custom ops, their fake (meta) implementations, and the
 # autograd rule that ties them together and forms the attention-sink gradient in torch. The
 # TileLang kernels themselves live in `dsv4_sparse_attn_fwd.py` and `dsv4_sparse_attn_bwd.py`; this
-# module is the only place that knows both exist.
+# module is the only place that knows both exist. The `cudnn_flashmla` backend calls FlashMLA's
+# sparse prefill forward and the cuDNN frontend's DSA backward from their wheels.
 
 # TileLang ships a libcudart stub that proxies to the real CUDA runtime via
 # dlsym(RTLD_DEFAULT, ...).  If the stub's own symbols are the first ones found
@@ -10,6 +11,7 @@
 # and the stub calls abort().  Pre-loading the real library with RTLD_GLOBAL
 # ensures dlsym finds it before the stub's own exports.
 import ctypes as _ctypes
+import math
 
 try:
     _ctypes.CDLL("libcudart.so", mode=_ctypes.RTLD_GLOBAL)
@@ -62,12 +64,7 @@ def num_tiles_covering_valid_slots(indices: torch.Tensor, tile_size: int) -> tor
     return last_valid_slot_idx // tile_size + 1
 
 
-def sparse_attn_shape_error(heads: int, kv_group: int, dim: int) -> str | None:
-    """The reason these kernels cannot serve this shape, or ``None`` if they can.
-
-    The forward, the backward and the constructor check in `deepseek_v4/attention.py` all need
-    the same answer, so the constraints live here only.
-    """
+def tilelang_shape_error(heads: int, kv_group: int, dim: int) -> str | None:
     # The backward's `preprocess` tiles the channel axis at `block_ND = 32` and reads whole
     # tiles, so a `dim` below that (or not a multiple of it) over-reads into the next head and
     # silently corrupts `Delta`, hence every gradient. This subsumes `atomic_addx4`'s own
@@ -96,32 +93,40 @@ def sparse_attn_shape_error(heads: int, kv_group: int, dim: int) -> str | None:
     return None
 
 
-@torch.library.custom_op("prime_rl::dsv4_sparse_attn", mutates_args=())
-def dsv4_sparse_attn(
+def cudnn_flashmla_shape_error(heads: int, kv_group: int, dim: int) -> str | None:
+    if kv_group != 1:
+        return f"FlashMLA and the cuDNN backward serve a single KV head, but this shape has {kv_group}"
+    if heads not in (64, 128):
+        return f"FlashMLA's sparse prefill serves 64 or 128 query heads, but this shape has {heads}"
+    if dim != 512:
+        return f"FlashMLA's sparse prefill and the cuDNN backward need head_dim 512 here, but it is {dim}"
+    if not torch.cuda.is_available():
+        return "the cudnn_flashmla backend needs a CUDA GPU"
+    capability = torch.cuda.get_device_capability()
+    if capability[0] != 9 and capability not in ((10, 0), (10, 3)):
+        return f"the cudnn_flashmla backend runs on SM90, SM100 and SM103, but this GPU is {capability}"
+    return None
+
+
+SHAPE_ERRORS = {"tilelang": tilelang_shape_error, "cudnn_flashmla": cudnn_flashmla_shape_error}
+
+
+def sparse_attn_shape_error(heads: int, kv_group: int, dim: int, backend: str = "tilelang") -> str | None:
+    return SHAPE_ERRORS[backend](heads, kv_group, dim)
+
+
+def tilelang_forward(
     q: torch.Tensor,
     kv: torch.Tensor,
     indices: torch.Tensor,
     sinks: torch.Tensor,
-    sm_scale: float | None = None,
-    block_I: int = 64,
-    num_stages: int = 2,
-    threads: int = 256,
+    sm_scale: float | None,
+    block_I: int,
+    num_stages: int,
+    threads: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    assert q.is_contiguous(), "q must be contiguous"
-    assert kv.is_contiguous(), "kv must be contiguous"
-    assert indices.is_contiguous(), "indices must be contiguous"
     batch, seq_len, heads, dim = q.shape
     _, _, kv_group, _ = kv.shape
-
-    assert kv.shape[-1] == dim, "q and kv must share the full channel dim; DS V4 has no score-only tail"
-    assert kv.shape[0] == batch
-    assert q.dtype == torch.bfloat16, (
-        f"the sparse attention kernel runs in bfloat16 only, but the queries are {q.dtype}"
-    )
-    shape_error = sparse_attn_shape_error(heads, kv_group, dim)
-    assert shape_error is None, shape_error
-    assert indices.shape[:3] == (batch, seq_len, kv_group)
-    assert sinks.shape == (heads,)
     assert SLOT_TILE % block_I == 0, (
         f"the slot axis is padded to a multiple of {SLOT_TILE}, so block_I must divide it, got {block_I}"
     )
@@ -143,6 +148,107 @@ def dsv4_sparse_attn(
     return out, lse
 
 
+# FlashMLA's SM90 sparse prefill reads slots two 64-wide tiles at a time, so it needs a multiple of 128, which
+# SM100 accepts too and which meets the clamped SM100 cuDNN backward's multiple of 64. The cuDNN backward
+# compiles one kernel for every width (`cudnn_dsa_patch`), so nothing pads further.
+CUDNN_FLASHMLA_SLOT_TILE = 128
+
+
+def cudnn_flashmla_slot_width(n_slots: int) -> int:
+    """The slot width the cudnn_flashmla kernels run `n_slots` slots at."""
+    return -(-n_slots // CUDNN_FLASHMLA_SLOT_TILE) * CUDNN_FLASHMLA_SLOT_TILE
+
+
+def flat_cudnn_flashmla_indices(indices: torch.Tensor, n_positions: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Tile-padded global indices into the batch-flattened KV, and each query's valid prefix length.
+
+    `indices` is `(batch, seq_len, 1, n_slots)` with batch-local positions. The results are
+    `(batch * seq_len, width)` and `(batch * seq_len,)`, both int32, with absent slots at -1.
+    """
+    batch, seq_len, _, n_slots = indices.shape
+    width = cudnn_flashmla_slot_width(n_slots)
+    padded = F.pad(indices, (0, width - n_slots), value=IGNORE_SLOT)
+    batch_offset = torch.arange(batch, device=indices.device, dtype=indices.dtype).view(batch, 1, 1, 1) * n_positions
+    is_valid_slot = padded >= 0
+    flat_indices = torch.where(is_valid_slot, padded + batch_offset, IGNORE_SLOT).view(batch * seq_len, width)
+    slot_numbers = torch.arange(1, width + 1, device=indices.device, dtype=torch.int32)
+    topk_length = torch.where(is_valid_slot, slot_numbers, 0).amax(dim=-1).view(batch * seq_len)
+    return flat_indices.contiguous(), topk_length
+
+
+def cudnn_flashmla_forward(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sinks: torch.Tensor,
+    sm_scale: float | None,
+    _block_I: int,
+    _num_stages: int,
+    _threads: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    from flash_mla import flash_mla_sparse_fwd
+
+    batch, seq_len, heads, dim = q.shape
+    n_positions = kv.shape[1]
+    if sm_scale is None:
+        sm_scale = dim**-0.5
+    flat_indices, topk_length = flat_cudnn_flashmla_indices(indices, n_positions)
+    sink_logits = sinks.float().contiguous()
+    out, _max_logits, sinkless_lse = flash_mla_sparse_fwd(
+        q.view(batch * seq_len, heads, dim),
+        kv.view(batch * n_positions, 1, dim),
+        flat_indices.unsqueeze(1),
+        sm_scale,
+        dim,
+        attn_sink=sink_logits,
+        topk_length=topk_length,
+    )
+    # FlashMLA's LSE is natural-log and leaves the sink out, and it is +inf for a query with no
+    # valid slot. The public LSE is TileLang's: base-2, sink included, the sink alone when empty.
+    lse = torch.where(torch.isfinite(sinkless_lse), torch.logaddexp(sinkless_lse, sink_logits), sink_logits) * LOG2E
+    return out.view(batch, seq_len, heads, dim), lse.view(batch, seq_len, heads)
+
+
+FORWARD_BACKENDS = {"tilelang": tilelang_forward, "cudnn_flashmla": cudnn_flashmla_forward}
+
+
+@torch.library.custom_op("prime_rl::dsv4_sparse_attn", mutates_args=())
+def dsv4_sparse_attn(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sinks: torch.Tensor,
+    sm_scale: float | None = None,
+    block_I: int = 64,
+    num_stages: int = 2,
+    threads: int = 256,
+    backend: str = "tilelang",
+    max_slots: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert backend in FORWARD_BACKENDS, (
+        f"unknown sparse attention backend {backend!r}, expected one of {sorted(FORWARD_BACKENDS)}"
+    )
+    assert q.is_contiguous(), "q must be contiguous"
+    assert kv.is_contiguous(), "kv must be contiguous"
+    assert indices.is_contiguous(), "indices must be contiguous"
+    batch, seq_len, heads, dim = q.shape
+    _, _, kv_group, _ = kv.shape
+
+    assert kv.shape[-1] == dim, "q and kv must share the full channel dim; DS V4 has no score-only tail"
+    assert kv.shape[0] == batch
+    assert q.dtype == torch.bfloat16, (
+        f"the sparse attention kernel runs in bfloat16 only, but the queries are {q.dtype}"
+    )
+    shape_error = sparse_attn_shape_error(heads, kv_group, dim, backend)
+    assert shape_error is None, shape_error
+    assert indices.shape[:3] == (batch, seq_len, kv_group)
+    assert sinks.shape == (heads,)
+    assert backend != "cudnn_flashmla" or max_slots is not None, (
+        "the cudnn_flashmla backward builds one kernel for every slot width up to max_slots, so it needs max_slots"
+    )
+    return FORWARD_BACKENDS[backend](q, kv, indices, sinks, sm_scale, block_I, num_stages, threads)
+
+
 # A fake must mirror the op's signature exactly, so it takes every argument even though only `q`,
 # the one argument that determines the output shapes, is read.
 @dsv4_sparse_attn.register_fake
@@ -155,34 +261,24 @@ def _dsv4_sparse_attn_fake(
     block_I: int = 64,
     num_stages: int = 2,
     threads: int = 256,
+    backend: str = "tilelang",
+    max_slots: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), q.new_empty(q.shape[:-1], dtype=torch.float32)
 
 
-@torch.library.custom_op("prime_rl::dsv4_sparse_attn_backward", mutates_args=())
-def dsv4_sparse_attn_backward(
+def tilelang_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,
+    sm_scale: float | None,
+    _max_slots: int | None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    assert q.is_contiguous(), "q must be contiguous"
-    assert kv.is_contiguous(), "kv must be contiguous"
-    assert indices.is_contiguous(), "indices must be contiguous"
-    assert lse.is_contiguous(), "lse must be contiguous"
-    grad_out = grad_out.contiguous()
     batch, seq_len, heads, dim = q.shape
     _, _, kv_group, _ = kv.shape
-    assert kv.shape[-1] == dim, "q and kv must share the full channel dim; DS V4 has no score-only tail"
-    assert kv.shape[0] == batch
-    # This op is public, so it repeats the forward's shape checks rather than trusting autograd.
-    shape_error = sparse_attn_shape_error(heads, kv_group, dim)
-    assert shape_error is None, shape_error
-    assert indices.shape[:3] == (batch, seq_len, kv_group)
-    assert lse.shape == (batch, seq_len, heads)
     indices = _pad_slots_to_tile(indices)
 
     preprocess_kernel = preprocess(heads, dim)
@@ -199,6 +295,97 @@ def dsv4_sparse_attn_backward(
     return dq, dkv, delta
 
 
+def cudnn_flashmla_backward(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    out: torch.Tensor,
+    grad_out: torch.Tensor,
+    indices: torch.Tensor,
+    lse: torch.Tensor,
+    sm_scale: float | None,
+    max_slots: int | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 import (
+        flash_attn_bwd_sm100_workspace_size,
+    )
+
+    from prime_rl.trainer.models.kernels.cudnn_dsa_patch import flash_attn_bwd_sm90, flash_attn_bwd_sm100
+
+    batch, seq_len, heads, dim = q.shape
+    n_positions = kv.shape[1]
+    flat_indices, topk_length = flat_cudnn_flashmla_indices(indices, n_positions)
+    flat_tensors = (
+        q.view(batch * seq_len, heads, dim),
+        kv.view(batch * n_positions, dim),
+        out.view(batch * seq_len, heads, dim),
+        grad_out.view(batch * seq_len, heads, dim),
+        # A sink-inclusive natural-log LSE with no sink makes cuDNN's `exp(s - lse)` exactly the
+        # sink-aware probability, where recovering the sinkless LSE would cancel catastrophically.
+        (lse * math.log(2.0)).view(batch * seq_len, heads),
+    )
+    if torch.cuda.get_device_capability(q.device)[0] == 9:
+        dq, dkv = flash_attn_bwd_sm90(
+            *flat_tensors,
+            softmax_scale=sm_scale,
+            topk_idxs=flat_indices,
+            topk_length=topk_length,
+            max_slots=cudnn_flashmla_slot_width(max_slots),
+        )
+    else:
+        # The SM100 kernel folds `attn_sink` into the LSE itself, and a -inf sink folds in nothing.
+        no_sink = torch.full((heads,), float("-inf"), dtype=torch.float32, device=q.device)
+        workspace_bytes = flash_attn_bwd_sm100_workspace_size(batch * seq_len, batch * n_positions, dim, heads)
+        dq, dkv, _ = flash_attn_bwd_sm100(
+            *flat_tensors,
+            no_sink,
+            flat_indices,
+            softmax_scale=sm_scale,
+            topk_length=topk_length,
+            workspace=torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device),
+            max_slots=cudnn_flashmla_slot_width(max_slots),
+        )
+    delta = preprocess(heads, dim)(out, grad_out)
+    return dq.view_as(q), dkv.view_as(kv), delta
+
+
+BACKWARD_BACKENDS = {"tilelang": tilelang_backward, "cudnn_flashmla": cudnn_flashmla_backward}
+
+
+@torch.library.custom_op("prime_rl::dsv4_sparse_attn_backward", mutates_args=())
+def dsv4_sparse_attn_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,
+    backend: str = "tilelang",
+    max_slots: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    assert backend in BACKWARD_BACKENDS, (
+        f"unknown sparse attention backend {backend!r}, expected one of {sorted(BACKWARD_BACKENDS)}"
+    )
+    assert q.is_contiguous(), "q must be contiguous"
+    assert kv.is_contiguous(), "kv must be contiguous"
+    assert indices.is_contiguous(), "indices must be contiguous"
+    assert lse.is_contiguous(), "lse must be contiguous"
+    grad_out = grad_out.contiguous()
+    batch, seq_len, heads, dim = q.shape
+    _, _, kv_group, _ = kv.shape
+    assert kv.shape[-1] == dim, "q and kv must share the full channel dim; DS V4 has no score-only tail"
+    assert kv.shape[0] == batch
+    # This op is public, so it repeats the forward's shape checks rather than trusting autograd.
+    shape_error = sparse_attn_shape_error(heads, kv_group, dim, backend)
+    assert shape_error is None, shape_error
+    assert indices.shape[:3] == (batch, seq_len, kv_group)
+    assert lse.shape == (batch, seq_len, heads)
+    assert backend != "cudnn_flashmla" or max_slots is not None, (
+        "the cudnn_flashmla backward builds one kernel for every slot width up to max_slots, so it needs max_slots"
+    )
+    return BACKWARD_BACKENDS[backend](q, kv, out, grad_out, indices, lse, sm_scale, max_slots)
+
+
 @dsv4_sparse_attn_backward.register_fake
 def _dsv4_sparse_attn_backward_fake(
     q: torch.Tensor,
@@ -208,15 +395,19 @@ def _dsv4_sparse_attn_backward_fake(
     indices: torch.Tensor,
     lse: torch.Tensor,
     sm_scale: float | None = None,
+    backend: str = "tilelang",
+    max_slots: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     return torch.empty_like(q), torch.empty_like(kv), torch.empty_like(lse)
 
 
 def _dsv4_sparse_attn_setup_context(ctx, inputs, output) -> None:
-    q, kv, indices, sinks, sm_scale, _block_I, _num_stages, _threads = inputs
+    q, kv, indices, sinks, sm_scale, _block_I, _num_stages, _threads, backend, max_slots = inputs
     out, lse = output
     ctx.save_for_backward(q, kv, out, indices, lse, sinks)
     ctx.sm_scale = sm_scale
+    ctx.backend = backend
+    ctx.max_slots = max_slots
     ctx.mark_non_differentiable(lse)
 
 
@@ -230,12 +421,14 @@ def _dsv4_sparse_attn_autograd_backward(ctx, grad_out: torch.Tensor, _grad_lse:
         indices,
         lse.detach(),
         ctx.sm_scale,
+        backend=ctx.backend,
+        max_slots=ctx.max_slots,
     )
     # dp_k/dsink = -p_k * p_sink, so do[d]/dsink = -p_sink * o[d] and the head's sink gradient
     # contracts to -p_sink * Delta. The sink logit is unscaled, hence no sm_scale factor.
     p_sink = torch.exp2(sinks.float().view(1, 1, -1) * LOG2E - lse)
     dsink = -(p_sink * delta).sum(dim=(0, 1)).to(sinks.dtype)
-    return dq, dkv, None, dsink, None, None, None, None
+    return dq, dkv, None, dsink, None, None, None, None, None, None
 
 
 dsv4_sparse_attn.register_autograd(_dsv4_sparse_attn_autograd_backward, setup_context=_dsv4_sparse_attn_setup_context)
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 00d38c83c4a..d3ffdfbe2ac 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_bwd.py
@@ -364,11 +364,16 @@ def cudnn_sparse_mla_backward(
 
         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
+        # A -inf sink leaves the softmax unchanged.
+        from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 import (
+            flash_attn_bwd_sm100,
+            flash_attn_bwd_sm100_workspace_size,
+        )
 
         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)
+        workspace_bytes = flash_attn_bwd_sm100_workspace_size(B * S, B * S_kv, D_qk, H)
+        workspace = torch.empty(workspace_bytes, dtype=torch.uint8, device=q.device)
+        dq, dkv, _ = flash_attn_bwd_sm100(*args, no_sink, flat_indices, softmax_scale=sm_scale, workspace=workspace)
     return dq.view_as(q), dkv.view_as(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 3bf6aa70e0c..39b310637a7 100644
--- a/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
+++ b/src/prime_rl/trainer/models/kernels/sparse_mla_fwd.py
@@ -235,8 +235,7 @@ 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 for both; ``cudnn_flashmla`` is FlashMLA forward + cuDNN backward
-    (SM90; SM100 / SM103 untested).
+    ``tilelang`` is TileLang for both; ``cudnn_flashmla`` is FlashMLA forward + cuDNN backward.
     """
     assert q.is_contiguous() and kv.is_contiguous() and indices.is_contiguous()
     batch, seq_len, heads, dim_plus_tail_dim = q.shape
diff --git a/tests/unit/train/models/test_deepseek_v4_kernels.py b/tests/unit/train/models/test_deepseek_v4_kernels.py
index e9da74d9047..3e2e885231c 100644
--- a/tests/unit/train/models/test_deepseek_v4_kernels.py
+++ b/tests/unit/train/models/test_deepseek_v4_kernels.py
@@ -1,4 +1,5 @@
 import copy
+import importlib.util
 import math
 from collections.abc import Callable
 from unittest.mock import MagicMock
@@ -10,7 +11,7 @@
 
 from prime_rl.trainer.models.deepseek_v4 import DeepseekV4Config
 from prime_rl.trainer.models.deepseek_v4 import attention as dsv4_attention
-from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Attention, PackedContext
+from prime_rl.trainer.models.deepseek_v4.attention import DeepseekV4Attention, PackedContext, max_sparse_attn_slots
 from prime_rl.trainer.models.deepseek_v4.hyperconnections import DeepseekV4HyperConnection, DeepseekV4UnweightedRMSNorm
 from prime_rl.trainer.models.deepseek_v4.rotary import DeepseekV4RotaryEmbedding
 from prime_rl.trainer.models.kernels.deepseek_v4 import IGNORE_SLOT, dsv4_mhc
@@ -201,6 +202,7 @@ def _compare_accumulated_grads(
 DIM = V4FLASH_MODEL["head_dim"]
 KV_GROUP = V4FLASH_MODEL["num_key_value_heads"]
 TOPK = V4FLASH_MODEL["sliding_window"] + V4FLASH_MODEL["index_topk"]
+MAX_SLOTS = max_sparse_attn_slots(V4FLASH_CONFIG)
 SM_SCALE = DIM**-0.5
 
 # The first is aligned to both of the backward's tile sizes (32 in `preprocess`, 64 in
@@ -209,6 +211,16 @@ def _compare_accumulated_grads(
 SHAPES = [(1, 256, 1024), (1, 200, 1000), (3, 128, 768)]
 SHAPE_IDS = ["aligned", "misaligned", "batched"]
 
+requires_cudnn_flashmla = pytest.mark.skipif(
+    importlib.util.find_spec("flash_mla") is None
+    or importlib.util.find_spec("cudnn") is None
+    or not torch.cuda.is_available()
+    or (torch.cuda.get_device_capability()[0] != 9 and torch.cuda.get_device_capability() not in ((10, 0), (10, 3))),
+    reason="the cudnn_flashmla backend needs the `flash-mla` extra, the cuDNN frontend and an SM90, SM100 or SM103 GPU",
+)
+
+DSV4_BACKENDS = ["tilelang", pytest.param("cudnn_flashmla", marks=requires_cudnn_flashmla)]
+
 # What a real query with a short window or a saturated top-k looks like; a masked slot still
 # costs a GEMM column.
 MASKED_FRACTION = 0.25
@@ -219,6 +231,9 @@ def _compare_accumulated_grads(
 OUT_RTOL = 1e-2
 # The LSE is float32 throughout on both sides.
 LSE_RTOL = 5e-7
+# FlashMLA returns a natural-log LSE without the sink, so the base-2 sink-inclusive one is two more
+# float32 roundings away; measured at up to 5.1e-7 on the benchmark corpus, about 4 ulps.
+LSE_RTOL_BY_BACKEND: dict[str, float] = {"cudnn_flashmla": 1e-6}
 DQ_RTOL = 1e-2
 # The vendored kernel this one forked from rounds `P` and `dP` to bfloat16 before the `dKV` GEMMs
 # while the oracle keeps them in float32, on top of the bfloat16 `kv` both sides share.
@@ -329,15 +344,16 @@ def _reference_lse(q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, sin
     return torch.cat([logits, sink_logits], dim=-1).logsumexp(dim=-1) * math.log2(math.e)
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @pytest.mark.parametrize(("batch", "seq_len", "seq_len_kv"), SHAPES, ids=SHAPE_IDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_kernel_forward_matches_the_dense_reference(batch, seq_len, seq_len_kv):
+def test_kernel_forward_matches_the_dense_reference(batch, seq_len, seq_len_kv, backend):
     """Output and log-sum-exp against the float32 gather oracle, which has identical semantics."""
     q, kv, indices, sinks = _inputs(batch, seq_len, seq_len_kv)
 
     with torch.no_grad():
-        out, lse = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE)
+        out, lse = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE, backend=backend, max_slots=MAX_SLOTS)
         # Float32 inputs to the oracle, which is what runs it in float32: it follows the dtype it
         # is handed. Widened here rather than inside it, so the exact answer is what the bound is
         # measured against instead of one rounded back to bfloat16.
@@ -347,12 +363,13 @@ def test_kernel_forward_matches_the_dense_reference(batch, seq_len, seq_len_kv):
     assert out.shape == q.shape and out.dtype == torch.bfloat16
     assert lse.shape == (batch, seq_len, HEADS) and lse.dtype == torch.float32
     _assert_relative(out, reference_out, OUT_RTOL, "output")
-    _assert_relative(lse, reference_lse, LSE_RTOL, "lse")
+    _assert_relative(lse, reference_lse, LSE_RTOL_BY_BACKEND.get(backend, LSE_RTOL), "lse")
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_kernel_pads_a_slot_count_its_tile_does_not_divide():
+def test_kernel_pads_a_slot_count_its_tile_does_not_divide(backend):
     """A caller states the slots it means and the kernel covers the difference to its own tile.
 
     The gather-slot axis is tiled at 64, but that is a fact about these kernels rather than
@@ -368,16 +385,17 @@ def test_kernel_pads_a_slot_count_its_tile_does_not_divide():
     indices = _build_indices(batch, seq_len, seq_len_kv, MASKED_FRACTION, topk=unaligned)
 
     with torch.no_grad():
-        out, _ = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE)
+        out, _ = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE, backend=backend, max_slots=MAX_SLOTS)
         reference = _dense_reference(q.float(), kv.float(), indices, sinks, SM_SCALE)
 
     assert out.shape == q.shape
     _assert_relative(out, reference, OUT_RTOL, "output")
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_fully_masked_query_reads_as_zero_keys():
+def test_fully_masked_query_reads_as_zero_keys(backend):
     """A query with no keys at all must emit exactly zero, on the sink term alone.
 
     The sink carries the softmax denominator by itself here, which is what keeps `lse` finite
@@ -394,7 +412,7 @@ def test_fully_masked_query_reads_as_zero_keys():
     indices[:, 0] = IGNORE_SLOT  # the first query gathers nothing
 
     with torch.no_grad():
-        out, lse = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE)
+        out, lse = dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE, backend=backend, max_slots=MAX_SLOTS)
 
     assert torch.equal(out[:, 0], torch.zeros_like(out[:, 0])), "a fully masked query must emit exactly zero"
     expected_lse = sinks.float() * math.log2(math.e)
@@ -452,20 +470,25 @@ def main(
     )
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @pytest.mark.parametrize(("batch", "seq_len", "seq_len_kv"), SHAPES, ids=SHAPE_IDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_kernel_backward_matches_autograd_through_the_reference(batch, seq_len, seq_len_kv):
+def test_kernel_backward_matches_autograd_through_the_reference(batch, seq_len, seq_len_kv, backend):
     """All three differentiable inputs, each against its own bound.
 
     `dsink` is the one term the kernel forms in torch rather than in tilelang, out of the `Delta`
-    the backward returns, so it is the assertion that would catch a wrong `Lse` convention.
+    the backward returns, so it is the assertion that would catch a wrong `Lse` convention. The
+    first query gathers nothing, so its `Lse` is the sink term alone on the way back too.
     """
     q, kv, indices, sinks = _inputs(batch, seq_len, seq_len_kv)
+    indices[:, 0] = IGNORE_SLOT
     kernel_q, kernel_kv, kernel_sinks = _leaves(q, kv, sinks)
     reference_q, reference_kv, reference_sinks = _float32_leaves(q, kv, sinks)
 
-    out, _lse = dsv4_sparse_attn(kernel_q, kernel_kv, indices, kernel_sinks, SM_SCALE)
+    out, _lse = dsv4_sparse_attn(
+        kernel_q, kernel_kv, indices, kernel_sinks, SM_SCALE, backend=backend, max_slots=MAX_SLOTS
+    )
     # One weight tensor for both losses, so the two backwards are the same function of the same
     # numbers and any difference belongs to the kernel.
     weight = torch.randn_like(out)
@@ -482,9 +505,45 @@ def test_kernel_backward_matches_autograd_through_the_reference(batch, seq_len,
     _assert_relative(kernel_sinks.grad, reference_sinks.grad, DSINK_RTOL, "dsink")
 
 
+@requires_cudnn_flashmla
+@requires_sparse_attn_kernel
+@requires_datacenter_gpu
+def test_cudnn_flashmla_backward_compiles_one_main_kernel_for_every_width():
+    """Slot widths from 128 to 8320 share one cuDNN main kernel, and each matches the reference.
+
+    The cuDNN frontend keys its main-kernel compile on the slot width. `cudnn_dsa_patch` keys it on a
+    clamp instead, through private frontend internals that an upgrade could move without an error. 8320 is
+    DeepSeek-V4-Flash's widest HCA layer, `128 + 1048576 // 128`.
+    """
+    from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm90 import flash_attn_bwd_sm90
+    from cudnn.deepseek_sparse_attention.sparse_attention_backward._interface_sm100 import flash_attn_bwd_sm100
+
+    batch, seq_len, max_slots = 1, 64, 8320
+    bwd = flash_attn_bwd_sm90 if torch.cuda.get_device_capability()[0] == 9 else flash_attn_bwd_sm100
+    bwd.compile_cache.clear()
+    for width in (128, 640, max_slots):
+        q, kv, _, sinks = _inputs(batch, seq_len, max_slots)
+        indices = _build_indices(batch, seq_len, max_slots, MASKED_FRACTION, topk=width)
+        kernel_q, kernel_kv, kernel_sinks = _leaves(q, kv, sinks)
+        reference_q, reference_kv, reference_sinks = _float32_leaves(q, kv, sinks)
+
+        out, _lse = dsv4_sparse_attn(
+            kernel_q, kernel_kv, indices, kernel_sinks, SM_SCALE, backend="cudnn_flashmla", max_slots=max_slots
+        )
+        weight = torch.randn_like(out)
+        (out * weight).sum().backward()
+        reference_out = _dense_reference(reference_q, reference_kv, indices, reference_sinks, SM_SCALE)
+        (reference_out * weight).sum().backward()
+
+        _assert_relative(kernel_q.grad, reference_q.grad, DQ_RTOL, f"dq at width {width}")
+        _assert_relative(kernel_kv.grad, reference_kv.grad, DKV_RTOL, f"dkv at width {width}")
+    assert len(bwd.compile_cache) == 1
+
+
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_kernel_traces_under_torch_compile():
+def test_kernel_traces_under_torch_compile(backend):
     """`torch.compile(fullgraph=True)` through the op, forward and backward.
 
     `apply_compile` in `prime_rl/trainer/model.py` compiles each decoder layer, so every real
@@ -498,7 +557,7 @@ def test_kernel_traces_under_torch_compile():
     compiled_leaves = _leaves(q, kv, sinks)
 
     def attend(q: torch.Tensor, kv: torch.Tensor, sinks: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        return dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE)
+        return dsv4_sparse_attn(q, kv, indices, sinks, SM_SCALE, backend=backend, max_slots=MAX_SLOTS)
 
     out, lse = attend(*eager_leaves)
     # One weight tensor for both losses, so the two backwards are the same function of the same
@@ -569,9 +628,9 @@ def _record_attention(monkeypatch) -> dict[str, torch.Tensor]:
     recorded: dict[str, torch.Tensor] = {}
     real_kernel = dsv4_attention.dsv4_sparse_attn
 
-    def kernel(q, kv_buf, indices, sinks, scale):
+    def kernel(q, kv_buf, indices, *args, **kwargs):
         recorded["kv_buf"], recorded["indices"] = kv_buf, indices
-        recorded["kernel"] = real_kernel(q, kv_buf, indices, sinks, scale)
+        recorded["kernel"] = real_kernel(q, kv_buf, indices, *args, **kwargs)
         return recorded["kernel"]
 
     monkeypatch.setattr(dsv4_attention, "dsv4_sparse_attn", kernel)
@@ -825,10 +884,11 @@ def test_absent_slots_are_marked_negative_rather_than_pointed_at_a_pad_row(doc_l
 EAGER_KERNEL_RTOL, EAGER_KERNEL_GRAD_RTOL = 1e-2, 5e-2
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @pytest.mark.parametrize("doc_lens", EAGER_KERNEL_DOC_LENS, ids=EAGER_KERNEL_DOC_IDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_sparse_attention_kernel_matches_eager(doc_lens, monkeypatch):
+def test_sparse_attention_kernel_matches_eager(doc_lens, backend, monkeypatch):
     """The fused kernel against the naive dense softmax, single-document and packed.
 
     Every other kernel test reaches eager only transitively: the kernel is compared to the gather
@@ -850,6 +910,7 @@ def test_sparse_attention_kernel_matches_eager(doc_lens, monkeypatch):
     """
     seq_len = sum(doc_lens)
     kernel_module = v4flash_attention(V4FLASH_CSA_LAYER, dtype=torch.bfloat16)
+    kernel_module.dsa_backend = backend
     eager_module = copy.deepcopy(kernel_module).float()
     eager_reference.use_eager_attention(eager_module)
 
@@ -861,9 +922,9 @@ def test_sparse_attention_kernel_matches_eager(doc_lens, monkeypatch):
     calls = []
     real_kernel = dsv4_attention.dsv4_sparse_attn
 
-    def counting_kernel(q, kv_buf, indices, sinks, scale):
+    def counting_kernel(q, *args, **kwargs):
         calls.append(q.shape[1])
-        return real_kernel(q, kv_buf, indices, sinks, scale)
+        return real_kernel(q, *args, **kwargs)
 
     monkeypatch.setattr(dsv4_attention, "dsv4_sparse_attn", counting_kernel)
 
@@ -886,9 +947,10 @@ def counting_kernel(q, kv_buf, indices, sinks, scale):
     _assert_relative(kernel_input.grad, eager_input.grad, EAGER_KERNEL_GRAD_RTOL, "hidden states gradient")
 
 
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
-def test_sparse_attention_kernel_packed_matches_unpacked(monkeypatch):
+def test_sparse_attention_kernel_packed_matches_unpacked(backend, monkeypatch):
     """The fused kernel path, end to end through one CSA layer, must respect documents.
 
     The same invariant its float32 neighbours assert, run in bfloat16 because that is the only
@@ -903,6 +965,7 @@ def test_sparse_attention_kernel_packed_matches_unpacked(monkeypatch):
     a fallback would leave this test asserting a property of the gather reference instead.
     """
     module = v4flash_attention(V4FLASH_CSA_LAYER, dtype=torch.bfloat16)
+    module.dsa_backend = backend
     packed = _packed_context(KERNEL_DOC_LENS, torch.bfloat16, V4FLASH_CONFIG)
     with torch.device("cuda"):
         hidden = torch.randn(1, sum(KERNEL_DOC_LENS), V4FLASH_MODEL["hidden_size"], dtype=torch.bfloat16)
@@ -911,9 +974,9 @@ def test_sparse_attention_kernel_packed_matches_unpacked(monkeypatch):
     calls = []
     real_kernel = dsv4_attention.dsv4_sparse_attn
 
-    def counting_kernel(q, kv_buf, indices, sinks, scale):
+    def counting_kernel(q, *args, **kwargs):
         calls.append(q.shape[1])
-        return real_kernel(q, kv_buf, indices, sinks, scale)
+        return real_kernel(q, *args, **kwargs)
 
     monkeypatch.setattr(dsv4_attention, "dsv4_sparse_attn", counting_kernel)
 
@@ -960,9 +1023,9 @@ def test_sparse_attention_kernel_trains_every_parameter(monkeypatch):
     calls = []
     real_kernel = dsv4_attention.dsv4_sparse_attn
 
-    def counting_kernel(q, kv_buf, indices, sinks, scale):
+    def counting_kernel(q, *args, **kwargs):
         calls.append(q.shape[1])
-        return real_kernel(q, kv_buf, indices, sinks, scale)
+        return real_kernel(q, *args, **kwargs)
 
     monkeypatch.setattr(dsv4_attention, "dsv4_sparse_attn", counting_kernel)
 
@@ -1092,9 +1155,10 @@ def _assert_within_the_bfloat16_floor(
 
 @requires_sparse_attn_kernel
 @requires_datacenter_gpu
+@pytest.mark.parametrize("backend", DSV4_BACKENDS)
 @pytest.mark.parametrize("doc_lens", PARITY_DOC_LENS, ids=PARITY_DOC_IDS)
 @pytest.mark.parametrize("layer_idx", V4FLASH_LAYERS, ids=V4FLASH_LAYER_IDS)
-def test_kernel_and_eager_consumers_agree_on_shared_weights(layer_idx, doc_lens):
+def test_kernel_and_eager_consumers_agree_on_shared_weights(layer_idx, doc_lens, backend):
     """The two consumers of one `SparseAttnInputs` must compute the same attention, and its gradient.
 
     They are handed the identical index tensor, so this is not about which keys a query reads,
@@ -1111,6 +1175,7 @@ def test_kernel_and_eager_consumers_agree_on_shared_weights(layer_idx, doc_lens)
     its slots out differently and only CSA had coverage.
     """
     module = v4flash_attention(layer_idx, dtype=torch.bfloat16)
+    module.dsa_backend = backend
     weights = module.state_dict()
     eager_bf16 = v4flash_attention(layer_idx, dtype=torch.bfloat16, eager=True)
     eager_bf16.load_state_dict(weights)