From 20cbd65b4dd532e602f5459b507487fa26123b38 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Wed, 9 Sep 2026 11:14:17 +0000
Subject: [PATCH 01/10] [dev] Make fused CSA Top-K deterministic across no-grad
 and grad forwards

The standalone cuDNN radix Top-K returns a different key set on every call
when scores tie at the K-th boundary. ReLU-scored indexers tie constantly:
every key whose score collapses to zero is interchangeable, so any row with
fewer than K positive scores draws an arbitrary subset. The ``deterministic``
flag only reached the compact indexer kernel, which older cuDNN Frontend
builds do not ship, so the fallback stayed non-repeatable. Independently,
Path C (no-grad) lowered window ids before compressed ids while Path B
(grad-enabled) lowered them compressed-first, so the two forwards fed
FlashMLA the same keys in different orders and its online softmax rounded
differently.

Re-scoring a batch without gradients therefore did not reproduce the
training forward. In RL fine-tuning this shows up as a PPO ratio that
drifts from 1 on roughly 0.5% of tokens and a nonzero clipped fraction
with unchanged weights; on GB300 with cuDNN Frontend 1.26 the two-call
disagreement reached 76% of rows on tie-heavy inputs.

Under ``deterministic_mode`` the fallback now selects keys from the dense
scores with a stable sort, resolving ties toward the smallest id, and the
no-grad paths use the training path's compressed-first lowering so both
forwards hand FlashMLA identical inputs.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../experimental_attention_variant/csa.py     |  13 +-
 .../csa_utils/fused_sparse_attention.py       |  41 +++++-
 .../test_attention_variant_csa.py             |  67 +++++++++
 .../test_csa_fused_sparse_attention.py        | 127 +++++++++++++++++-
 4 files changed, 240 insertions(+), 8 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py
index 59a5fef0d71..98964ea138d 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa.py
+++ b/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,10 @@ def _forward_fused_indexer_inference(
             deterministic=self.config.deterministic_mode,
         )
         compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1)
+        # Compressed ids first, matching the training path's key order so both
+        # forwards drive FlashMLA identically.
         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 +2659,11 @@ def _forward_fused_indexer_inference_thd(
                 deterministic=self.config.deterministic_mode,
             )
 
