diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index e20bfb85833..51b3e2b3631 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, @@ -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, 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 dd6a2e32988..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 @@ -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)] 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..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 @@ -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}) 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 ffee14fa9c9..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) ) @@ -4310,7 +4318,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) @@ -4320,6 +4328,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() @@ -4827,11 +4879,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 @@ -4855,8 +4908,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:] @@ -4902,7 +4955,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, @@ -5566,3 +5619,39 @@ 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)) + + 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)