[dev] Make fused CSA Top-K deterministic across no-grad and grad forwards

megatron/core/transformer/experimental_attention_variant/csa.py

@@ -29,6 +29,7 @@
     BSHDCompactIndexerWorkspace,
     FusedCSAIndexerSparseAttnFromTopkFunc,
     THDCompactIndexerWorkspace,
+    _compact_flat_topk_idxs,
     batch_of_row,
     bshd_compact_indexer_available,
     build_flat_topk_idxs,
@@ -2202,8 +2203,13 @@ def _forward_fused_indexer_inference(
             deterministic=self.config.deterministic_mode,
         )
         compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1)
+        # Same lowering as ``FusedCSAIndexerSparseAttnFunc.forward`` (the
+        # grad-enabled path): compressed ids, then window ids, globalized and
+        # compacted. FlashMLA's online softmax accumulates in key order, so a
+        # forward-only pass reproduces the training forward only if it hands
+        # the kernel the same id sequence.
         flat_idxs, flat_tlen = build_flat_topk_idxs(
-            window_idxs, compress_topk_idxs, batch_size=b, compact=True
+            compress_topk_idxs, window_idxs, batch_size=b, compact=True
         )
         nvtx_range_pop("compressed_indices")
 
@@ -2656,6 +2662,17 @@ def _forward_fused_indexer_inference_thd(
                 deterministic=self.config.deterministic_mode,
             )
 
