[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,
@@ -2238,8 +2239,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")
 
@@ -2692,6 +2698,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,
@@ -2702,11 +2719,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

@@ -1530,6 +1530,52 @@ 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:
+    """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.
+
+    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.
+        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.
+    """
+    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), _NEG_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
+
+
 # The balanced CP path fails closed above the shared limit before reaching this
 # compatibility warning; see cp_utils.compute_cp_indexer_topk.
 
@@ -1847,10 +1893,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(
@@ -1923,8 +1972,9 @@ def indexer_topk(
             available. False retains dense scoring for the balanced indexer's
             existing unpadded synthetic layouts.
         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.
 
@@ -2508,6 +2558,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:
@@ -3013,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)]

tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py

@@ -2314,6 +2314,134 @@ 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,
+            "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():
+            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)
+
+    # 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(
+            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
+        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})