[Fix] Make flash_mla_cudnn SparseMLA backward deterministic under XTUNER_DETERMINISTIC

tests/ops/test_flash_mla_cudnn_sparse_mla.py

@@ -5,6 +5,9 @@
     test_forward_backward_matches_torch_reference         前反向与 torch 参考一致
     test_rejects_576_head_dim_without_matching_value_dim  维度白名单不接受错配的 value_dim
     test_accepts_512_head_dim_whitelist_entry             NoPE 的 (512, 512) 在白名单内
+TestDeterministicSparseMLABackward
+    test_deterministic_mode_backward_bitwise_repeatable   确定性模式下反向逐位可重复
+    test_deterministic_backward_matches_torch_reference   确定性参考反向与 torch 参考一致
 """
 
 import importlib.util
@@ -106,3 +109,60 @@ def test_accepts_512_head_dim_whitelist_entry(self):
         kv = torch.randn(4, 1, 512, device="cuda", dtype=torch.bfloat16)
         indices = torch.zeros(4, 1, 64, device="cuda", dtype=torch.int32)
         validate_sparse_mla_inputs(q, kv, indices, value_dim=512)  # must not raise
+
+
+class TestDeterministicSparseMLABackward:
+    """XTUNER_DETERMINISTIC 下 flash_mla_cudnn 的 fork 分片归约确定性反向。
+
+    cuDNN 稀疏注意力反向的 dkv 是无序原子累加，同输入两次反向不逐位一致，CI
+    双跑 grad_norm 因此抖动；确定性模式改走 fork 分片归约（cudnn_dsa_det.py：
+    每 CTA 写自己的 FP32 分片、按固定顺序归约，零跨 CTA 原子竞争）。
+    """
+
+    @pytest.fixture
+    def deterministic_mode(self, monkeypatch):
+        from xtuner.v1.ops.sparse_mla import cudnn_dsa
+
+        # 反向分支在调用时读模块级常量，monkeypatch 即可覆盖 import 期绑定。
+        monkeypatch.setattr(cudnn_dsa, "XTUNER_DETERMINISTIC", True)
+        previous = torch.are_deterministic_algorithms_enabled()
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        yield
+        torch.use_deterministic_algorithms(previous)
+
+    def _run_backward(self, q0, kv0, indices, grad_output):
+        q = q0.detach().clone().requires_grad_()
+        kv = kv0.detach().clone().requires_grad_()
+        out = sparse_mla(
+            q, kv, indices.to(torch.int32), scaling=q0.shape[-1] ** -0.5, value_dim=512, backend="flash_mla_cudnn"
+        )
+        out.raw_output.backward(grad_output)
+        return q.grad, kv.grad
+
+    @pytest.mark.skipif(not _flash_mla_cudnn_available(), reason="requires FlashMLA + cuDNN DSA runtimes")
+    def test_deterministic_mode_backward_bitwise_repeatable(self, deterministic_mode):
+        # 同输入两次反向必须逐位一致——cuDNN 反向在此会因 dkv 原子序不同而失败。
+        q0, kv0, indices = _nope_sparse_mla_inputs()
+        grad_output = torch.randn(64, 64, 512, device="cuda", dtype=torch.bfloat16)
+
+        dq1, dkv1 = self._run_backward(q0, kv0, indices, grad_output)
+        dq2, dkv2 = self._run_backward(q0, kv0, indices, grad_output)
+
+        assert torch.equal(dq1, dq2)
+        assert torch.equal(dkv1, dkv2)
+
+    @pytest.mark.skipif(not _flash_mla_cudnn_available(), reason="requires FlashMLA + cuDNN DSA runtimes")
+    def test_deterministic_backward_matches_torch_reference(self, deterministic_mode):
+        # 参考反向只是换累加方式，数值语义必须保持与 torch 参考一致（容差同主测试）。
+        q0, kv0, indices = _nope_sparse_mla_inputs()
+        scaling = 1 / (q0.shape[-1] ** 0.5)
+        q_ref = q0.detach().clone().requires_grad_()
+        kv_ref = kv0.detach().clone().requires_grad_()
+        expected = sparse_mla(q_ref, kv_ref, indices, scaling=scaling, value_dim=512, backend="torch")
+        grad_output = torch.randn_like(expected.raw_output)
+        expected.raw_output.backward(grad_output)
+
+        dq, dkv = self._run_backward(q0, kv0, indices, grad_output)
+
+        torch.testing.assert_close(dq, q_ref.grad, atol=BF16_ATOL, rtol=BF16_RTOL)
+        torch.testing.assert_close(dkv, kv_ref.grad, atol=DKV_ATOL, rtol=DKV_RTOL)

xtuner/v1/ops/sparse_mla/_cudnn_dsa_bwd_det_subclass.py

@@ -0,0 +1,96 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""Deterministic DSA backward kernel via subclassing the cuDNN CuTe-DSL kernel.
+
+Instead of vendoring the 2300-line ``dsa_bwd_sm90.py`` from the cudnn package, this
+subclass overrides only ``scatter_dkv_atomic`` so each CTA writes its dKV contribution
+into its own slice of a widened FP32 accumulator (plain single-writer atomics per
+address) instead of racing all CTAs on one shared buffer. The driver
+(``cudnn_dsa_det.py``) then reduces the slices in a fixed order. Mechanism follows FA2's
+deterministic backward (split buffer + ordered reduction).
+
+Fragility note: the override mirrors cudnn 1.28.0's method body; if the cudnn package
+upgrades and changes that method's internals, ``flash_mla_cudnn.py``'s deterministic
+dispatch must fall back (guarded by the import/capability check there).
+"""
+
+import cutlass
+import cutlass.cute as cute
+from cudnn.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm90 import (
+    FlashAttentionDSABackwardSm90,
+)
+from cudnn.deepseek_sparse_attention.utils.sm90.primitives import (
+    atomic_add_fp32x4,
+    make_acc_tensor_mn_view,
+)
+from cutlass import Int32, const_expr
+
+
+class FlashAttentionDSABackwardSm90Det(FlashAttentionDSABackwardSm90):
+    """cuDNN SM90 DSA backward with per-CTA dKV accumulator slices.
+
+    The driver must set ``_det_slice_size`` (elements per slice, i.e.
+    ``seqlen_k_rounded * head_dim_rounded``) and ``_det_n_splits`` on the instance before
+    ``cute.compile``; they are read at trace time and baked in as constants. The widened
+    accumulator keeps the stock 3D layout with ``n_splits * slice_size`` elements in the
+    last dim, so the stock ``__call__``/transpose paths are unchanged.
+    """
+
+    _det_slice_size: int = 0
+    _det_n_splits: int = 1
+
+    @cute.jit
+    def scatter_dkv_atomic(
+        self,
+        acc: cute.Tensor,
+        chunk_idx: cutlass.Constexpr[int],
+        mTopkIdxs_cur: cute.Tensor,
+        mdKVaccum_cur: cute.Tensor,
+        n_block: Int32,
+        topK: Int32,
+        thr_mma: cute.TiledMma,
+        tidx: Int32,
+    ):
+        """Stock method + one change: target_ptr gains a per-CTA slice offset.
+
+        Every address within a slice still has a single writer (the stock kernel's
+        fragment partition is exclusive and duplicate topk entries accumulate in
+        program order on the same thread), so the per-slice atomics are exact, and the
+        cross-CTA nondeterminism disappears because CTAs no longer share slices.
+        """
+        tile_shape = (self.tile_n, self.hdim_chunk)
+        cDKV = cute.make_identity_tensor(tile_shape if const_expr(not self.dKV_swapAB) else tile_shape[::-1])
+        tScDKV_mn = make_acc_tensor_mn_view(thr_mma.partition_C(cDKV), transpose=self.dKV_swapAB)
+        ROW = 0 if const_expr(not self.dKV_swapAB) else 1
+
+        acc_mn = make_acc_tensor_mn_view(acc, transpose=self.dKV_swapAB)
+        nrow = const_expr(cute.size(acc_mn.shape[0]))
+        ncol = const_expr(cute.size(acc_mn.shape[1]))
+
+        rank = tidx % 4
+        # DET: this CTA's private slice base within the widened accumulator.
+        split_offset = (cute.arch.block_idx()[0] % self._det_n_splits) * self._det_slice_size
+
+        for r in cutlass.range_constexpr(nrow):
+            local_row = tScDKV_mn[r, 0][ROW]
+            global_topk_row = n_block * self.tile_n + local_row
+            if global_topk_row < topK:
+                global_kv_row = mTopkIdxs_cur[global_topk_row]
+                should_accumulate_dkv = const_expr(True)
+                if const_expr(not self.have_topk_length):
+                    should_accumulate_dkv = global_kv_row >= 0
+                if should_accumulate_dkv:
+                    row_base = global_kv_row * self.tile_hdim + chunk_idx * self.hdim_chunk
+
+                    # Each group of 4 MN-view cols maps to one N-tile.
+                    # Interleaved: fake_pos = c4 * 16 + rank * 4
+                    for c4 in cutlass.range_constexpr(ncol // 4):
+                        fake_offset = const_expr(c4) * 16 + rank * 4
+                        c = const_expr(c4 * 4)
+                        target_ptr = mdKVaccum_cur.iterator + split_offset + row_base + fake_offset
+                        atomic_add_fp32x4(
+                            acc_mn[r, c + 0],
+                            acc_mn[r, c + 1],
+                            acc_mn[r, c + 2],
+                            acc_mn[r, c + 3],
+                            target_ptr,
+                        )

xtuner/v1/ops/sparse_mla/cudnn_dsa.py

@@ -3,6 +3,8 @@
 import torch
 from torch import Tensor
 
+from xtuner.v1.utils import XTUNER_DETERMINISTIC
+
 from .protocol import SparseMLAOutputs
 from .tilelang import validate_sparse_mla_inputs
 
@@ -99,15 +101,36 @@ def _cudnn_dsa_sparse_mla_backward_op(
     softmax_lse: Tensor,
     scaling: float,
 ) -> tuple[Tensor, Tensor]:
-    from cudnn.deepseek_sparse_attention.sparse_attention_backward import sparse_attention_backward_wrapper
-
     if kv.shape[1] != 1 or indices.shape[1] != 1:
         raise RuntimeError("cuDNN DSA SparseMLA backward currently supports kv_group=1 only.")
 
     # topk_length limits cuDNN to a physical prefix. KPool can put -1 before
     # valid pools or between pools and the appended tail, so compact first.
     # Forward keeps its original order; backward retains every valid duplicate.
     topk_idxs, topk_length = _compact_cudnn_topk_indices(indices[:, 0, :])
+
+    if XTUNER_DETERMINISTIC:
+        # Deterministic mode: same cuDNN kernel, but each CTA writes its own FP32
+        # dkv slice instead of racing atomics on a shared buffer; slices are reduced
+        # in fixed order afterwards (see cudnn_dsa_det.py). Bitwise repeatable at
+        # ~1.6x the stock cost. Covers every backend that routes through this op
+        # (flash_mla_cudnn and cudnn_dsa).
+        from .cudnn_dsa_det import flash_attn_bwd_sm90_multilaunch
+
+        dq, dkv = flash_attn_bwd_sm90_multilaunch(
+            q.contiguous(),
+            kv[:, 0, :].contiguous(),
+            raw_output.contiguous(),
+            grad_output.contiguous(),
+            softmax_lse.contiguous(),
+            topk_idxs,
+            topk_length,
+            softmax_scale=scaling,
+        )
+        return dq, dkv.unsqueeze(1)
+
+    from cudnn.deepseek_sparse_attention.sparse_attention_backward import sparse_attention_backward_wrapper
+
     attn_sink = torch.full((q.shape[1],), float("-inf"), dtype=torch.float32, device=q.device)
 
     outputs = sparse_attention_backward_wrapper(