+        # ``build_attention_indices`` has two output layouts. Its default puts
+        # window ids first and returns the list already compacted with a
+        # ``topk_length``. Its indexer-loss layout puts the compressed ids
+        # first at fixed positions (so the grad-enabled path can slice the
+        # window segment for the dense teacher) and leaves compaction to the
+        # caller, returning ``None`` for the length. The grad-enabled path in
+        # ``FusedCSAIndexerSparseAttnFunc.forward`` uses the latter; take the
+        # same layout here and compact it ourselves so both forwards hand
+        # FlashMLA an identical id sequence. With no compressed ids the two
+        # layouts coincide, so keep the default and its precomputed length.
+        compressed_first = topk_indices_cmp.shape[-1] > 0
         flat_idxs, flat_tlen, _, _ = thd_layout_kernels.build_attention_indices(
             cu_seqlens_q,
             0,
@@ -2666,11 +2683,14 @@ def _forward_fused_indexer_inference_thd(
             topk_indices_cmp.shape[-1],
             topk_indices_cmp,
             cu_seqlens_compressed=cu_seqlens_compressed,
+            for_indexer_loss=compressed_first,
             compressed_base=compressed_base,
             compressed_rows=compressed_rows,
             compressed_is_sequence_major=True,
             output_alignment=get_flash_mla_topk_alignment(),
         )
+        if compressed_first:
+            flat_idxs, flat_tlen = _compact_flat_topk_idxs(flat_idxs)
         output = csa_sparse_attn(
             query,
             kv_full_thd,

megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py

@@ -1525,6 +1525,31 @@ def csa_sparse_attn(
 # ---------------------------------------------------------------------------
 
 
+def _stable_topk_indices(scores: Tensor, seq_lens: Tensor, topk_k: int) -> Tensor:
+    """Select the ``topk_k`` highest-scoring key ids per row with a fixed tie order.
+
+    Rows may only draw from their first ``seq_lens[row]`` key columns; a row with
+    fewer valid keys than ``topk_k`` is padded with ``-1``. Exact score ties are
+    resolved toward the smallest key id and the selected ids are returned in
+    descending-score order, so identical inputs always yield identical ids. The
+    radix Top-K kernel does not order equal scores, which matters for ReLU-scored
+    indexers where many keys share a score of exactly zero.
+
+    Args:
+        scores: ``(rows, sk)`` fp32 indexer scores; masked positions hold ``-inf``.
+        seq_lens: ``(rows,)`` int32 number of candidate key columns per row.
+        topk_k: number of ids to select, at most ``sk``.
+
+    Returns:
+        ``(rows, topk_k)`` int32 key ids, ``-1`` where a row has no more valid keys.
+    """
+    columns = torch.arange(scores.shape[-1], device=scores.device, dtype=seq_lens.dtype)
+    candidates = scores.masked_fill(columns.unsqueeze(0) >= seq_lens.unsqueeze(1), float("-inf"))
+    sorted_scores, order = torch.sort(candidates, dim=-1, descending=True, stable=True)
+    selected = order[:, :topk_k].to(torch.int32)
+    return selected.masked_fill(torch.isneginf(sorted_scores[:, :topk_k]), -1)
+
+
 def _indexer_topk_core(
     q: Tensor,
     k: Tensor,
@@ -1829,10 +1854,13 @@ def _indexer_topk_core(
 
     # ---------------- Shared: radix top-K + pad-to-topk -----------------
     topk_k = min(topk, sk)
-    tk_result = _DSA.indexer_top_k_wrapper(
-        scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False
-    )
-    topk_indices = tk_result["indices"]  # (total_q, topk_k) int32
+    if deterministic:
+        topk_indices = _stable_topk_indices(scores_flat, seq_lens, topk_k)
+    else:
+        tk_result = _DSA.indexer_top_k_wrapper(
+            scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False
+        )
+        topk_indices = tk_result["indices"]  # (total_q, topk_k) int32
 
     if is_thd:
         topk_indices, topk_length = thd_indexer_kernels.sanitize_topk(
@@ -1901,8 +1929,9 @@ def indexer_topk(
             outputs are allocated by each dispatch. A matching workspace is
             required while capturing the compact path.
         deterministic: resolve exact-value ties at the K-th boundary toward
-            the smallest local KV indices. The output slot order remains
-            unspecified.
+            the smallest local KV indices. Compact dispatch leaves the output
+            slot order unspecified; the standalone Top-K fallback returns ids in
+            descending-score order.
         return_softmax: also return the compact kernel's Top-K softmax. The
             third return is ``None`` when compact dispatch is unavailable.
 

tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py

@@ -2314,6 +2314,73 @@ def test_b1_sbhd_thd_parity_inference_path_c(self):
             f"{(out_sbhd.float() - out_thd.float()).abs().max().item():.4e}"
         )
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_thd_fused_no_grad_and_grad_forwards_share_key_order(self, monkeypatch):
+        """Path C (no autograd) and Path B (grad-enabled) hand FlashMLA the same
+        compacted key ids in the same slot order, so re-scoring a batch without
+        gradients reproduces the training forward's attention exactly."""
+        from megatron.core.transformer.experimental_attention_variant.csa_utils import (
+            fused_sparse_attention as fsa,
+        )
+
+        try:
+            fsa._ensure_dsa_namespace()
+        except ImportError:
+            pytest.skip("cuDNN Frontend DSA namespace not available")
+
+        captured = []
+
+        def fake_flash(q, kv, topk_idxs, softmax_scale, d_v=512, **kwargs):
+            topk_length = kwargs.get("topk_length")
+            captured.append(
+                (topk_idxs.clone(), None if topk_length is None else topk_length.clone())
+            )
+            lse = torch.zeros(q.shape[0], q.shape[1], dtype=torch.float32, device=q.device)
+            return torch.zeros_like(q), lse, None
+
+        monkeypatch.setattr(fsa, "_csa_fwd_flash_mla", fake_flash)
+
+        overrides = {
+            "dsa_kernel_backend": "cudnn",
+            "deterministic_mode": True,
+            "dsa_indexer_loss_coeff": 0.0,
+        }
+        saved = {name: getattr(self.config, name) for name in overrides}
+        for name, value in overrides.items():
+            setattr(self.config, name, value)
+        try:
+            csa = self._build_csa(compress_ratio=4)
+            torch.manual_seed(7)
+            query, key, value, x, qr, packed = self._make_thd_inputs([96, 40, 72])
+            csa.train()
+            with torch.no_grad():
+                csa(
+                    query=query,
+                    key=key,
+                    value=value,
+                    attention_mask=None,
+                    x=x,
+                    qr=qr,
+                    packed_seq_params=packed,
+                )
+            csa(
+                query=query,
+                key=key,
+                value=value,
+                attention_mask=None,
+                x=x,
+                qr=qr,
+                packed_seq_params=packed,
+            )
+        finally:
+            for name, value in saved.items():
+                setattr(self.config, name, value)
+
+        assert len(captured) == 2
+        (no_grad_idxs, no_grad_lengths), (grad_idxs, grad_lengths) = captured
+        assert torch.equal(no_grad_lengths, grad_lengths)
+        assert torch.equal(no_grad_idxs, grad_idxs)
+
 
 # ===========================================================================
 # _apply_rope direct THD tests (4 corners: ratio={1, >1} × fused={False, True})

tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py

@@ -4162,7 +4162,7 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs):
                 max_seqlen_q=3,
                 max_seqlen_kv=2,
                 q_causal_offsets=q_causal_offsets,
-                deterministic=True,
+                deterministic=False,
             )
         # THD return shape: (total_q, topk) + (total_q,).
         assert topk_idxs.shape == (total_q, 2)
@@ -4172,6 +4172,50 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs):
         seq_lens = fake_dsa.indexer_top_k_wrapper.call_args.args[1]
         assert torch.equal(seq_lens, torch.tensor([1, 1, 2, 0, 0], device=q.device))
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_indexer_topk_thd_deterministic_fallback_skips_radix_kernel(
+        self, reset_lazy_kernel_state
+    ):
+        """Without a compact wrapper, ``deterministic=True`` selects ids from the
+        dense scores with a fixed tie order instead of the radix Top-K kernel."""
+        q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs()
+        total_q = q.shape[0]
+        q_causal_offsets = torch.tensor([5, 0], dtype=torch.int32, device=q.device)
+
+        fake_dsa = MagicMock(name='_DSA_thd_deterministic_stub')
+        fake_dsa.indexer_forward_top_k_wrapper = None
+        # Causal lengths are [1, 1, 2, 0, 0]. Row 0 scores its out-of-range key
+        # highest, row 2 has two valid keys with equal scores, rows 3/4 have none.
+        scores = torch.zeros(total_q, 2, dtype=torch.float32, device=q.device)
+        scores[0, 1] = 1.0
+        scores[3:] = float('-inf')
+        fake_dsa.indexer_forward_wrapper.return_value = {'scores': scores}
+        dk._DSA = fake_dsa
+
+        with pytest.warns(RuntimeWarning, match="Compact indexer.*falling back"):
+            topk_idxs, topk_len = indexer_topk(
+                q,
+                k,
+                w,
+                topk=2,
+                ratio=4,
+                cu_seqlens_q=cu_q,
+                cu_seqlens_kv=cu_kv,
+                max_seqlen_q=3,
+                max_seqlen_kv=2,
+                q_causal_offsets=q_causal_offsets,
+                deterministic=True,
+            )
+
+        fake_dsa.indexer_top_k_wrapper.assert_not_called()
+        expected = torch.tensor(
+            [[0, -1], [0, -1], [0, 1], [-1, -1], [-1, -1]], dtype=torch.int32, device=q.device
+        )
+        assert torch.equal(topk_idxs, expected)
+        assert torch.equal(
+            topk_len, torch.tensor([1, 1, 2, 0, 0], dtype=torch.int32, device=q.device)
+        )
+
     @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
     def test_indexer_topk_thd_compact_dispatch(self, reset_lazy_kernel_state):
         q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs()
@@ -5418,3 +5462,32 @@ def test_public_surface(self):
         assert issubclass(CSASparseAttnFunc, torch.autograd.Function)
         assert issubclass(FusedCSAIndexerSparseAttnFunc, torch.autograd.Function)
         assert issubclass(FusedCSAIndexerSparseAttnFromTopkFunc, torch.autograd.Function)
+
+
+# ---------------------------------------------------------------------------
+# _stable_topk_indices — deterministic fallback for the standalone Top-K
+# ---------------------------------------------------------------------------
+
+
+class TestStableTopkIndices:
+    """``_stable_topk_indices`` orders exact ties toward the smallest key id and
+    pads rows that run out of valid keys with ``-1``."""
+
+    def test_ties_resolve_toward_smallest_index(self):
+        scores = torch.tensor([[0.0, 3.0, 0.0, 3.0, 1.0, 0.0]])
+        seq_lens = torch.tensor([6], dtype=torch.int32)
+        out = dk._stable_topk_indices(scores, seq_lens, topk_k=4)
+        assert torch.equal(out, torch.tensor([[1, 3, 4, 0]], dtype=torch.int32))
+
+    def test_seq_lens_mask_and_padding(self):
+        scores = torch.tensor([[5.0, 4.0, 3.0, 2.0], [5.0, 4.0, 3.0, 2.0], [5.0, 4.0, 3.0, 2.0]])
+        seq_lens = torch.tensor([4, 2, 0], dtype=torch.int32)
+        out = dk._stable_topk_indices(scores, seq_lens, topk_k=3)
+        expected = torch.tensor([[0, 1, 2], [0, 1, -1], [-1, -1, -1]], dtype=torch.int32)
+        assert torch.equal(out, expected)
+
+    def test_masked_scores_are_never_selected(self):
+        scores = torch.tensor([[float('-inf'), 0.0, float('-inf'), 0.0]])
+        seq_lens = torch.tensor([4], dtype=torch.int32)
+        out = dk._stable_topk_indices(scores, seq_lens, topk_k=3)
+        assert torch.equal(out, torch.tensor([[1, 3, -1]], dtype=torch.int32))