+        # The training path lowers these ids compressed-first (the indexer-loss
+        # layout) and compacts them afterwards. Use the same lowering here so a
+        # no-grad forward hands FlashMLA exactly the key order a grad-enabled
+        # forward does; the online-softmax result depends on that order.
+        compressed_first = topk_indices_cmp.shape[-1] > 0
         flat_idxs, flat_tlen, _, _ = thd_layout_kernels.build_attention_indices(
             cu_seqlens_q,
             0,
@@ -2666,11 +2674,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,
diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
index 1be7733436c..af25c172df5 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
+++ b/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.
 
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
index d5873a66fb3..507f547d9a5 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
+++ b/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})
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
index db370fc80cb..3dc1cbc457b 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
+++ b/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,84 @@ 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))
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_real_kernel_fallback_is_repeatable_under_ties(self, reset_lazy_kernel_state):
+        """Identical inputs must select identical ids across calls even when the
+        ReLU-scored indexer leaves many keys tied at exactly zero."""
+        try:
+            dk._ensure_dsa_namespace()
+        except ImportError:
+            pytest.skip("cuDNN Frontend DSA namespace not available")
+        if getattr(dk._DSA, "indexer_forward_top_k_wrapper", None) is not None:
+            pytest.skip("compact indexer wrapper available; this test targets the fallback")
+
+        torch.manual_seed(0)
+        nh, hd, ratio, topk = 8, 64, 4, 8
+        seg_lens = [512, 300, 1024]
+        cu_q = _make_cu_seqlens(seg_lens, device='cuda')
+        cu_kv = _make_cu_seqlens([n // ratio for n in seg_lens], device='cuda')
+        total_q, total_k = int(cu_q[-1]), int(cu_kv[-1])
+        # Mostly negative q.k products collapse to zero after the ReLU, so most
+        # rows have fewer than ``topk`` positive scores and tie at the boundary.
+        q = torch.randn(total_q, nh, hd, dtype=torch.bfloat16, device='cuda').abs()
+        k = torch.randn(total_k, hd, dtype=torch.bfloat16, device='cuda') - 1.0
+        w = torch.rand(total_q, nh, dtype=torch.bfloat16, device='cuda')
+
+        def select():
+            idxs, lengths = indexer_topk(
+                q,
+                k,
+                w,
+                topk=topk,
+                ratio=ratio,
+                cu_seqlens_q=cu_q,
+                cu_seqlens_kv=cu_kv,
+                max_seqlen_q=max(seg_lens),
+                max_seqlen_kv=max(seg_lens) // ratio,
+                deterministic=True,
+            )
+            return idxs, lengths
+
+        first_idxs, first_lengths = select()
+        second_idxs, second_lengths = select()
+        assert torch.equal(first_idxs, second_idxs)
+        assert torch.equal(first_lengths, second_lengths)
+        # Ties are broken toward the smallest key id: each row's selection is
+        # a prefix-closed set under equal scores, so no selected id may exceed
+        # an unselected id of equal score. Check the simplest consequence: the
+        # selected ids are unique and lie within the row's causal range.
+        valid = first_idxs >= 0
+        assert torch.all(first_idxs[valid] < total_k)
+        for row in range(0, total_q, 97):
+            ids = first_idxs[row][first_idxs[row] >= 0].tolist()
+            assert len(ids) == len(set(ids))

From 6e5d9ce5acee0c70f8c59c69084c9d6e037fb962 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Fri, 11 Sep 2026 09:24:30 +0000
Subject: [PATCH 02/10] Name the grad-enabled lowering the no-grad CSA paths
 mirror

Point the two no-grad comments at FusedCSAIndexerSparseAttnFunc.forward and
say why the id sequence matters to FlashMLA, instead of the vaguer
"training path's key order".

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../experimental_attention_variant/csa.py       | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py
index 98964ea138d..ec463ba6b9d 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa.py
@@ -2203,8 +2203,11 @@ def _forward_fused_indexer_inference(
             deterministic=self.config.deterministic_mode,
         )
         compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1)
-        # Compressed ids first, matching the training path's key order so both
-        # forwards drive FlashMLA identically.
+        # 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(
             compress_topk_idxs, window_idxs, batch_size=b, compact=True
         )
@@ -2659,10 +2662,12 @@ def _forward_fused_indexer_inference_thd(
                 deterministic=self.config.deterministic_mode,
             )
 
-        # The training path lowers these ids compressed-first (the indexer-loss
-        # layout) and compacts them afterwards. Use the same lowering here so a
-        # no-grad forward hands FlashMLA exactly the key order a grad-enabled
-        # forward does; the online-softmax result depends on that order.
+        # Same lowering as ``FusedCSAIndexerSparseAttnFunc.forward`` (the
+        # grad-enabled path): ``build_attention_indices`` in indexer-loss mode
+        # emits compressed ids before window ids, and the list is compacted
+        # afterwards. 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.
         compressed_first = topk_indices_cmp.shape[-1] > 0
         flat_idxs, flat_tlen, _, _ = thd_layout_kernels.build_attention_indices(
             cu_seqlens_q,

From 4edec16563a731319e16980af4dff7d86192665e Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Fri, 11 Sep 2026 10:06:59 +0000
Subject: [PATCH 03/10] Describe both build_attention_indices layouts at the
 THD no-grad site

State what the default and indexer-loss layouts return and why the no-grad
path compacts the latter itself, rather than restating the FlashMLA
ordering argument already made for the SBHD site.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../experimental_attention_variant/csa.py        | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py
index ec463ba6b9d..daa50e668b2 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa.py
@@ -2662,12 +2662,16 @@ def _forward_fused_indexer_inference_thd(
                 deterministic=self.config.deterministic_mode,
             )
 
-        # Same lowering as ``FusedCSAIndexerSparseAttnFunc.forward`` (the
-        # grad-enabled path): ``build_attention_indices`` in indexer-loss mode
-        # emits compressed ids before window ids, and the list is compacted
-        # afterwards. 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.
+        # ``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,

From 92473a3d70b6edc66e7c79054a35fed11bb8816f Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Fri, 11 Sep 2026 15:28:47 +0000
Subject: [PATCH 04/10] Drop the real-kernel Top-K repeatability test

With deterministic=True the selection after the cuDNN score kernel is a
plain stable sort, so asserting that two calls agree only re-tests
torch.sort. The mocked bypass test and the module-level key-order test
already cover the integration this was meant to guard.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../test_csa_fused_sparse_attention.py        | 52 -------------------
 1 file changed, 52 deletions(-)

diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
index 3dc1cbc457b..721af5d7310 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
@@ -5491,55 +5491,3 @@ def test_masked_scores_are_never_selected(self):
         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))
-
-    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
-    def test_real_kernel_fallback_is_repeatable_under_ties(self, reset_lazy_kernel_state):
-        """Identical inputs must select identical ids across calls even when the
-        ReLU-scored indexer leaves many keys tied at exactly zero."""
-        try:
-            dk._ensure_dsa_namespace()
-        except ImportError:
-            pytest.skip("cuDNN Frontend DSA namespace not available")
-        if getattr(dk._DSA, "indexer_forward_top_k_wrapper", None) is not None:
-            pytest.skip("compact indexer wrapper available; this test targets the fallback")
-
-        torch.manual_seed(0)
-        nh, hd, ratio, topk = 8, 64, 4, 8
-        seg_lens = [512, 300, 1024]
-        cu_q = _make_cu_seqlens(seg_lens, device='cuda')
-        cu_kv = _make_cu_seqlens([n // ratio for n in seg_lens], device='cuda')
-        total_q, total_k = int(cu_q[-1]), int(cu_kv[-1])
-        # Mostly negative q.k products collapse to zero after the ReLU, so most
-        # rows have fewer than ``topk`` positive scores and tie at the boundary.
-        q = torch.randn(total_q, nh, hd, dtype=torch.bfloat16, device='cuda').abs()
-        k = torch.randn(total_k, hd, dtype=torch.bfloat16, device='cuda') - 1.0
-        w = torch.rand(total_q, nh, dtype=torch.bfloat16, device='cuda')
-
-        def select():
-            idxs, lengths = indexer_topk(
-                q,
-                k,
-                w,
-                topk=topk,
-                ratio=ratio,
-                cu_seqlens_q=cu_q,
-                cu_seqlens_kv=cu_kv,
-                max_seqlen_q=max(seg_lens),
-                max_seqlen_kv=max(seg_lens) // ratio,
-                deterministic=True,
-            )
-            return idxs, lengths
-
-        first_idxs, first_lengths = select()
-        second_idxs, second_lengths = select()
-        assert torch.equal(first_idxs, second_idxs)
-        assert torch.equal(first_lengths, second_lengths)
-        # Ties are broken toward the smallest key id: each row's selection is
-        # a prefix-closed set under equal scores, so no selected id may exceed
-        # an unselected id of equal score. Check the simplest consequence: the
-        # selected ids are unique and lie within the row's causal range.
-        valid = first_idxs >= 0
-        assert torch.all(first_idxs[valid] < total_k)
-        for row in range(0, total_q, 97):
-            ids = first_idxs[row][first_idxs[row] >= 0].tolist()
-            assert len(ids) == len(set(ids))

From c78fef6ce2a3a212f471a2ab09cf0f114d9ff8e9 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Tue, 15 Sep 2026 06:43:42 +0000
Subject: [PATCH 05/10] Sort the deterministic Top-K fallback in row slabs

_stable_topk_indices sorted the whole (rows, sk) score matrix at once, so its
temporaries (masked scores, sorted scores, int64 order) grew with the token
count: about 9 GiB at 32K tokens with compression ratio 4. Rows are
independent, so sort them in slabs sized to keep those buffers under a fixed
budget. The selected ids and their order are unchanged.

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../csa_utils/fused_sparse_attention.py       | 30 +++++++++++++++----
 .../test_csa_fused_sparse_attention.py        |  7 +++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
index af25c172df5..a25fc965a5e 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
@@ -1525,6 +1525,11 @@ def csa_sparse_attn(
 # ---------------------------------------------------------------------------
 
 
+# Upper bound on the transient buffers one ``_stable_topk_indices`` sort may
+# hold at once: the masked score copy, the sorted scores and the int64 order.
+_STABLE_TOPK_SORT_BYTES = 256 << 20
+
+
 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.
 
@@ -1535,6 +1540,9 @@ def _stable_topk_indices(scores: Tensor, seq_lens: Tensor, topk_k: int) -> Tenso
     radix Top-K kernel does not order equal scores, which matters for ReLU-scored
     indexers where many keys share a score of exactly zero.
 
+    Rows are independent, so they are sorted in slabs sized to keep the sort's
+    temporary buffers under ``_STABLE_TOPK_SORT_BYTES`` regardless of ``rows``.
+
     Args:
         scores: ``(rows, sk)`` fp32 indexer scores; masked positions hold ``-inf``.
         seq_lens: ``(rows,)`` int32 number of candidate key columns per row.
@@ -1543,11 +1551,23 @@ def _stable_topk_indices(scores: Tensor, seq_lens: Tensor, topk_k: int) -> Tenso
     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)
+    rows, sk = scores.shape
+    columns = torch.arange(sk, device=scores.device, dtype=seq_lens.dtype)
+    bytes_per_row = sk * (2 * scores.element_size() + 8)
+    slab = max(1, min(rows, _STABLE_TOPK_SORT_BYTES // bytes_per_row))
+    selected = torch.empty(rows, topk_k, dtype=torch.int32, device=scores.device)
+    for start in range(0, rows, slab):
+        stop = min(start + slab, rows)
+        candidates = scores[start:stop].masked_fill(
+            columns.unsqueeze(0) >= seq_lens[start:stop].unsqueeze(1), float("-inf")
+        )
+        sorted_scores, order = torch.sort(candidates, dim=-1, descending=True, stable=True)
+        selected[start:stop] = (
+            order[:, :topk_k]
+            .to(torch.int32)
+            .masked_fill(torch.isneginf(sorted_scores[:, :topk_k]), -1)
+        )
+    return selected
 
 
 def _indexer_topk_core(
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
index 721af5d7310..90e24d8163e 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
@@ -5491,3 +5491,10 @@ def test_masked_scores_are_never_selected(self):
         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))
+
+    def test_row_slabs_match_one_sort(self, monkeypatch):
+        scores = torch.relu(torch.randn(7, 12, generator=torch.Generator().manual_seed(0)))
+        seq_lens = torch.tensor([12, 5, 0, 12, 9, 1, 12], dtype=torch.int32)
+        whole = dk._stable_topk_indices(scores, seq_lens, topk_k=4)
+        monkeypatch.setattr(dk, "_STABLE_TOPK_SORT_BYTES", 1)
+        assert torch.equal(dk._stable_topk_indices(scores, seq_lens, topk_k=4), whole)

From f6802bc324c74d05cd3f9ad2a2d239d3449feb00 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Tue, 15 Sep 2026 14:14:52 +0000
Subject: [PATCH 06/10] Check no-grad and grad CSA forwards agree on the real
 kernels

The stubbed key-order test shows both forwards hand FlashMLA the same id
sequence; this one runs the real cuDNN indexer and FlashMLA sparse kernel
under deterministic_mode and asserts the layer outputs are bit-identical
across grad modes and repeatable within each. It needs SM100 for the pinned
FlashMLA build, so the file is marked for the GB200 unit-test bucket and the
test skips elsewhere.

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../test_attention_variant_csa.py             | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
index 507f547d9a5..ac2819f3463 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
@@ -2381,6 +2381,59 @@ def fake_flash(q, kv, topk_idxs, softmax_scale, d_v=512, **kwargs):
         assert torch.equal(no_grad_lengths, grad_lengths)
         assert torch.equal(no_grad_idxs, grad_idxs)
 
+    # The marker routes this file to the GB200 unit-test bucket, the only CI
+    # runner with SM100 for the pinned FlashMLA kernels.
+    @pytest.mark.launch_on_gb200
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+    def test_thd_real_kernel_no_grad_forward_matches_grad_forward(self, monkeypatch):
+        """With the real cuDNN indexer and FlashMLA kernels, a forward-only pass
+        reproduces the grad-enabled forward bit for bit under deterministic_mode."""
+        from megatron.core.transformer.experimental_attention_variant.csa_utils import (
+            fused_sparse_attention as fsa,
+        )
+
+        if torch.cuda.get_device_capability()[0] < 10:
+            pytest.skip("pinned FlashMLA sparse kernels require SM100+")
+        try:
+            fsa._ensure_dsa_namespace()
+            fsa._ensure_flash_mla()
+        except ImportError as error:
+            pytest.skip(str(error))
+
+        # FlashMLA's sparse kernel is built for the 512-wide DSA value head.
+        config = _make_mla_config(
+            v_head_dim=512, csa_compress_ratios=[4], dsa_indexer_loss_coeff=0.0
+        )
+        config.dsa_kernel_backend = "cudnn"
+        config.deterministic_mode = True
+        monkeypatch.setattr(self, "config", config)
+        csa = self._build_csa(compress_ratio=4).train()
+        torch.manual_seed(7)
+        query, key, value, x, qr, packed = self._make_thd_inputs([96, 40, 72])
+        inputs = dict(
+            query=query,
+            key=key,
+            value=value,
+            attention_mask=None,
+            x=x,
+            qr=qr,
+            packed_seq_params=packed,
+        )
+
+        def forward(grad_enabled):
+            with torch.set_grad_enabled(grad_enabled):
+                return csa(**inputs).detach()
+
+        # First calls compile and autotune; compare only steady-state outputs.
+        forward(False)
+        forward(True)
+        no_grad_1, no_grad_2 = forward(False), forward(False)
+        grad_1, grad_2 = forward(True), forward(True)
+
+        torch.testing.assert_close(no_grad_1, no_grad_2, rtol=0, atol=0)
+        torch.testing.assert_close(grad_1, grad_2, rtol=0, atol=0)
+        torch.testing.assert_close(no_grad_1, grad_1, rtol=0, atol=0)
+
 
 # ===========================================================================
 # _apply_rope direct THD tests (4 corners: ratio={1, >1} × fused={False, True})

From 6fc209b5f3715b3d171192e1992a52e3bce795ca Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Tue, 15 Sep 2026 16:41:08 +0000
Subject: [PATCH 07/10] Use supported shapes for fused CSA parity tests

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../test_attention_variant_csa.py                      | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
index ac2819f3463..8dead6e6083 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py
@@ -2344,6 +2344,8 @@ def fake_flash(q, kv, topk_idxs, softmax_scale, d_v=512, **kwargs):
             "dsa_kernel_backend": "cudnn",
             "deterministic_mode": True,
             "dsa_indexer_loss_coeff": 0.0,
+            "dsa_indexer_n_heads": 64,
+            "dsa_indexer_head_dim": 128,
         }
         saved = {name: getattr(self.config, name) for name in overrides}
         for name, value in overrides.items():
@@ -2402,7 +2404,13 @@ def test_thd_real_kernel_no_grad_forward_matches_grad_forward(self, monkeypatch)
 
         # FlashMLA's sparse kernel is built for the 512-wide DSA value head.
         config = _make_mla_config(
-            v_head_dim=512, csa_compress_ratios=[4], dsa_indexer_loss_coeff=0.0
+            num_layers=1,
+            num_attention_heads=64,
+            v_head_dim=512,
+            csa_compress_ratios=[4],
+            dsa_indexer_n_heads=64,
+            dsa_indexer_head_dim=128,
+            dsa_indexer_loss_coeff=0.0,
         )
         config.dsa_kernel_backend = "cudnn"
         config.deterministic_mode = True

From 003e3341a57d0789328b12953af3c436438520ac Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Tue, 15 Sep 2026 16:41:09 +0000
Subject: [PATCH 08/10] Mask raw THD padding before sparse attention forward

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../csa_utils/fused_sparse_attention.py         |  4 ++++
 .../test_csa_fused_sparse_attention.py          | 17 +++++++++--------
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
index f11ba7ab208..f8b369ca95c 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
@@ -2557,6 +2557,10 @@ def forward(
         # longer needs that segment boundary because FlashMLA's partial
         # ``lse_indexer`` is not used, so compact the full attention set once
         # and share the resulting prefix length with forward and backward.
+        if padding_row_mask is not None:
+            # Padding queries are sink-only in forward; backward supplies a
+            # harmless tile after masking their output gradients.
+            global_idxs.masked_fill_(padding_row_mask.unsqueeze(-1), -1)
         if is_thd and thd_compressed_is_sequence_major:
             logical_window_width = int(thd_window_size)
         else:
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
index 405fc284adf..b5e478c6c21 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
@@ -4871,11 +4871,12 @@ def test_thd_single_segment_matches_sbhd_b1(
     def test_raw_thd_tail_padding_backward_uses_dummy_tile(self, reset_lazy_kernel_state):
         """Raw THD tail padding supplies a harmless tile to DSA backward.
 
-        CUDA graphs keep Q-side tensors at a static capacity, which may leave
-        physical rows beyond the last packed-sequence endpoint. Raw THD
-        lowering emits all ``-1`` indices for those rows, so FlashMLA forward
-        sees ``topk_length == 0``. DSA backward instead requires at least one
-        tile; the fused wrapper must pass length 1 after sanitizing the index.
+        CUDA graphs keep Q-side tensors at a static capacity. Physical sequence
+        lengths include the padding; unpadded lengths identify the real rows.
+        Raw THD lowering emits all ``-1`` indices for padding, so FlashMLA
+        forward sees ``topk_length == 0``. DSA backward instead requires at
+        least one tile; the fused wrapper must pass length 1 after sanitizing
+        the index.
         """
         _skip_if_real_kernels_unavailable(need_flash_mla=True)
         s = self.SHAPES
@@ -4899,8 +4900,8 @@ def test_raw_thd_tail_padding_backward_uses_dummy_tile(self, reset_lazy_kernel_s
         k_indexer = torch.randn(n_comp, s['idx_hd'], dtype=torch.bfloat16, device=dev)
         weights = torch.randn(total_q, s['idx_nh'], dtype=torch.bfloat16, device=dev)
 
-        cu_q = _make_cu_seqlens([real_q], device=dev)
-        cu_q_unpadded = cu_q.clone()
+        cu_q = _make_cu_seqlens([total_q], device=dev)
+        cu_q_unpadded = _make_cu_seqlens([real_q], device=dev)
         cu_kv = _make_cu_seqlens([kv_offset], device=dev)
         cu_comp_idx = _make_cu_seqlens([n_comp], device=dev)
         compressed_kv = kv_full.detach()[kv_offset:]
@@ -4946,7 +4947,7 @@ def capture_sparse_attention_backward(*args, **kwargs):
                 cu_seqlens_q=cu_q,
                 cu_seqlens_kv=cu_kv,
                 cu_seqlens_compressed_idx=cu_comp_idx,
-                max_seqlen_q=real_q,
+                max_seqlen_q=total_q,
                 max_seqlen_compressed_idx=n_comp,
                 compressed_kv=compressed_kv,
                 cu_seqlens_q_unpadded=cu_q_unpadded,

From da81780e02dad91449b2cbf1e4040e075a0f8c48 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Wed, 16 Sep 2026 11:13:28 +0000
Subject: [PATCH 09/10] Reuse negative infinity constant in stable CSA Top-K

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../csa_utils/fused_sparse_attention.py                        | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
index f8b369ca95c..5f20445d259 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
@@ -1533,6 +1533,7 @@ def csa_sparse_attn(
 # Upper bound on the transient buffers one ``_stable_topk_indices`` sort may
 # hold at once: the masked score copy, the sorted scores and the int64 order.
 _STABLE_TOPK_SORT_BYTES = 256 << 20
+_NEG_INF = float("-inf")
 
 
 def _stable_topk_indices(scores: Tensor, seq_lens: Tensor, topk_k: int) -> Tensor:
@@ -1564,7 +1565,7 @@ def _stable_topk_indices(scores: Tensor, seq_lens: Tensor, topk_k: int) -> Tenso
     for start in range(0, rows, slab):
         stop = min(start + slab, rows)
         candidates = scores[start:stop].masked_fill(
-            columns.unsqueeze(0) >= seq_lens[start:stop].unsqueeze(1), float("-inf")
+            columns.unsqueeze(0) >= seq_lens[start:stop].unsqueeze(1), _NEG_INF
         )
         sorted_scores, order = torch.sort(candidates, dim=-1, descending=True, stable=True)
         selected[start:stop] = (

From b57ad25b1f8891d7c2e631b4c64377d0e100f280 Mon Sep 17 00:00:00 2001
From: Bo Li <22713281+bobboli@users.noreply.github.com>
Date: Thu, 17 Sep 2026 07:28:31 +0000
Subject: [PATCH 10/10] Mask CP padding before fused sparse attention forward

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
---
 .../csa_utils/fused_sparse_attention.py            |  3 +++
 .../test_csa_fused_sparse_attention.py             | 14 +++++++++++---
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
index 5f20445d259..9165fe88d57 100644
--- a/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
+++ b/megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py
@@ -3067,6 +3067,9 @@ def forward(
 
         # Preserve the fixed window suffix for the dense teacher before
         # compacting the complete attention index set.
+        if q_padding_mask is not None:
+            # Padding queries attend only to the sink, as in raw THD lowering.
+            topk_idxs.masked_fill_(q_padding_mask.unsqueeze(-1), -1)
         if logical_window_width is None:
             logical_window_width = topk_idxs.shape[-1] - indexer_topk
         window_topk_idxs = topk_idxs[:, indexer_topk : indexer_topk + int(logical_window_width)]
diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
index b5e478c6c21..3004d54086a 100644
--- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
+++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_sparse_attention.py
@@ -2791,7 +2791,6 @@ def dense_indexer_backward_wrapper(q, w, k, *args, **kwargs):
 
     def test_backward_reuses_compact_indices_and_length(self, monkeypatch):
         inputs = self._inputs()
-        inputs['topk_idxs'][1] = -1
         for name in ('query', 'kv_full', 'attn_sink', 'q_indexer', 'k_indexer', 'weights'):
             inputs[name].requires_grad_(True)
 
@@ -2799,8 +2798,10 @@ def test_backward_reuses_compact_indices_and_length(self, monkeypatch):
         q_padding_mask = torch.tensor([False, True, False, False])
         seen = {}
 
-        def fake_flash(query, *args, **kwargs):
-            del args, kwargs
+        def fake_flash(query, kv_full, topk_idxs, softmax_scale, **kwargs):
+            del kv_full, softmax_scale
+            seen['forward_topk'] = topk_idxs.detach().clone()
+            seen['forward_topk_length'] = kwargs['topk_length'].detach().clone()
             return torch.zeros_like(query), torch.full((total_q, num_heads), 3.0), None
 
         class FakeDSA:
@@ -2851,6 +2852,13 @@ def sparse_attention_backward_wrapper(
         )
         (output.sum() + loss).backward()
 
+        torch.testing.assert_close(
+            seen['forward_topk'],
+            torch.tensor([[4, 0], [-1, -1], [5, 2], [5, 3]], dtype=torch.int32),
+        )
+        torch.testing.assert_close(
+            seen['forward_topk_length'], torch.tensor([2, 0, 2, 2], dtype=torch.int32)
+        )
         torch.testing.assert_close(
             seen['topk'], torch.tensor([[4, 0], [0, 0], [5, 2], [5, 3]], dtype=torch.int32)
         )