From 3b6e69a8460f68229dbbdb3702f4223c97faad61 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Fri, 7 Aug 2026 08:28:22 -0700 Subject: [PATCH 1/5] Fix fused CSA indexer loss denominator Signed-off-by: Hongxiao Bai --- .../csa_utils/fused_sparse_attention.py | 295 +++++++++++++- .../test_csa_fused_sparse_attention.py | 384 +++++++++++++++--- .../test_dsv4_hybrid_native_parity.py | 98 ++--- 3 files changed, 632 insertions(+), 145 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 250aa713167..92572b2c533 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 @@ -36,6 +36,8 @@ _flash_mla_sparse_fwd = None _DSA = None +_CSA_TEACHER_LSE_CHUNK_MAX_BYTES = 1024 * 1024 * 1024 + def _ensure_flash_mla(): """Lazily import the FlashMLA sparse-forward kernel. @@ -183,6 +185,234 @@ def batch_of_row(cu_seqlens_q: Tensor, total_q: Optional[int] = None) -> Tensor: ) +def _teacher_lse_chunk_rows( + num_rows: int, batch: int, heads: int, keys: int, extra_bytes_per_row: int = 0 +) -> int: + """Bound temporary storage used to recompute a teacher LSE.""" + bytes_per_row = max(1, batch) * max(1, heads) * max(1, keys) * 4 + extra_bytes_per_row + return max(1, min(num_rows, _CSA_TEACHER_LSE_CHUNK_MAX_BYTES // bytes_per_row)) + + +@torch.no_grad() +def _compute_csa_non_compressed_lse( + query: Tensor, kv_full: Tensor, attn_sink: Tensor, window_indices: Tensor, softmax_scale: float +) -> Tensor: + """Return per-head LSE for the CSA sliding-window and sink mass. + + Inputs use FlashMLA's flat layout and global KV indices. The returned + tensor has shape ``[total_q, heads]``. Invalid indices contribute no + mass, while ``attn_sink`` always participates in the denominator. + """ + if query.ndim != 3: + raise ValueError(f"query must have shape [total_q, heads, dim], got {tuple(query.shape)}") + if kv_full.ndim != 2 or kv_full.shape[1] != query.shape[2]: + raise ValueError( + "kv_full must have shape [total_kv, dim] matching query, " + f"got {tuple(kv_full.shape)} and {tuple(query.shape)}" + ) + if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: + raise ValueError( + "window_indices must have shape [total_q, window], " + f"got {tuple(window_indices.shape)} for total_q={query.shape[0]}" + ) + if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: + raise ValueError( + f"attn_sink must contain {query.shape[1]} head values, got {tuple(attn_sink.shape)}" + ) + if not (query.device == kv_full.device == attn_sink.device == window_indices.device): + raise ValueError("query, kv_full, attn_sink, and window_indices must share a device") + if kv_full.shape[0] == 0 and window_indices.numel() > 0: + raise ValueError("window indices cannot address an empty KV tensor") + + total_q, num_heads, head_dim = query.shape + window_width = window_indices.shape[1] + # In addition to the FP32 score matrix, window recomputation materializes + # gathered KV in its source dtype and an FP32 cast consumed by einsum. + gathered_kv_bytes = window_width * head_dim * (kv_full.element_size() + 4) + query_float_bytes = num_heads * head_dim * 4 + chunk_rows = _teacher_lse_chunk_rows( + total_q, 1, num_heads, window_width, gathered_kv_bytes + query_float_bytes + ) + sink = attn_sink.detach().float().view(1, num_heads) + lse_chunks = [] + for start in range(0, total_q, chunk_rows): + end = min(start + chunk_rows, total_q) + indices = window_indices[start:end].to(dtype=torch.int64) + valid = (indices >= 0) & (indices < kv_full.shape[0]) + if window_width == 0: + window_lse = torch.full( + (end - start, num_heads), float("-inf"), device=query.device, dtype=torch.float32 + ) + else: + safe_indices = indices.clamp(min=0, max=max(kv_full.shape[0] - 1, 0)) + gathered_kv = ( + kv_full.detach() + .index_select(0, safe_indices.reshape(-1)) + .reshape(end - start, window_width, head_dim) + ) + window_logits = torch.einsum( + "rhd,rkd->rhk", query[start:end].detach().float(), gathered_kv.float() + ) + window_logits = (window_logits * softmax_scale).masked_fill( + ~valid.unsqueeze(1), float("-inf") + ) + window_lse = torch.logsumexp(window_logits, dim=-1) + lse_chunks.append(torch.logaddexp(window_lse, sink)) + + if not lse_chunks: + return torch.empty((0, num_heads), device=query.device, dtype=torch.float32) + return torch.cat(lse_chunks, dim=0) + + +@torch.no_grad() +def _compute_dense_csa_teacher_lse( + query: Tensor, + compressed_kv: Tensor, + non_compressed_lse: Tensor, + softmax_scale: float, + ratio: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Return the full CSA teacher LSE for dense indexer loss. + + ``non_compressed_lse`` contains sliding-window plus sink mass. This + helper recomputes the per-head LSE over every causally valid compressed + key and combines the two masses. SBHD callers pass BSHD/BKHD tensors; + packed callers pass THD tensors plus cumulative-length metadata. + """ + if ratio <= 0: + raise ValueError(f"ratio must be positive, got {ratio}") + + if query.ndim == 4: + b, sq, num_heads, head_dim = query.shape + if compressed_kv.ndim == 4: + if compressed_kv.shape[2] != 1: + raise ValueError("CSA dense teacher expects exactly one compressed KV head") + compressed_kv = compressed_kv.squeeze(2) + if compressed_kv.ndim != 3 or compressed_kv.shape[0] != b: + raise ValueError( + "SBHD compressed_kv must have shape [batch, seqlen_k, dim], " + f"got {tuple(compressed_kv.shape)}" + ) + if compressed_kv.shape[2] != head_dim: + raise ValueError("query and compressed_kv head dimensions must match") + if tuple(non_compressed_lse.shape) != (b, sq, num_heads): + raise ValueError( + "SBHD non_compressed_lse must have shape " + f"[{b}, {sq}, {num_heads}], got {tuple(non_compressed_lse.shape)}" + ) + + sk = compressed_kv.shape[1] + compressed_lse = torch.empty_like(non_compressed_lse, dtype=torch.float32) + compressed_kv_float = compressed_kv.detach().float() + key_positions = torch.arange(sk, device=query.device).view(1, 1, 1, sk) + chunk_rows = _teacher_lse_chunk_rows(sq, b, num_heads, sk) + for start in range(0, sq, chunk_rows): + end = min(start + chunk_rows, sq) + scores = torch.einsum( + "bqhd,bkd->bqhk", query[:, start:end].detach().float(), compressed_kv_float + ) + scores.mul_(softmax_scale) + visible_k = torch.div( + torch.arange(start + 1, end + 1, device=query.device), ratio, rounding_mode="floor" + ).view(1, end - start, 1, 1) + scores.masked_fill_(key_positions >= visible_k, float("-inf")) + compressed_lse[:, start:end] = torch.logsumexp(scores, dim=-1) + elif query.ndim == 3: + total_q, num_heads, head_dim = query.shape + if compressed_kv.ndim == 3: + if compressed_kv.shape[1] != 1: + raise ValueError("CSA dense teacher expects exactly one compressed KV head") + compressed_kv = compressed_kv.squeeze(1) + if compressed_kv.ndim != 2 or compressed_kv.shape[1] != head_dim: + raise ValueError( + "THD compressed_kv must have shape [total_k, dim] matching query, " + f"got {tuple(compressed_kv.shape)} and {tuple(query.shape)}" + ) + if tuple(non_compressed_lse.shape) != (total_q, num_heads): + raise ValueError( + "THD non_compressed_lse must have shape " + f"[{total_q}, {num_heads}], got {tuple(non_compressed_lse.shape)}" + ) + if any( + value is None for value in (cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv) + ): + raise ValueError("THD dense teacher LSE requires cumulative lengths and max lengths") + if cu_seqlens_q.shape != cu_seqlens_kv.shape: + raise ValueError("THD query and compressed-KV cumulative lengths must have equal shape") + + num_sequences = cu_seqlens_q.shape[0] - 1 + if q_causal_offsets is None: + q_causal_offsets = torch.zeros( + num_sequences, device=query.device, dtype=cu_seqlens_q.dtype + ) + if q_causal_offsets.shape != (num_sequences,): + raise ValueError( + f"q_causal_offsets must have shape [{num_sequences}], " + f"got {tuple(q_causal_offsets.shape)}" + ) + + max_q = int(max_seqlen_q) + max_k = int(max_seqlen_kv) + q_padded = query.new_zeros((num_sequences, max_q, num_heads, head_dim)) + k_padded = compressed_kv.new_zeros((num_sequences, max_k, head_dim)) + + q_rows = torch.arange(total_q, device=query.device, dtype=torch.int64) + q_batch = batch_of_row(cu_seqlens_q, total_q=total_q) + q_positions = q_rows - cu_seqlens_q[q_batch].to(torch.int64) + q_padded[q_batch, q_positions] = query + + total_k = compressed_kv.shape[0] + k_rows = torch.arange(total_k, device=query.device, dtype=torch.int64) + k_batch = batch_of_row(cu_seqlens_kv, total_q=total_k) + k_positions = k_rows - cu_seqlens_kv[k_batch].to(torch.int64) + k_padded[k_batch, k_positions] = compressed_kv + + q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(torch.int64) + k_lens = (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).to(torch.int64) + key_positions = torch.arange(max_k, device=query.device).view(1, 1, 1, max_k) + compressed_lse_padded = torch.full( + (num_sequences, max_q, num_heads), + float("-inf"), + device=query.device, + dtype=torch.float32, + ) + k_padded_float = k_padded.detach().float() + q_causal_offsets_i64 = q_causal_offsets.to(torch.int64) + chunk_rows = _teacher_lse_chunk_rows(max_q, num_sequences, num_heads, max_k) + for start in range(0, max_q, chunk_rows): + end = min(start + chunk_rows, max_q) + scores = torch.einsum( + "bqhd,bkd->bqhk", q_padded[:, start:end].detach().float(), k_padded_float + ) + scores.mul_(softmax_scale) + local_q = torch.arange(start, end, device=query.device, dtype=torch.int64) + visible_k = torch.div( + local_q.view(1, -1) + q_causal_offsets_i64.view(-1, 1) + 1, + ratio, + rounding_mode="floor", + ).clamp(max=max_k) + valid = key_positions < torch.minimum(visible_k, k_lens.view(-1, 1)).view( + num_sequences, end - start, 1, 1 + ) + valid = valid & (local_q.view(1, -1, 1, 1) < q_lens.view(-1, 1, 1, 1)) + scores.masked_fill_(~valid, float("-inf")) + compressed_lse_padded[:, start:end] = torch.logsumexp(scores, dim=-1) + compressed_lse = compressed_lse_padded[q_batch, q_positions] + else: + raise ValueError( + "query must be BSHD [batch, seqlen_q, heads, dim] or " + f"THD [total_q, heads, dim], got {tuple(query.shape)}" + ) + + return torch.logaddexp(non_compressed_lse.detach().float(), compressed_lse) + + def local_to_global_flat( local_idxs: Tensor, batch_size: int, @@ -1142,14 +1372,18 @@ def forward( else: q_flat = query.reshape(sq * b, np_, d) kv_flat = kv_full.reshape(skv * b, d) - out_flat, lse, lse_indexer = _csa_fwd_flash_mla( + # The partial indexer LSE excludes the window (and FlashMLA excludes + # the sink from every returned LSE), so it is not a valid teacher + # denominator. Sparse loss uses the full LSE below; dense loss + # recomputes the all-compressed denominator. + out_flat, lse, _ = _csa_fwd_flash_mla( q_flat, kv_flat, global_idxs, softmax_scale, attn_sink=attn_sink, topk_length=None, - indexer_topk=indexer_topk, + indexer_topk=0, ) # ---- 4b. Derive padding-row mask for loss exclusion. ----------------- @@ -1179,11 +1413,18 @@ def forward( assert compressed_kv is not None, "compressed_kv is required for THD" q_attn_det = query.detach() k_attn_compressed_det = compressed_kv.detach() - lse_indexer_det = lse_indexer.detach() + sparse_teacher_lse = torch.logaddexp( + lse.detach().float(), attn_sink.detach().float().view(1, np_) + ) else: q_attn_det = query.detach().permute(1, 0, 2, 3).contiguous() k_attn_compressed_det = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous() - lse_indexer_det = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2) + sparse_teacher_lse = ( + torch.logaddexp(lse.detach().float(), attn_sink.detach().float().view(1, np_)) + .reshape(sq, b, np_) + .permute(1, 0, 2) + .contiguous() + ) # Invalidate padding rows for the loss/backward path. The sparse # attention (steps 3-4) has already built global_idxs from the @@ -1219,7 +1460,7 @@ def forward( target = _compute_attn_target( q_attn_det, k_attn_compressed_det, - lse_indexer_det, + sparse_teacher_lse, topk_for_target, softmax_scale, qhead_per_kv_head=np_, @@ -1245,10 +1486,27 @@ def forward( max_seqlen_q=int(max_seqlen_q), max_seqlen_kv=int(max_seqlen_compressed_idx), ) + non_compressed_lse_flat = _compute_csa_non_compressed_lse( + q_flat, kv_flat, attn_sink, global_idxs[:, indexer_topk:], softmax_scale + ) + if is_thd: + non_compressed_lse = non_compressed_lse_flat + else: + non_compressed_lse = ( + non_compressed_lse_flat.reshape(sq, b, np_).permute(1, 0, 2).contiguous() + ) + dense_teacher_lse = _compute_dense_csa_teacher_lse( + q_attn_det, + k_attn_compressed_det, + non_compressed_lse, + softmax_scale, + ratio, + **dense_attn_kwargs, + ) attn_score, attn_l1norm = _compute_dense_attn_score( q_attn_det, k_attn_compressed_det.unsqueeze(k_unsqueeze_dim), - lse_indexer_det, + dense_teacher_lse, qhead_per_kv_head=np_, softmax_scale=softmax_scale, ratio=ratio, @@ -1533,14 +1791,16 @@ def forward( total_comp = k_indexer.shape[0] indexer_topk = indexer_topk_idxs.shape[-1] - out_flat, lse, lse_indexer = _csa_fwd_flash_mla( + # Do not request FlashMLA's partial indexer LSE: it omits both the + # window and sink masses required by the CSA teacher. + out_flat, lse, _ = _csa_fwd_flash_mla( query, kv_full, topk_idxs, softmax_scale, attn_sink=attn_sink, topk_length=None, - indexer_topk=indexer_topk, + indexer_topk=0, ) bwd_loss_coeff = loss_coeff * total_q / loss_divisor @@ -1564,7 +1824,7 @@ def forward( target = _compute_attn_target( query.detach(), compressed_kv.detach(), - lse_indexer.detach(), + torch.logaddexp(lse.detach().float(), attn_sink.detach().float().view(1, np_)), indexer_topk_idxs_for_loss, softmax_scale, qhead_per_kv_head=np_, @@ -1610,10 +1870,25 @@ def forward( if q_padding_mask is not None: index_score = index_score.masked_fill(q_padding_mask.unsqueeze(-1), float("-inf")) index_lse = index_lse.masked_fill(q_padding_mask, float("-inf")) + non_compressed_lse = _compute_csa_non_compressed_lse( + query, kv_full, attn_sink, topk_idxs[:, indexer_topk:], softmax_scale + ) + dense_teacher_lse = _compute_dense_csa_teacher_lse( + query.detach(), + compressed_kv.detach(), + non_compressed_lse, + softmax_scale, + ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + ) attn_score, attn_l1norm = _compute_dense_attn_score( query.detach(), compressed_kv.detach().unsqueeze(1), - lse_indexer.detach(), + dense_teacher_lse, qhead_per_kv_head=np_, softmax_scale=softmax_scale, ratio=ratio, 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 751fe5bd7e8..2921da2c465 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 @@ -35,6 +35,8 @@ CSASparseAttnFunc, FusedCSAIndexerSparseAttnFromTopkFunc, FusedCSAIndexerSparseAttnFunc, + _compute_csa_non_compressed_lse, + _compute_dense_csa_teacher_lse, _csa_fwd_flash_mla, _ensure_dsa_namespace, _ensure_flash_mla, @@ -289,6 +291,107 @@ def fake_compactify(global_idxs): ), "(b) length tensor differs between CPU fallback and cuDNN kernel" +# --------------------------------------------------------------------------- +# CSA teacher LSE helpers +# --------------------------------------------------------------------------- + + +class TestCsaTeacherLse: + """Full-denominator teacher LSE helpers; CPU execution is sufficient.""" + + def test_non_compressed_lse_includes_window_and_sink(self): + query = torch.tensor([[[1.0, 0.0], [0.0, 1.0]], [[0.5, 0.5], [1.0, -1.0]]]) + kv_full = torch.tensor([[1.0, 0.0], [0.0, 2.0], [2.0, -1.0]]) + window_indices = torch.tensor([[0, -1], [1, 2]], dtype=torch.int32) + sink = torch.tensor([0.25, -0.5]) + + actual = _compute_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=0.5 + ) + + gathered = kv_full.index_select(0, window_indices.clamp_min(0).reshape(-1)).reshape(2, 2, 2) + logits = torch.einsum("rhd,rkd->rhk", query, gathered) * 0.5 + logits = logits.masked_fill((window_indices < 0).unsqueeze(1), float("-inf")) + expected = torch.logaddexp(torch.logsumexp(logits, dim=-1), sink.view(1, -1)) + torch.testing.assert_close(actual, expected) + + all_invalid = torch.full_like(window_indices, -1) + sink_only = _compute_csa_non_compressed_lse( + query, kv_full, sink, all_invalid, softmax_scale=0.5 + ) + torch.testing.assert_close(sink_only, sink.view(1, -1).expand_as(sink_only)) + + def test_dense_sbhd_lse_adds_every_causal_compressed_key(self): + query = torch.tensor( + [ + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 1.0], [1.0, -1.0]], + [[0.5, 1.0], [-1.0, 0.5]], + [[2.0, -1.0], [0.25, 0.75]], + ] + ] + ) + compressed_kv = torch.tensor([[[1.0, 0.0], [0.0, 2.0]]]) + non_compressed_lse = torch.tensor([[[0.25, -0.5], [0.5, 0.75], [-0.25, 0.0], [1.0, -1.0]]]) + + actual = _compute_dense_csa_teacher_lse( + query, compressed_kv, non_compressed_lse, softmax_scale=0.5, ratio=2 + ) + + expected = torch.empty_like(actual) + for q_pos in range(query.shape[1]): + visible = (q_pos + 1) // 2 + if visible == 0: + compressed_lse = torch.full((query.shape[2],), float("-inf")) + else: + logits = ( + torch.einsum("hd,kd->hk", query[0, q_pos], compressed_kv[0, :visible]) * 0.5 + ) + compressed_lse = torch.logsumexp(logits, dim=-1) + expected[0, q_pos] = torch.logaddexp(non_compressed_lse[0, q_pos], compressed_lse) + torch.testing.assert_close(actual, expected) + + def test_dense_thd_lse_honors_segment_offsets(self): + query = torch.tensor( + [[[1.0], [2.0]], [[2.0], [1.0]], [[3.0], [0.5]], [[1.5], [-1.0]], [[0.25], [4.0]]] + ) + compressed_kv = torch.tensor([[1.0], [2.0], [3.0]]) + cu_q = torch.tensor([0, 2, 5], dtype=torch.int32) + cu_k = torch.tensor([0, 1, 3], dtype=torch.int32) + q_offsets = torch.tensor([1, 0], dtype=torch.int32) + non_compressed_lse = torch.arange(10, dtype=torch.float32).reshape(5, 2) * 0.1 + + actual = _compute_dense_csa_teacher_lse( + query, + compressed_kv, + non_compressed_lse, + softmax_scale=0.25, + ratio=2, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_k, + max_seqlen_q=3, + max_seqlen_kv=2, + q_causal_offsets=q_offsets, + ) + + expected = torch.empty_like(actual) + for row in range(query.shape[0]): + batch = 0 if row < 2 else 1 + q_pos = row - int(cu_q[batch]) + visible = min( + (q_pos + int(q_offsets[batch]) + 1) // 2, int(cu_k[batch + 1] - cu_k[batch]) + ) + keys = compressed_kv[int(cu_k[batch]) : int(cu_k[batch]) + visible] + if visible == 0: + compressed_lse = torch.full((query.shape[1],), float("-inf")) + else: + logits = torch.einsum("hd,kd->hk", query[row], keys) * 0.25 + compressed_lse = torch.logsumexp(logits, dim=-1) + expected[row] = torch.logaddexp(non_compressed_lse[row], compressed_lse) + torch.testing.assert_close(actual, expected) + + # --------------------------------------------------------------------------- # _kl_loss_from_target_predict # --------------------------------------------------------------------------- @@ -1214,7 +1317,7 @@ def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel # ---- (a) forward pass-through (no grads needed) ------------------ inputs = self._make_inputs() - _, flash_stub_a = _install_full_dsa_mock( + fake_dsa_a, flash_stub_a = _install_full_dsa_mock( b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] ) output_a, _ = fused_csa_indexer_sparse_attn( @@ -1233,6 +1336,16 @@ def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel s['sq'], s['b'], s['np_'] * s['d'] ) assert torch.equal(output_a, expected_a), "(a) forward value pass-through" + target_lse = fake_dsa_a.sparse_attn_score_recompute_wrapper.call_args.args[2] + expected_target_lse = ( + torch.logaddexp(flash_stub_a.last_lse, inputs['attn_sink'].view(1, s['np_'])) + .reshape(s['sq'], s['b'], s['np_']) + .permute(1, 0, 2) + ) + torch.testing.assert_close(target_lse, expected_target_lse) + assert ( + flash_stub_a.call_args.kwargs['indexer_topk'] == 0 + ), "(a) partial indexer LSE should not be requested" # ---- (b) backward grad propagation ------------------------------- dk._DSA = None # fresh mocks @@ -1419,7 +1532,7 @@ def target_score_fn(B, S, K, dev): ), f"got {indexer_loss.item()}, expected {expected}" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state): + def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state, monkeypatch): """Combined coverage for the dense-loss path's two non-numerical properties (assertion blocks self-label on failure): @@ -1439,6 +1552,15 @@ def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state softmax_scale = 0.5 idx_scale = 0.125 loss_coeff = 1.0 + expected_teacher_lse = torch.full( + (s['b'], s['sq'], s['np_']), 37.0, dtype=torch.float32, device='cuda' + ) + + def fake_dense_teacher_lse(*args, **kwargs): + del args, kwargs + return expected_teacher_lse + + monkeypatch.setattr(dk, '_compute_dense_csa_teacher_lse', fake_dense_teacher_lse) # ---- (a) forward kernel selection + arg shapes ------------------- inputs_a = self._make_inputs() @@ -1472,6 +1594,9 @@ def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state s['d'], ), "(a) dense attn score: k shape (h_kv=1)" assert lse_arg.shape == (s['b'], s['sq'], s['np_']), "(a) dense attn score: lse shape" + assert torch.equal( + lse_arg, expected_teacher_lse + ), "(a) dense path did not use full teacher LSE" assert sm_arg == softmax_scale, "(a) dense attn score: positional softmax_scale" assert attn_call.kwargs['qhead_per_kv_head'] == s['np_'] assert attn_call.kwargs['ratio'] == ratio @@ -1529,6 +1654,170 @@ def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state ), f"(b) {name}: grad does not equal full({value})" +class TestFusedIndexerSparseAttnFromTopk: + """Full-denominator plumbing for the caller-supplied top-k fused path.""" + + @staticmethod + def _inputs(): + total_q, num_heads, head_dim = 4, 2, 2 + idx_heads, idx_dim, total_comp = 2, 3, 2 + return dict( + query=torch.randn(total_q, num_heads, head_dim), + kv_full=torch.randn(6, head_dim), + attn_sink=torch.tensor([0.25, -0.5]), + topk_idxs=torch.tensor([[4, 0], [4, 1], [5, 2], [5, 3]], dtype=torch.int32), + q_indexer=torch.randn(total_q, idx_heads, idx_dim), + k_indexer=torch.randn(total_comp, idx_dim), + weights=torch.randn(total_q, idx_heads), + indexer_topk_idxs=torch.tensor([[0], [0], [1], [1]], dtype=torch.int32), + compressed_kv=torch.randn(total_comp, head_dim), + ) + + def test_sparse_loss_uses_full_flash_lse_plus_sink(self, monkeypatch): + inputs = self._inputs() + total_q, num_heads, _ = inputs['query'].shape + full_lse = torch.arange(total_q * num_heads, dtype=torch.float32).reshape( + total_q, num_heads + ) + partial_lse = full_lse + 100.0 + seen = {} + + def fake_flash( + query, + kv_full, + topk_idxs, + softmax_scale, + d_v=512, + attn_sink=None, + topk_length=None, + indexer_topk=0, + ): + del kv_full, topk_idxs, softmax_scale, d_v, attn_sink, topk_length + seen['indexer_topk'] = indexer_topk + return torch.zeros_like(query), full_lse, partial_lse + + class FakeDSA: + @staticmethod + def sparse_indexer_score_recompute_wrapper(q, k, w, topk, **kwargs): + del q, k, w, kwargs + return {'predict': torch.ones_like(topk, dtype=torch.float32)} + + @staticmethod + def sparse_attn_score_recompute_wrapper(q, k, lse, topk, scale, **kwargs): + del q, k, scale, kwargs + seen['teacher_lse'] = lse.detach().clone() + return {'target': torch.ones_like(topk, dtype=torch.float32)} + + @staticmethod + def indexer_backward_wrapper(q, w, k, *args, **kwargs): + del args, kwargs + return { + 'd_index_q': torch.zeros_like(q), + 'd_weights': torch.zeros_like(w), + 'd_index_k': torch.zeros_like(k), + } + + monkeypatch.setattr(dk, '_ensure_dsa_namespace', lambda: None) + monkeypatch.setattr(dk, '_csa_fwd_flash_mla', fake_flash) + monkeypatch.setattr(dk, '_DSA', FakeDSA) + + FusedCSAIndexerSparseAttnFromTopkFunc.apply( + *inputs.values(), + 1.0, + 1.0, + 1.0, + float(total_q), + True, + 2, + total_q, + ( + torch.tensor([0, total_q], dtype=torch.int32), + torch.tensor([0, inputs['k_indexer'].shape[0]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + ), + None, + ) + + expected = torch.logaddexp(full_lse, inputs['attn_sink'].view(1, num_heads)).unsqueeze(0) + torch.testing.assert_close(seen['teacher_lse'], expected) + assert seen['indexer_topk'] == 0 + + def test_dense_loss_passes_recomputed_full_teacher_lse(self, monkeypatch): + inputs = self._inputs() + total_q, num_heads, _ = inputs['query'].shape + max_seqlen_k = 2 + sentinel_lse = torch.full((total_q, num_heads), 37.0) + seen = {} + + def fake_flash(query, *args, **kwargs): + del args, kwargs + return torch.zeros_like(query), torch.zeros(total_q, num_heads), None + + def fake_non_compressed(query, kv_full, sink, window_indices, scale): + del query, kv_full, sink, scale + torch.testing.assert_close(window_indices, inputs['topk_idxs'][:, 1:]) + return torch.full((total_q, num_heads), 5.0) + + def fake_dense_teacher(query, compressed_kv, non_compressed_lse, *args, **kwargs): + del query, compressed_kv, args, kwargs + torch.testing.assert_close(non_compressed_lse, torch.full((total_q, num_heads), 5.0)) + seen['dense_teacher_called'] = True + return sentinel_lse + + class FakeDSA: + @staticmethod + def dense_indexer_score_recompute_wrapper(q, k, w, **kwargs): + del q, k, w, kwargs + return { + 'out': torch.zeros(total_q, max_seqlen_k, dtype=torch.float32), + 'denom': torch.zeros(total_q, dtype=torch.float32), + } + + @staticmethod + def dense_attn_score_recompute_wrapper(q, k, lse, scale, **kwargs): + del q, k, scale, kwargs + seen['teacher_lse'] = lse.detach().clone() + return { + 'out': torch.full((total_q, max_seqlen_k), 0.5), + 'denom': torch.ones(total_q), + } + + @staticmethod + def dense_indexer_backward_wrapper(q, w, k, *args, **kwargs): + del args, kwargs + return { + 'd_index_q': torch.zeros_like(q), + 'd_weights': torch.zeros_like(w), + 'd_index_k': torch.zeros_like(k), + } + + monkeypatch.setattr(dk, '_ensure_dsa_namespace', lambda: None) + monkeypatch.setattr(dk, '_csa_fwd_flash_mla', fake_flash) + monkeypatch.setattr(dk, '_compute_csa_non_compressed_lse', fake_non_compressed) + monkeypatch.setattr(dk, '_compute_dense_csa_teacher_lse', fake_dense_teacher) + monkeypatch.setattr(dk, '_DSA', FakeDSA) + + FusedCSAIndexerSparseAttnFromTopkFunc.apply( + *inputs.values(), + 1.0, + 1.0, + 1.0, + float(total_q), + False, + 2, + total_q, + ( + torch.tensor([0, total_q], dtype=torch.int32), + torch.tensor([0, max_seqlen_k], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + ), + None, + ) + + assert seen['dense_teacher_called'] + torch.testing.assert_close(seen['teacher_lse'], sentinel_lse) + + # --------------------------------------------------------------------------- # Real-kernel parity tests (cuDNN + optional FlashMLA) # --------------------------------------------------------------------------- @@ -2271,10 +2560,8 @@ class TestRealKernelFusedIndexerSparseAttn: def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): """Real dense path's KL loss value matches the all-PyTorch reference - on the same inputs. The reference uses an analytical - ``logsumexp(QK*scale, ratio mask)`` for ``lse_indexer`` (FlashMLA - emits its own internal lse_indexer that differs slightly), so the - tolerance is wider than for the kernel-only ``KLLossDense`` test. + on the same inputs, including the sliding-window and sink mass in + the dense teacher denominator. """ _skip_if_real_kernels_unavailable(need_flash_mla=True) s = self.SHAPES @@ -2322,41 +2609,18 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): q_attn_bshd = query.permute(1, 0, 2, 3).contiguous().float() k_attn_bsd = kv_full[kv_offset:].permute(1, 0, 2).contiguous().float() - # The reference uses FlashMLA's actual normalization because it - # includes the selected top-k positions and the per-head sink term. - from megatron.core.transformer.experimental_attention_variant.csa_utils.fused_sparse_attention import ( - _csa_fwd_flash_mla, - _indexer_topk_core, - ) - - # Run indexer + FlashMLA to capture the same ``lse_indexer`` the fused - # path consumes internally. - effective_topk = min(s['indexer_topk'], s['n_comp']) - q_idx_bshd_bf = q_indexer.permute(1, 0, 2, 3).contiguous() - k_idx_bsd_bf = k_indexer.permute(1, 0, 2).contiguous() - w_bsh_bf = weights.permute(1, 0, 2).contiguous() - if s['indexer_softmax_scale'] != 1.0: - w_bsh_scaled_bf = (w_bsh_bf.float() * s['indexer_softmax_scale']).to(w_bsh_bf.dtype) - else: - w_bsh_scaled_bf = w_bsh_bf - topk_indices_cmp, _, _ = _indexer_topk_core( - q_idx_bshd_bf, k_idx_bsd_bf, w_bsh_scaled_bf, effective_topk, s['ratio'] - ) - compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1) - combined_local = torch.cat([compress_topk_idxs, win_idxs], dim=-1) - global_idxs = local_to_global_flat(combined_local, s['b']) + window_global_idxs = local_to_global_flat(win_idxs, s['b']) q_flat = query.reshape(s['sq'] * s['b'], s['np_'], s['d']) kv_flat = kv_full.reshape(s['skv'] * s['b'], s['d']) - _, _, lse_indexer = _csa_fwd_flash_mla( - q_flat, - kv_flat, - global_idxs, - s['softmax_scale'], - attn_sink=attn_sink, - topk_length=None, - indexer_topk=effective_topk, + non_compressed_lse = _compute_csa_non_compressed_lse( + q_flat, kv_flat, attn_sink, window_global_idxs, s['softmax_scale'] + ) + non_compressed_lse = ( + non_compressed_lse.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2).contiguous() + ) + teacher_lse = _compute_dense_csa_teacher_lse( + q_attn_bshd, k_attn_bsd, non_compressed_lse, s['softmax_scale'], s['ratio'] ) - lse_indexer_bsqh = lse_indexer.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2) loss_ref = _ref_dense_indexer_loss( q_idx_bshd, @@ -2364,7 +2628,7 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): w_bsh, q_attn_bshd, k_attn_bsd, - lse_indexer_bsqh, + teacher_lse, indexer_softmax_scale=s['indexer_softmax_scale'], attn_softmax_scale=s['softmax_scale'], ratio=s['ratio'], @@ -3255,9 +3519,9 @@ class TestRealKernelDenseIndexerBackward: gradients for ``q_indexer`` / ``k_indexer`` / ``weights``, then compares each against PyTorch autograd through the matching analytical ``_kl_loss_from_dense_scores`` formulation. ``attn_score`` / - ``attn_l1norm`` / ``lse_indexer`` are captured from the kernel and - treated as constants on the reference side (the attention-side - backward is a separate kernel, out of scope here). + ``attn_l1norm`` / full CSA teacher LSE are treated as constants on the + reference side (the attention-side backward is a separate kernel, out + of scope here). Mirrors the gradient parity check done by ``test_dsv4_hybrid_native_parity::test_dsv4_hybrid_attention_matches_native_reference`` @@ -3314,11 +3578,10 @@ def test_real_dense_backward_grad_matches_autograd(self, reset_lazy_kernel_state dk_kernel = k_idx_real.grad.detach().clone() dw_kernel = w_real.grad.detach().clone() - # ---- Reference: capture the kernel's attn-side / lse_indexer (treated - # as constants) and run autograd through the analytical dense KL. + # ---- Reference: recompute the full CSA teacher denominator (treated + # as constant) and run autograd through the analytical dense KL. from megatron.core.transformer.experimental_attention_variant.csa_utils.fused_sparse_attention import ( _compute_dense_attn_score, - _csa_fwd_flash_mla, _indexer_topk_core, _kl_loss_from_dense_scores, ) @@ -3335,36 +3598,25 @@ def _sbhd_to_bshd(q_sbhd, k_sbd, w_sbh, sm_scale): effective_topk = min(s['indexer_topk'], s['n_comp']) with torch.no_grad(): - q_idx_bshd_bf, k_idx_bsd_bf, _, w_bsh_scaled_bf = _sbhd_to_bshd( - q_idx_init, k_idx_init, w_init, s['indexer_softmax_scale'] - ) - topk_indices_cmp, _, _ = _indexer_topk_core( - q_idx_bshd_bf, k_idx_bsd_bf, w_bsh_scaled_bf, effective_topk, s['ratio'] - ) - compress_topk_idxs = torch.where( - topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1 - ) - combined_local = torch.cat([compress_topk_idxs, win_idxs], dim=-1) - global_idxs = local_to_global_flat(combined_local, s['b']) + window_global_idxs = local_to_global_flat(win_idxs, s['b']) q_flat = query.reshape(s['sq'] * s['b'], s['np_'], s['d']) kv_flat = kv_full.reshape(s['skv'] * s['b'], s['d']) - _, _, lse_indexer = _csa_fwd_flash_mla( - q_flat, - kv_flat, - global_idxs, - s['softmax_scale'], - attn_sink=attn_sink, - topk_length=None, - indexer_topk=effective_topk, + non_compressed_lse = _compute_csa_non_compressed_lse( + q_flat, kv_flat, attn_sink, window_global_idxs, s['softmax_scale'] + ) + non_compressed_lse = ( + non_compressed_lse.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2).contiguous() ) - lse_indexer_bsqh = lse_indexer.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2) q_attn_bshd = query.permute(1, 0, 2, 3).contiguous() k_attn_bsd = kv_full[kv_offset:].permute(1, 0, 2).contiguous() + teacher_lse = _compute_dense_csa_teacher_lse( + q_attn_bshd, k_attn_bsd, non_compressed_lse, s['softmax_scale'], s['ratio'] + ) attn_score_const, attn_l1norm_const = _compute_dense_attn_score( q_attn_bshd, k_attn_bsd.unsqueeze(2), - lse_indexer_bsqh, + teacher_lse, qhead_per_kv_head=s['np_'], softmax_scale=s['softmax_scale'], ratio=s['ratio'], diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py index 28445701845..c1c30dffeb2 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -355,80 +355,38 @@ def _native_fused_sparse_indexer_loss( index_scores: torch.Tensor, topk_indices: torch.Tensor, query: torch.Tensor, - compressed_kv: torch.Tensor, + kv_full: torch.Tensor, attn_sink: torch.Tensor, + window_indices: torch.Tensor, + compressed_kv_offset: int, softmax_scale: float, loss_coeff: float, sparse_loss: bool, calculate_per_token_loss: bool, ) -> torch.Tensor: - batch_size, seqlen, _ = topk_indices.size() - num_heads, head_dim = query.size(2), query.size(3) - n_compressed = compressed_kv.size(0) - - sink = attn_sink.detach().view(1, num_heads, 1, 1).float() - q = query.detach().permute(1, 2, 0, 3).float() - compressed_kv_t = compressed_kv.detach().permute(1, 0, 2) - - if sparse_loss: - safe_indices = topk_indices.clamp(min=0).long() - valid = topk_indices >= 0 - row_valid = valid.any(dim=-1, keepdim=True) - - predict_logits = torch.gather(index_scores, dim=-1, index=safe_indices) - predict_logits = predict_logits.masked_fill(~valid, float("-inf")) - predict_logits = predict_logits.masked_fill(~row_valid, 0.0) - predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) - predict = predict * row_valid.float() - - selected_kv = torch.gather( - compressed_kv_t.unsqueeze(1).expand(-1, seqlen, -1, -1), - dim=2, - index=safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), - ) - attn_scores = torch.einsum("bhsd,bskd->bhsk", q, selected_kv.float()) - attn_scores = attn_scores * softmax_scale - attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) - else: - # Dense loss: KL is computed over the FULL compressed-KV axis (not - # just topk). Index-side and attention-side both use the kernel's - # ratio-causal mask, which we derive analytically from the - # compress_ratio (= seqlen / n_compressed): position k of the - # compressed-KV is valid for query row q iff k < (q + 1) // ratio. - compress_ratio = seqlen // n_compressed - k_idx = torch.arange(n_compressed, device=index_scores.device) - valid_per_q = ( - torch.arange(1, seqlen + 1, device=index_scores.device) // compress_ratio - ).clamp(max=n_compressed) - finite_pos = k_idx.view(1, 1, -1) < valid_per_q.view(1, -1, 1) # (1, sq, n_compressed) - finite_pos = finite_pos.expand(batch_size, -1, -1) - row_valid = finite_pos.any(dim=-1, keepdim=True) - - predict_logits = index_scores.masked_fill(~finite_pos, float("-inf")) - predict_logits = predict_logits.masked_fill(~row_valid, 0.0) - predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) - predict = predict * row_valid.float() - - attn_scores = torch.einsum("bhsd,bkd->bhsk", q, compressed_kv_t.float()) - attn_scores = attn_scores * softmax_scale - attn_mask = finite_pos.unsqueeze(1).expand(-1, num_heads, -1, -1) - attn_scores = attn_scores.masked_fill(~attn_mask, float("-inf")) - - score_max = torch.max(attn_scores.max(dim=-1, keepdim=True).values, sink) - exp_scores = torch.exp(attn_scores - score_max) - exp_sink = torch.exp(sink - score_max) - attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - target = attn_probs.sum(dim=1) - target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-10) - target = target * row_valid.float() - - eps = torch.finfo(torch.float32).tiny - target = target.clamp(min=eps) - predict = predict.clamp(min=eps) - kl_per_row = (target * (torch.log(target) - torch.log(predict))).sum(dim=-1) - kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) - loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() - return loss_coeff * loss + seqlen = query.shape[0] + n_compressed = index_scores.shape[-1] + compress_ratio = seqlen // n_compressed + compressed_positions = torch.arange(n_compressed, device=query.device).view(1, -1) + visible_compressed = torch.arange(1, seqlen + 1, device=query.device).view(-1, 1) + causal_mask = torch.where( + compressed_positions < visible_compressed // compress_ratio, 0.0, float("-inf") + ) + causal_mask = causal_mask.unsqueeze(0).expand(query.shape[1], -1, -1) + return _native_unfused_sparse_indexer_loss( + index_scores, + topk_indices, + query, + kv_full, + attn_sink, + window_indices, + compressed_kv_offset, + softmax_scale, + loss_coeff, + sparse_loss, + causal_mask, + calculate_per_token_loss, + ) def _native_unfused_sparse_indexer_loss( @@ -749,8 +707,10 @@ def forward( index_scores, topk_compressed, query, - compressed_kv, + kv_full, self.attn_sink, + window_idxs, + offset, self.softmax_scale, self.indexer_loss_coeff, self.indexer_use_sparse_loss, From e983f7291448147976f12374fd2c6d2fd2c56a40 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Fri, 7 Aug 2026 10:03:43 -0700 Subject: [PATCH 2/5] Fuse CSA teacher LSE computation Signed-off-by: Hongxiao Bai --- .../csa_teacher_lse.py | 493 ++++++++++++++++++ .../csa_utils/fused_sparse_attention.py | 108 +++- .../test_csa_fused_sparse_attention.py | 154 +++++- 3 files changed, 738 insertions(+), 17 deletions(-) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py b/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py new file mode 100644 index 00000000000..5aaa19599e2 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py @@ -0,0 +1,493 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Memory-efficient Triton kernels for the dense CSA teacher denominator.""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import Tensor + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + triton = None + tl = None + _TRITON_AVAILABLE = False + + +if _TRITON_AVAILABLE: + + @triton.jit + def _csa_window_lse_kernel( + query, + full_kv, + window_indices, + attn_sink, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_kv_row: tl.constexpr, + stride_kv_dim: tl.constexpr, + stride_idx_row: tl.constexpr, + stride_idx_col: tl.constexpr, + stride_sink: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + total_kv: tl.constexpr, + window_width: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Compute ``log(exp(sink) + sum_window(exp(q @ k * scale)))``.""" + query_row = tl.program_id(0) + head_block = tl.program_id(1) + + head_offsets = head_block * BLOCK_H + tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, BLOCK_D) + head_mask = head_offsets < num_heads + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_row * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=head_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + attn_sink + head_offsets * stride_sink, mask=head_mask, other=-float("inf") + ).to(tl.float32) + running_sum = tl.where(head_mask & (running_max > -float("inf")), 1.0, 0.0) + + for key_start in range(0, window_width, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + index_mask = key_offsets < window_width + global_indices = tl.load( + window_indices + query_row * stride_idx_row + key_offsets * stride_idx_col, + mask=index_mask, + other=-1, + ) + valid_keys = index_mask & (global_indices >= 0) & (global_indices < total_kv) + safe_indices = tl.where(valid_keys, global_indices, 0) + + k_offsets = dim_offsets[:, None] * stride_kv_dim + safe_indices[None, :] * stride_kv_row + k = tl.load( + full_kv + k_offsets, mask=dim_mask[:, None] & valid_keys[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = head_mask[:, None] & valid_keys[None, :] + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_row * stride_out_row + head_offsets * stride_out_head, + lse, + mask=head_mask, + ) + + @triton.jit + def _csa_compressed_lse_sbhd_kernel( + query, + compressed_kv, + non_compressed_lse, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_batch: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_batch: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + batch_size: tl.constexpr, + seqlen_q: tl.constexpr, + seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for fixed-shape SBHD input.""" + query_blocks = tl.cdiv(seqlen_q, BLOCK_Q) + batch = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + query_offsets = query_block * BLOCK_Q + row_offsets // BLOCK_H + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + flat_query_rows = query_offsets * batch_size + batch + row_mask = (query_offsets < seqlen_q) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + flat_query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + flat_query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = (query_offsets + 1) // ratio + + for key_start in range(0, seqlen_k, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + key_mask = key_offsets < seqlen_k + k_offsets = ( + batch * stride_k_batch + + dim_offsets[:, None] * stride_k_dim + + key_offsets[None, :] * stride_k_row + ) + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (key_offsets[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + output_offsets = ( + batch * stride_out_batch + + query_offsets * stride_out_row + + head_offsets * stride_out_head + ) + tl.store(output + output_offsets, lse, mask=row_mask) + + @triton.jit + def _csa_compressed_lse_thd_kernel( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_sequences: tl.constexpr, + max_seqlen_q: tl.constexpr, + max_seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + HAS_Q_CAUSAL_OFFSETS: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for packed THD input.""" + query_blocks = tl.cdiv(max_seqlen_q, BLOCK_Q) + sequence = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + query_start = tl.load(cu_seqlens_q + sequence) + query_end = tl.load(cu_seqlens_q + sequence + 1) + key_start_offset = tl.load(cu_seqlens_k + sequence) + key_end_offset = tl.load(cu_seqlens_k + sequence + 1) + query_length = query_end - query_start + key_length = key_end_offset - key_start_offset + causal_offset = 0 + if HAS_Q_CAUSAL_OFFSETS: + causal_offset = tl.load(q_causal_offsets + sequence) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + local_queries = query_block * BLOCK_Q + row_offsets // BLOCK_H + query_rows = query_start + local_queries + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + row_mask = (local_queries < query_length) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = tl.minimum((local_queries + causal_offset + 1) // ratio, key_length) + + for key_start in range(0, max_seqlen_k, BLOCK_K): + local_keys = key_start + tl.arange(0, BLOCK_K) + key_mask = local_keys < key_length + k_offsets = (key_start_offset + local_keys[None, :]) * stride_k_row + dim_offsets[ + :, None + ] * stride_k_dim + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (local_keys[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_rows * stride_out_row + head_offsets * stride_out_head, + lse, + mask=row_mask, + ) + + +def can_use_fused_csa_teacher_lse( + query: Tensor, full_kv: Tensor, compressed_kv: Tensor, attn_sink: Tensor, window_indices: Tensor +) -> bool: + """Return whether the Triton teacher-LSE kernels support these tensors.""" + tensors = (query, full_kv, compressed_kv, attn_sink, window_indices) + if not _TRITON_AVAILABLE or not all(tensor.is_cuda for tensor in tensors): + return False + if query.dtype not in (torch.bfloat16, torch.float16): + return False + if full_kv.dtype != query.dtype or compressed_kv.dtype != query.dtype: + return False + if query.ndim != 3 or full_kv.ndim != 2 or compressed_kv.ndim not in (2, 3): + return False + if query.shape[-1] != full_kv.shape[-1] or query.shape[-1] != compressed_kv.shape[-1]: + return False + if query.shape[-1] < 16 or query.shape[-1] > 512: + return False + if query.stride(-1) != 1 or full_kv.stride(-1) != 1 or compressed_kv.stride(-1) != 1: + return False + if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: + return False + if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: + return False + return window_indices.dtype in (torch.int32, torch.int64) + + +@torch.no_grad() +def fused_csa_teacher_lse( + query: Tensor, + full_kv: Tensor, + compressed_kv: Tensor, + attn_sink: Tensor, + window_indices: Tensor, + softmax_scale: float, + ratio: int, + *, + batch_size: Optional[int] = None, + seqlen_q: Optional[int] = None, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_k: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Compute the full dense CSA teacher LSE without a score-matrix temporary. + + ``query`` and ``full_kv`` use FlashMLA's flat-global layout. Sliding-window + indices address ``full_kv`` directly. ``compressed_kv`` is B/K/D for SBHD + or packed K/D for THD. The result is B/S/H for SBHD and T/H for THD. + """ + if ratio <= 0: + raise ValueError(f"ratio must be positive, got {ratio}") + if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, attn_sink, window_indices): + raise ValueError("unsupported tensor layout or dtype for fused CSA teacher LSE") + + total_q, num_heads, head_dim = query.shape + block_d = max(16, triton.next_power_of_2(head_dim)) + window_block_h = min(128, max(16, triton.next_power_of_2(num_heads))) + window_block_k = min(64, max(16, triton.next_power_of_2(max(1, window_indices.shape[1])))) + window_num_stages = 1 if window_block_h == 128 and block_d == 512 else 2 + compressed_block_h = 16 + compressed_block_k = 32 + compressed_block_q = 8 + + non_compressed_lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + window_grid = (total_q, triton.cdiv(num_heads, window_block_h)) + with torch.cuda.device(query.device): + _csa_window_lse_kernel[window_grid]( + query, + full_kv, + window_indices, + attn_sink, + non_compressed_lse, + query.stride(0), + query.stride(1), + query.stride(2), + full_kv.stride(0), + full_kv.stride(1), + window_indices.stride(0), + window_indices.stride(1), + attn_sink.stride(0), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + softmax_scale, + num_heads, + head_dim, + full_kv.shape[0], + window_indices.shape[1], + BLOCK_H=window_block_h, + BLOCK_D=block_d, + BLOCK_K=window_block_k, + num_warps=8, + num_stages=window_num_stages, + ) + + if cu_seqlens_q is None: + if batch_size is None or seqlen_q is None: + raise ValueError("SBHD fused CSA teacher LSE requires batch_size and seqlen_q") + if compressed_kv.ndim != 3 or compressed_kv.shape[0] != batch_size: + raise ValueError("SBHD compressed_kv must have shape [batch, seqlen_k, dim]") + if total_q != batch_size * seqlen_q: + raise ValueError("flat query length must equal batch_size * seqlen_q") + + output = torch.empty( + (batch_size, seqlen_q, num_heads), device=query.device, dtype=torch.float32 + ) + compressed_grid = ( + batch_size * triton.cdiv(seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + _csa_compressed_lse_sbhd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + compressed_kv.stride(2), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + softmax_scale, + batch_size, + seqlen_q, + compressed_kv.shape[1], + num_heads, + head_dim, + ratio, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + + if any(value is None for value in (cu_seqlens_k, max_seqlen_q, max_seqlen_k)): + raise ValueError("THD fused CSA teacher LSE requires packed-sequence metadata") + if compressed_kv.ndim != 2: + raise ValueError("THD compressed_kv must have shape [total_k, dim]") + if cu_seqlens_q.shape != cu_seqlens_k.shape: + raise ValueError("THD query and compressed-KV cumulative lengths must match") + + output = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + num_sequences = cu_seqlens_q.numel() - 1 + compressed_grid = ( + num_sequences * triton.cdiv(max_seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + q_offsets_arg = q_causal_offsets if q_causal_offsets is not None else cu_seqlens_q + _csa_compressed_lse_thd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_offsets_arg, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + softmax_scale, + num_sequences, + max_seqlen_q, + max_seqlen_k, + num_heads, + head_dim, + ratio, + HAS_Q_CAUSAL_OFFSETS=q_causal_offsets is not None, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + + +__all__ = ["can_use_fused_csa_teacher_lse", "fused_csa_teacher_lse"] 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 92572b2c533..1ca4f5cb32e 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 @@ -28,6 +28,8 @@ import torch from torch import Tensor +from .csa_teacher_lse import can_use_fused_csa_teacher_lse, fused_csa_teacher_lse + # --------------------------------------------------------------------------- # Lazy kernel imports # --------------------------------------------------------------------------- @@ -413,6 +415,86 @@ def _compute_dense_csa_teacher_lse( return torch.logaddexp(non_compressed_lse.detach().float(), compressed_lse) +@torch.no_grad() +def _compute_full_csa_teacher_lse( + query: Tensor, + query_flat: Tensor, + full_kv_flat: Tensor, + compressed_kv: Tensor, + attn_sink: Tensor, + window_indices: Tensor, + softmax_scale: float, + ratio: int, + *, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_kv: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Return the full CSA teacher LSE through Triton or the eager reference. + + The Triton path streams both key domains through online reductions and + emits only the per-row/per-head LSE. Unsupported layouts and dtypes retain + the eager implementation as a correctness fallback. + """ + if can_use_fused_csa_teacher_lse( + query_flat, full_kv_flat, compressed_kv, attn_sink, window_indices + ): + if query.ndim == 4: + batch, seqlen_q = query.shape[:2] + return fused_csa_teacher_lse( + query_flat, + full_kv_flat, + compressed_kv, + attn_sink, + window_indices, + softmax_scale, + ratio, + batch_size=batch, + seqlen_q=seqlen_q, + ) + return fused_csa_teacher_lse( + query_flat, + full_kv_flat, + compressed_kv, + attn_sink, + window_indices, + softmax_scale, + ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_kv, + q_causal_offsets=q_causal_offsets, + ) + + non_compressed_lse_flat = _compute_csa_non_compressed_lse( + query_flat, full_kv_flat, attn_sink, window_indices, softmax_scale + ) + if query.ndim == 4: + batch, seqlen_q, num_heads = query.shape[:3] + non_compressed_lse = ( + non_compressed_lse_flat.reshape(seqlen_q, batch, num_heads) + .permute(1, 0, 2) + .contiguous() + ) + else: + non_compressed_lse = non_compressed_lse_flat + return _compute_dense_csa_teacher_lse( + query, + compressed_kv, + non_compressed_lse, + softmax_scale, + ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + q_causal_offsets=q_causal_offsets, + ) + + def local_to_global_flat( local_idxs: Tensor, batch_size: int, @@ -1486,19 +1568,13 @@ def forward( max_seqlen_q=int(max_seqlen_q), max_seqlen_kv=int(max_seqlen_compressed_idx), ) - non_compressed_lse_flat = _compute_csa_non_compressed_lse( - q_flat, kv_flat, attn_sink, global_idxs[:, indexer_topk:], softmax_scale - ) - if is_thd: - non_compressed_lse = non_compressed_lse_flat - else: - non_compressed_lse = ( - non_compressed_lse_flat.reshape(sq, b, np_).permute(1, 0, 2).contiguous() - ) - dense_teacher_lse = _compute_dense_csa_teacher_lse( + dense_teacher_lse = _compute_full_csa_teacher_lse( q_attn_det, + q_flat, + kv_flat, k_attn_compressed_det, - non_compressed_lse, + attn_sink, + global_idxs[:, indexer_topk:], softmax_scale, ratio, **dense_attn_kwargs, @@ -1870,13 +1946,13 @@ def forward( if q_padding_mask is not None: index_score = index_score.masked_fill(q_padding_mask.unsqueeze(-1), float("-inf")) index_lse = index_lse.masked_fill(q_padding_mask, float("-inf")) - non_compressed_lse = _compute_csa_non_compressed_lse( - query, kv_full, attn_sink, topk_idxs[:, indexer_topk:], softmax_scale - ) - dense_teacher_lse = _compute_dense_csa_teacher_lse( + dense_teacher_lse = _compute_full_csa_teacher_lse( query.detach(), + query, + kv_full, compressed_kv.detach(), - non_compressed_lse, + attn_sink, + topk_idxs[:, indexer_topk:], softmax_scale, ratio, cu_seqlens_q=cu_seqlens_q, 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 2921da2c465..a71952ee257 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 @@ -50,6 +50,10 @@ indexer_topk, local_to_global_flat, ) +from megatron.core.transformer.experimental_attention_variant.csa_teacher_lse import ( + can_use_fused_csa_teacher_lse, + fused_csa_teacher_lse, +) # --------------------------------------------------------------------------- # Test fixtures / helpers @@ -321,6 +325,154 @@ def test_non_compressed_lse_includes_window_and_sink(self): ) torch.testing.assert_close(sink_only, sink.view(1, -1).expand_as(sink_only)) + +class TestFusedCsaTeacherLse: + """Triton online-LSE kernels match the score-matrix reference.""" + + @staticmethod + def _require_kernel(query, full_kv, compressed_kv, sink, window_indices): + if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, sink, window_indices): + pytest.skip("Triton CSA teacher-LSE kernel is unavailable") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_sbhd_matches_eager_reference(self): + torch.manual_seed(123) + device = "cuda" + seqlen_q, batch, num_heads, head_dim = 17, 2, 16, 32 + seqlen_full, seqlen_compressed, window_width, ratio = 23, 9, 5, 3 + scale = head_dim**-0.5 + + query = torch.randn( + seqlen_q, batch, num_heads, head_dim, device=device, dtype=torch.bfloat16 + ) + full_kv = torch.randn(seqlen_full, batch, head_dim, device=device, dtype=torch.bfloat16) + compressed_kv = torch.randn( + batch, seqlen_compressed, head_dim, device=device, dtype=torch.bfloat16 + ) + sink = torch.randn(num_heads, device=device, dtype=torch.float32) + window_local = torch.randint( + 0, seqlen_full, (batch, seqlen_q, window_width), device=device, dtype=torch.int32 + ) + window_local[:, ::4, -1] = -1 + window_global = local_to_global_flat(window_local, batch) + query_flat = query.reshape(seqlen_q * batch, num_heads, head_dim) + full_kv_flat = full_kv.reshape(seqlen_full * batch, head_dim) + self._require_kernel(query_flat, full_kv_flat, compressed_kv, sink, window_global) + + actual = fused_csa_teacher_lse( + query_flat, + full_kv_flat, + compressed_kv, + sink, + window_global, + scale, + ratio, + batch_size=batch, + seqlen_q=seqlen_q, + ) + + non_compressed = _compute_csa_non_compressed_lse( + query_flat, full_kv_flat, sink, window_global, scale + ) + non_compressed = ( + non_compressed.reshape(seqlen_q, batch, num_heads).permute(1, 0, 2).contiguous() + ) + expected = _compute_dense_csa_teacher_lse( + query.permute(1, 0, 2, 3).contiguous(), compressed_kv, non_compressed, scale, ratio + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=1e-2) + + empty_window = window_global[:, :0] + sink_lse = sink.view(1, 1, num_heads).expand(batch, seqlen_q, num_heads) + sink_and_compressed_expected = _compute_dense_csa_teacher_lse( + query.permute(1, 0, 2, 3).contiguous(), compressed_kv, sink_lse, scale, ratio + ) + sink_and_compressed_actual = fused_csa_teacher_lse( + query_flat, + full_kv_flat, + compressed_kv, + sink, + empty_window, + scale, + ratio, + batch_size=batch, + seqlen_q=seqlen_q, + ) + torch.testing.assert_close( + sink_and_compressed_actual, sink_and_compressed_expected, atol=2e-2, rtol=1e-2 + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_matches_eager_reference_with_offsets(self): + torch.manual_seed(456) + device = "cuda" + query_lengths = [5, 7] + full_kv_lengths = [8, 10] + compressed_lengths = [3, 4] + num_heads, head_dim, window_width, ratio = 16, 32, 4, 2 + total_q = sum(query_lengths) + scale = head_dim**-0.5 + + query = torch.randn(total_q, num_heads, head_dim, device=device, dtype=torch.bfloat16) + full_kv = torch.randn(sum(full_kv_lengths), head_dim, device=device, dtype=torch.bfloat16) + compressed_kv = torch.randn( + sum(compressed_lengths), head_dim, device=device, dtype=torch.bfloat16 + ) + sink = torch.randn(num_heads, device=device, dtype=torch.float32) + cu_q = _make_cu_seqlens(query_lengths, device=device) + cu_full = _make_cu_seqlens(full_kv_lengths, device=device) + cu_compressed = _make_cu_seqlens(compressed_lengths, device=device) + q_offsets = torch.tensor([1, 0], device=device, dtype=torch.int32) + + window_global = torch.empty(total_q, window_width, device=device, dtype=torch.int32) + query_start = 0 + for sequence, query_length in enumerate(query_lengths): + local = torch.randint( + 0, + full_kv_lengths[sequence], + (query_length, window_width), + device=device, + dtype=torch.int32, + ) + window_global[query_start : query_start + query_length] = local + cu_full[sequence] + query_start += query_length + window_global[::3, -1] = -1 + self._require_kernel(query, full_kv, compressed_kv, sink, window_global) + + actual = fused_csa_teacher_lse( + query, + full_kv, + compressed_kv, + sink, + window_global, + scale, + ratio, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_compressed, + max_seqlen_q=max(query_lengths), + max_seqlen_k=max(compressed_lengths), + q_causal_offsets=q_offsets, + ) + + non_compressed = _compute_csa_non_compressed_lse(query, full_kv, sink, window_global, scale) + expected = _compute_dense_csa_teacher_lse( + query, + compressed_kv, + non_compressed, + scale, + ratio, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_compressed, + max_seqlen_q=max(query_lengths), + max_seqlen_kv=max(compressed_lengths), + q_causal_offsets=q_offsets, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=1e-2) + + +class TestDenseCsaTeacherLseReference: + """Eager compressed-key LSE reference for SBHD and THD.""" + def test_dense_sbhd_lse_adds_every_causal_compressed_key(self): query = torch.tensor( [ @@ -1560,7 +1712,7 @@ def fake_dense_teacher_lse(*args, **kwargs): del args, kwargs return expected_teacher_lse - monkeypatch.setattr(dk, '_compute_dense_csa_teacher_lse', fake_dense_teacher_lse) + monkeypatch.setattr(dk, '_compute_full_csa_teacher_lse', fake_dense_teacher_lse) # ---- (a) forward kernel selection + arg shapes ------------------- inputs_a = self._make_inputs() From 9a616ec4620fea9a702ef85999ead59fc2910b4e Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Sat, 8 Aug 2026 04:21:40 -0700 Subject: [PATCH 3/5] Inline CSA teacher LSE kernels Signed-off-by: Hongxiao Bai --- .../csa_teacher_lse.py | 493 ------------------ .../csa_utils/fused_sparse_attention.py | 482 ++++++++++++++++- .../test_csa_fused_sparse_attention.py | 6 +- 3 files changed, 483 insertions(+), 498 deletions(-) delete mode 100644 megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py b/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py deleted file mode 100644 index 5aaa19599e2..00000000000 --- a/megatron/core/transformer/experimental_attention_variant/csa_teacher_lse.py +++ /dev/null @@ -1,493 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Memory-efficient Triton kernels for the dense CSA teacher denominator.""" - -from __future__ import annotations - -from typing import Optional - -import torch -from torch import Tensor - -try: - import triton - import triton.language as tl - - _TRITON_AVAILABLE = True -except ImportError: - triton = None - tl = None - _TRITON_AVAILABLE = False - - -if _TRITON_AVAILABLE: - - @triton.jit - def _csa_window_lse_kernel( - query, - full_kv, - window_indices, - attn_sink, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_kv_row: tl.constexpr, - stride_kv_dim: tl.constexpr, - stride_idx_row: tl.constexpr, - stride_idx_col: tl.constexpr, - stride_sink: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - total_kv: tl.constexpr, - window_width: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Compute ``log(exp(sink) + sum_window(exp(q @ k * scale)))``.""" - query_row = tl.program_id(0) - head_block = tl.program_id(1) - - head_offsets = head_block * BLOCK_H + tl.arange(0, BLOCK_H) - dim_offsets = tl.arange(0, BLOCK_D) - head_mask = head_offsets < num_heads - dim_mask = dim_offsets < head_dim - - q_offsets = ( - query_row * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=head_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - attn_sink + head_offsets * stride_sink, mask=head_mask, other=-float("inf") - ).to(tl.float32) - running_sum = tl.where(head_mask & (running_max > -float("inf")), 1.0, 0.0) - - for key_start in range(0, window_width, BLOCK_K): - key_offsets = key_start + tl.arange(0, BLOCK_K) - index_mask = key_offsets < window_width - global_indices = tl.load( - window_indices + query_row * stride_idx_row + key_offsets * stride_idx_col, - mask=index_mask, - other=-1, - ) - valid_keys = index_mask & (global_indices >= 0) & (global_indices < total_kv) - safe_indices = tl.where(valid_keys, global_indices, 0) - - k_offsets = dim_offsets[:, None] * stride_kv_dim + safe_indices[None, :] * stride_kv_row - k = tl.load( - full_kv + k_offsets, mask=dim_mask[:, None] & valid_keys[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = head_mask[:, None] & valid_keys[None, :] - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - tl.store( - output + query_row * stride_out_row + head_offsets * stride_out_head, - lse, - mask=head_mask, - ) - - @triton.jit - def _csa_compressed_lse_sbhd_kernel( - query, - compressed_kv, - non_compressed_lse, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_k_batch: tl.constexpr, - stride_k_row: tl.constexpr, - stride_k_dim: tl.constexpr, - stride_noncomp_row: tl.constexpr, - stride_noncomp_head: tl.constexpr, - stride_out_batch: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - batch_size: tl.constexpr, - seqlen_q: tl.constexpr, - seqlen_k: tl.constexpr, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - ratio: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Add the causal compressed-key mass for fixed-shape SBHD input.""" - query_blocks = tl.cdiv(seqlen_q, BLOCK_Q) - batch = tl.program_id(0) // query_blocks - query_block = tl.program_id(0) % query_blocks - head_block = tl.program_id(1) - - row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) - query_offsets = query_block * BLOCK_Q + row_offsets // BLOCK_H - head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H - flat_query_rows = query_offsets * batch_size + batch - row_mask = (query_offsets < seqlen_q) & (head_offsets < num_heads) - dim_offsets = tl.arange(0, BLOCK_D) - dim_mask = dim_offsets < head_dim - - q_offsets = ( - flat_query_rows[:, None] * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - non_compressed_lse - + flat_query_rows * stride_noncomp_row - + head_offsets * stride_noncomp_head, - mask=row_mask, - other=-float("inf"), - ).to(tl.float32) - running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) - visible_keys = (query_offsets + 1) // ratio - - for key_start in range(0, seqlen_k, BLOCK_K): - key_offsets = key_start + tl.arange(0, BLOCK_K) - key_mask = key_offsets < seqlen_k - k_offsets = ( - batch * stride_k_batch - + dim_offsets[:, None] * stride_k_dim - + key_offsets[None, :] * stride_k_row - ) - k = tl.load( - compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = ( - row_mask[:, None] - & key_mask[None, :] - & (key_offsets[None, :] < visible_keys[:, None]) - ) - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - output_offsets = ( - batch * stride_out_batch - + query_offsets * stride_out_row - + head_offsets * stride_out_head - ) - tl.store(output + output_offsets, lse, mask=row_mask) - - @triton.jit - def _csa_compressed_lse_thd_kernel( - query, - compressed_kv, - non_compressed_lse, - cu_seqlens_q, - cu_seqlens_k, - q_causal_offsets, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_k_row: tl.constexpr, - stride_k_dim: tl.constexpr, - stride_noncomp_row: tl.constexpr, - stride_noncomp_head: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - num_sequences: tl.constexpr, - max_seqlen_q: tl.constexpr, - max_seqlen_k: tl.constexpr, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - ratio: tl.constexpr, - HAS_Q_CAUSAL_OFFSETS: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Add the causal compressed-key mass for packed THD input.""" - query_blocks = tl.cdiv(max_seqlen_q, BLOCK_Q) - sequence = tl.program_id(0) // query_blocks - query_block = tl.program_id(0) % query_blocks - head_block = tl.program_id(1) - - query_start = tl.load(cu_seqlens_q + sequence) - query_end = tl.load(cu_seqlens_q + sequence + 1) - key_start_offset = tl.load(cu_seqlens_k + sequence) - key_end_offset = tl.load(cu_seqlens_k + sequence + 1) - query_length = query_end - query_start - key_length = key_end_offset - key_start_offset - causal_offset = 0 - if HAS_Q_CAUSAL_OFFSETS: - causal_offset = tl.load(q_causal_offsets + sequence) - - row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) - local_queries = query_block * BLOCK_Q + row_offsets // BLOCK_H - query_rows = query_start + local_queries - head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H - row_mask = (local_queries < query_length) & (head_offsets < num_heads) - dim_offsets = tl.arange(0, BLOCK_D) - dim_mask = dim_offsets < head_dim - - q_offsets = ( - query_rows[:, None] * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - non_compressed_lse - + query_rows * stride_noncomp_row - + head_offsets * stride_noncomp_head, - mask=row_mask, - other=-float("inf"), - ).to(tl.float32) - running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) - visible_keys = tl.minimum((local_queries + causal_offset + 1) // ratio, key_length) - - for key_start in range(0, max_seqlen_k, BLOCK_K): - local_keys = key_start + tl.arange(0, BLOCK_K) - key_mask = local_keys < key_length - k_offsets = (key_start_offset + local_keys[None, :]) * stride_k_row + dim_offsets[ - :, None - ] * stride_k_dim - k = tl.load( - compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = ( - row_mask[:, None] - & key_mask[None, :] - & (local_keys[None, :] < visible_keys[:, None]) - ) - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - tl.store( - output + query_rows * stride_out_row + head_offsets * stride_out_head, - lse, - mask=row_mask, - ) - - -def can_use_fused_csa_teacher_lse( - query: Tensor, full_kv: Tensor, compressed_kv: Tensor, attn_sink: Tensor, window_indices: Tensor -) -> bool: - """Return whether the Triton teacher-LSE kernels support these tensors.""" - tensors = (query, full_kv, compressed_kv, attn_sink, window_indices) - if not _TRITON_AVAILABLE or not all(tensor.is_cuda for tensor in tensors): - return False - if query.dtype not in (torch.bfloat16, torch.float16): - return False - if full_kv.dtype != query.dtype or compressed_kv.dtype != query.dtype: - return False - if query.ndim != 3 or full_kv.ndim != 2 or compressed_kv.ndim not in (2, 3): - return False - if query.shape[-1] != full_kv.shape[-1] or query.shape[-1] != compressed_kv.shape[-1]: - return False - if query.shape[-1] < 16 or query.shape[-1] > 512: - return False - if query.stride(-1) != 1 or full_kv.stride(-1) != 1 or compressed_kv.stride(-1) != 1: - return False - if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: - return False - if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: - return False - return window_indices.dtype in (torch.int32, torch.int64) - - -@torch.no_grad() -def fused_csa_teacher_lse( - query: Tensor, - full_kv: Tensor, - compressed_kv: Tensor, - attn_sink: Tensor, - window_indices: Tensor, - softmax_scale: float, - ratio: int, - *, - batch_size: Optional[int] = None, - seqlen_q: Optional[int] = None, - cu_seqlens_q: Optional[Tensor] = None, - cu_seqlens_k: Optional[Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_k: Optional[int] = None, - q_causal_offsets: Optional[Tensor] = None, -) -> Tensor: - """Compute the full dense CSA teacher LSE without a score-matrix temporary. - - ``query`` and ``full_kv`` use FlashMLA's flat-global layout. Sliding-window - indices address ``full_kv`` directly. ``compressed_kv`` is B/K/D for SBHD - or packed K/D for THD. The result is B/S/H for SBHD and T/H for THD. - """ - if ratio <= 0: - raise ValueError(f"ratio must be positive, got {ratio}") - if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, attn_sink, window_indices): - raise ValueError("unsupported tensor layout or dtype for fused CSA teacher LSE") - - total_q, num_heads, head_dim = query.shape - block_d = max(16, triton.next_power_of_2(head_dim)) - window_block_h = min(128, max(16, triton.next_power_of_2(num_heads))) - window_block_k = min(64, max(16, triton.next_power_of_2(max(1, window_indices.shape[1])))) - window_num_stages = 1 if window_block_h == 128 and block_d == 512 else 2 - compressed_block_h = 16 - compressed_block_k = 32 - compressed_block_q = 8 - - non_compressed_lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) - window_grid = (total_q, triton.cdiv(num_heads, window_block_h)) - with torch.cuda.device(query.device): - _csa_window_lse_kernel[window_grid]( - query, - full_kv, - window_indices, - attn_sink, - non_compressed_lse, - query.stride(0), - query.stride(1), - query.stride(2), - full_kv.stride(0), - full_kv.stride(1), - window_indices.stride(0), - window_indices.stride(1), - attn_sink.stride(0), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - softmax_scale, - num_heads, - head_dim, - full_kv.shape[0], - window_indices.shape[1], - BLOCK_H=window_block_h, - BLOCK_D=block_d, - BLOCK_K=window_block_k, - num_warps=8, - num_stages=window_num_stages, - ) - - if cu_seqlens_q is None: - if batch_size is None or seqlen_q is None: - raise ValueError("SBHD fused CSA teacher LSE requires batch_size and seqlen_q") - if compressed_kv.ndim != 3 or compressed_kv.shape[0] != batch_size: - raise ValueError("SBHD compressed_kv must have shape [batch, seqlen_k, dim]") - if total_q != batch_size * seqlen_q: - raise ValueError("flat query length must equal batch_size * seqlen_q") - - output = torch.empty( - (batch_size, seqlen_q, num_heads), device=query.device, dtype=torch.float32 - ) - compressed_grid = ( - batch_size * triton.cdiv(seqlen_q, compressed_block_q), - triton.cdiv(num_heads, compressed_block_h), - ) - _csa_compressed_lse_sbhd_kernel[compressed_grid]( - query, - compressed_kv, - non_compressed_lse, - output, - query.stride(0), - query.stride(1), - query.stride(2), - compressed_kv.stride(0), - compressed_kv.stride(1), - compressed_kv.stride(2), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - output.stride(0), - output.stride(1), - output.stride(2), - softmax_scale, - batch_size, - seqlen_q, - compressed_kv.shape[1], - num_heads, - head_dim, - ratio, - BLOCK_Q=compressed_block_q, - BLOCK_H=compressed_block_h, - BLOCK_D=block_d, - BLOCK_K=compressed_block_k, - num_warps=8, - num_stages=2, - ) - return output - - if any(value is None for value in (cu_seqlens_k, max_seqlen_q, max_seqlen_k)): - raise ValueError("THD fused CSA teacher LSE requires packed-sequence metadata") - if compressed_kv.ndim != 2: - raise ValueError("THD compressed_kv must have shape [total_k, dim]") - if cu_seqlens_q.shape != cu_seqlens_k.shape: - raise ValueError("THD query and compressed-KV cumulative lengths must match") - - output = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) - num_sequences = cu_seqlens_q.numel() - 1 - compressed_grid = ( - num_sequences * triton.cdiv(max_seqlen_q, compressed_block_q), - triton.cdiv(num_heads, compressed_block_h), - ) - q_offsets_arg = q_causal_offsets if q_causal_offsets is not None else cu_seqlens_q - _csa_compressed_lse_thd_kernel[compressed_grid]( - query, - compressed_kv, - non_compressed_lse, - cu_seqlens_q, - cu_seqlens_k, - q_offsets_arg, - output, - query.stride(0), - query.stride(1), - query.stride(2), - compressed_kv.stride(0), - compressed_kv.stride(1), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - output.stride(0), - output.stride(1), - softmax_scale, - num_sequences, - max_seqlen_q, - max_seqlen_k, - num_heads, - head_dim, - ratio, - HAS_Q_CAUSAL_OFFSETS=q_causal_offsets is not None, - BLOCK_Q=compressed_block_q, - BLOCK_H=compressed_block_h, - BLOCK_D=block_d, - BLOCK_K=compressed_block_k, - num_warps=8, - num_stages=2, - ) - return output - - -__all__ = ["can_use_fused_csa_teacher_lse", "fused_csa_teacher_lse"] 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 1ca4f5cb32e..fd57c46be3b 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 @@ -28,7 +28,487 @@ import torch from torch import Tensor -from .csa_teacher_lse import can_use_fused_csa_teacher_lse, fused_csa_teacher_lse +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + triton = None + tl = None + _TRITON_AVAILABLE = False + + +if _TRITON_AVAILABLE: + + @triton.jit + def _csa_window_lse_kernel( + query, + full_kv, + window_indices, + attn_sink, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_kv_row: tl.constexpr, + stride_kv_dim: tl.constexpr, + stride_idx_row: tl.constexpr, + stride_idx_col: tl.constexpr, + stride_sink: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + total_kv: tl.constexpr, + window_width: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Compute ``log(exp(sink) + sum_window(exp(q @ k * scale)))``.""" + query_row = tl.program_id(0) + head_block = tl.program_id(1) + + head_offsets = head_block * BLOCK_H + tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, BLOCK_D) + head_mask = head_offsets < num_heads + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_row * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=head_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + attn_sink + head_offsets * stride_sink, mask=head_mask, other=-float("inf") + ).to(tl.float32) + running_sum = tl.where(head_mask & (running_max > -float("inf")), 1.0, 0.0) + + for key_start in range(0, window_width, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + index_mask = key_offsets < window_width + global_indices = tl.load( + window_indices + query_row * stride_idx_row + key_offsets * stride_idx_col, + mask=index_mask, + other=-1, + ) + valid_keys = index_mask & (global_indices >= 0) & (global_indices < total_kv) + safe_indices = tl.where(valid_keys, global_indices, 0) + + k_offsets = dim_offsets[:, None] * stride_kv_dim + safe_indices[None, :] * stride_kv_row + k = tl.load( + full_kv + k_offsets, mask=dim_mask[:, None] & valid_keys[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = head_mask[:, None] & valid_keys[None, :] + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_row * stride_out_row + head_offsets * stride_out_head, + lse, + mask=head_mask, + ) + + @triton.jit + def _csa_compressed_lse_sbhd_kernel( + query, + compressed_kv, + non_compressed_lse, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_batch: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_batch: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + batch_size: tl.constexpr, + seqlen_q: tl.constexpr, + seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for fixed-shape SBHD input.""" + query_blocks = tl.cdiv(seqlen_q, BLOCK_Q) + batch = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + query_offsets = query_block * BLOCK_Q + row_offsets // BLOCK_H + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + flat_query_rows = query_offsets * batch_size + batch + row_mask = (query_offsets < seqlen_q) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + flat_query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + flat_query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = (query_offsets + 1) // ratio + + for key_start in range(0, seqlen_k, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + key_mask = key_offsets < seqlen_k + k_offsets = ( + batch * stride_k_batch + + dim_offsets[:, None] * stride_k_dim + + key_offsets[None, :] * stride_k_row + ) + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (key_offsets[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + output_offsets = ( + batch * stride_out_batch + + query_offsets * stride_out_row + + head_offsets * stride_out_head + ) + tl.store(output + output_offsets, lse, mask=row_mask) + + @triton.jit + def _csa_compressed_lse_thd_kernel( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_sequences: tl.constexpr, + max_seqlen_q: tl.constexpr, + max_seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + HAS_Q_CAUSAL_OFFSETS: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for packed THD input.""" + query_blocks = tl.cdiv(max_seqlen_q, BLOCK_Q) + sequence = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + query_start = tl.load(cu_seqlens_q + sequence) + query_end = tl.load(cu_seqlens_q + sequence + 1) + key_start_offset = tl.load(cu_seqlens_k + sequence) + key_end_offset = tl.load(cu_seqlens_k + sequence + 1) + query_length = query_end - query_start + key_length = key_end_offset - key_start_offset + causal_offset = 0 + if HAS_Q_CAUSAL_OFFSETS: + causal_offset = tl.load(q_causal_offsets + sequence) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + local_queries = query_block * BLOCK_Q + row_offsets // BLOCK_H + query_rows = query_start + local_queries + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + row_mask = (local_queries < query_length) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = tl.minimum((local_queries + causal_offset + 1) // ratio, key_length) + + for key_start in range(0, max_seqlen_k, BLOCK_K): + local_keys = key_start + tl.arange(0, BLOCK_K) + key_mask = local_keys < key_length + k_offsets = (key_start_offset + local_keys[None, :]) * stride_k_row + dim_offsets[ + :, None + ] * stride_k_dim + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (local_keys[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_rows * stride_out_row + head_offsets * stride_out_head, + lse, + mask=row_mask, + ) + + +def can_use_fused_csa_teacher_lse( + query: Tensor, full_kv: Tensor, compressed_kv: Tensor, attn_sink: Tensor, window_indices: Tensor +) -> bool: + """Return whether the Triton teacher-LSE kernels support these tensors.""" + tensors = (query, full_kv, compressed_kv, attn_sink, window_indices) + if not _TRITON_AVAILABLE or not all(tensor.is_cuda for tensor in tensors): + return False + if query.dtype not in (torch.bfloat16, torch.float16): + return False + if full_kv.dtype != query.dtype or compressed_kv.dtype != query.dtype: + return False + if query.ndim != 3 or full_kv.ndim != 2 or compressed_kv.ndim not in (2, 3): + return False + if query.shape[-1] != full_kv.shape[-1] or query.shape[-1] != compressed_kv.shape[-1]: + return False + if query.shape[-1] < 16 or query.shape[-1] > 512: + return False + if query.stride(-1) != 1 or full_kv.stride(-1) != 1 or compressed_kv.stride(-1) != 1: + return False + if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: + return False + if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: + return False + return window_indices.dtype in (torch.int32, torch.int64) + + +@torch.no_grad() +def fused_csa_teacher_lse( + query: Tensor, + full_kv: Tensor, + compressed_kv: Tensor, + attn_sink: Tensor, + window_indices: Tensor, + softmax_scale: float, + ratio: int, + *, + batch_size: Optional[int] = None, + seqlen_q: Optional[int] = None, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_k: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Compute the full dense CSA teacher LSE without a score-matrix temporary. + + ``query`` and ``full_kv`` use FlashMLA's flat-global layout. Sliding-window + indices address ``full_kv`` directly. ``compressed_kv`` is B/K/D for SBHD + or packed K/D for THD. The result is B/S/H for SBHD and T/H for THD. + """ + if ratio <= 0: + raise ValueError(f"ratio must be positive, got {ratio}") + if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, attn_sink, window_indices): + raise ValueError("unsupported tensor layout or dtype for fused CSA teacher LSE") + + total_q, num_heads, head_dim = query.shape + block_d = max(16, triton.next_power_of_2(head_dim)) + window_block_h = min(128, max(16, triton.next_power_of_2(num_heads))) + window_block_k = min(64, max(16, triton.next_power_of_2(max(1, window_indices.shape[1])))) + window_num_stages = 1 if window_block_h == 128 and block_d == 512 else 2 + compressed_block_h = 16 + compressed_block_k = 32 + compressed_block_q = 8 + + non_compressed_lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + window_grid = (total_q, triton.cdiv(num_heads, window_block_h)) + with torch.cuda.device(query.device): + _csa_window_lse_kernel[window_grid]( + query, + full_kv, + window_indices, + attn_sink, + non_compressed_lse, + query.stride(0), + query.stride(1), + query.stride(2), + full_kv.stride(0), + full_kv.stride(1), + window_indices.stride(0), + window_indices.stride(1), + attn_sink.stride(0), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + softmax_scale, + num_heads, + head_dim, + full_kv.shape[0], + window_indices.shape[1], + BLOCK_H=window_block_h, + BLOCK_D=block_d, + BLOCK_K=window_block_k, + num_warps=8, + num_stages=window_num_stages, + ) + + if cu_seqlens_q is None: + if batch_size is None or seqlen_q is None: + raise ValueError("SBHD fused CSA teacher LSE requires batch_size and seqlen_q") + if compressed_kv.ndim != 3 or compressed_kv.shape[0] != batch_size: + raise ValueError("SBHD compressed_kv must have shape [batch, seqlen_k, dim]") + if total_q != batch_size * seqlen_q: + raise ValueError("flat query length must equal batch_size * seqlen_q") + + output = torch.empty( + (batch_size, seqlen_q, num_heads), device=query.device, dtype=torch.float32 + ) + compressed_grid = ( + batch_size * triton.cdiv(seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + _csa_compressed_lse_sbhd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + compressed_kv.stride(2), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + softmax_scale, + batch_size, + seqlen_q, + compressed_kv.shape[1], + num_heads, + head_dim, + ratio, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + + if any(value is None for value in (cu_seqlens_k, max_seqlen_q, max_seqlen_k)): + raise ValueError("THD fused CSA teacher LSE requires packed-sequence metadata") + if compressed_kv.ndim != 2: + raise ValueError("THD compressed_kv must have shape [total_k, dim]") + if cu_seqlens_q.shape != cu_seqlens_k.shape: + raise ValueError("THD query and compressed-KV cumulative lengths must match") + + output = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + num_sequences = cu_seqlens_q.numel() - 1 + compressed_grid = ( + num_sequences * triton.cdiv(max_seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + q_offsets_arg = q_causal_offsets if q_causal_offsets is not None else cu_seqlens_q + _csa_compressed_lse_thd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_offsets_arg, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + softmax_scale, + num_sequences, + max_seqlen_q, + max_seqlen_k, + num_heads, + head_dim, + ratio, + HAS_Q_CAUSAL_OFFSETS=q_causal_offsets is not None, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + # --------------------------------------------------------------------------- # Lazy kernel imports 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 a71952ee257..df718cad99c 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 @@ -45,15 +45,13 @@ _kl_loss_from_target_predict, batch_of_row, build_flat_topk_idxs, + can_use_fused_csa_teacher_lse, csa_sparse_attn, fused_csa_indexer_sparse_attn, + fused_csa_teacher_lse, indexer_topk, local_to_global_flat, ) -from megatron.core.transformer.experimental_attention_variant.csa_teacher_lse import ( - can_use_fused_csa_teacher_lse, - fused_csa_teacher_lse, -) # --------------------------------------------------------------------------- # Test fixtures / helpers From ab1736c293f60365bdc97630e6939bbe02b2a029 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Sat, 8 Aug 2026 06:09:40 -0700 Subject: [PATCH 4/5] Extract CSA teacher LSE kernels Signed-off-by: Hongxiao Bai --- .../csa_utils/csa_teacher_lse.py | 493 ++++++++++++++++++ .../csa_utils/fused_sparse_attention.py | 482 +---------------- .../test_csa_fused_sparse_attention.py | 6 +- 3 files changed, 498 insertions(+), 483 deletions(-) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa_utils/csa_teacher_lse.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa_utils/csa_teacher_lse.py b/megatron/core/transformer/experimental_attention_variant/csa_utils/csa_teacher_lse.py new file mode 100644 index 00000000000..5aaa19599e2 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_utils/csa_teacher_lse.py @@ -0,0 +1,493 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Memory-efficient Triton kernels for the dense CSA teacher denominator.""" + +from __future__ import annotations + +from typing import Optional + +import torch +from torch import Tensor + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + triton = None + tl = None + _TRITON_AVAILABLE = False + + +if _TRITON_AVAILABLE: + + @triton.jit + def _csa_window_lse_kernel( + query, + full_kv, + window_indices, + attn_sink, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_kv_row: tl.constexpr, + stride_kv_dim: tl.constexpr, + stride_idx_row: tl.constexpr, + stride_idx_col: tl.constexpr, + stride_sink: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + total_kv: tl.constexpr, + window_width: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Compute ``log(exp(sink) + sum_window(exp(q @ k * scale)))``.""" + query_row = tl.program_id(0) + head_block = tl.program_id(1) + + head_offsets = head_block * BLOCK_H + tl.arange(0, BLOCK_H) + dim_offsets = tl.arange(0, BLOCK_D) + head_mask = head_offsets < num_heads + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_row * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=head_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + attn_sink + head_offsets * stride_sink, mask=head_mask, other=-float("inf") + ).to(tl.float32) + running_sum = tl.where(head_mask & (running_max > -float("inf")), 1.0, 0.0) + + for key_start in range(0, window_width, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + index_mask = key_offsets < window_width + global_indices = tl.load( + window_indices + query_row * stride_idx_row + key_offsets * stride_idx_col, + mask=index_mask, + other=-1, + ) + valid_keys = index_mask & (global_indices >= 0) & (global_indices < total_kv) + safe_indices = tl.where(valid_keys, global_indices, 0) + + k_offsets = dim_offsets[:, None] * stride_kv_dim + safe_indices[None, :] * stride_kv_row + k = tl.load( + full_kv + k_offsets, mask=dim_mask[:, None] & valid_keys[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = head_mask[:, None] & valid_keys[None, :] + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_row * stride_out_row + head_offsets * stride_out_head, + lse, + mask=head_mask, + ) + + @triton.jit + def _csa_compressed_lse_sbhd_kernel( + query, + compressed_kv, + non_compressed_lse, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_batch: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_batch: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + batch_size: tl.constexpr, + seqlen_q: tl.constexpr, + seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for fixed-shape SBHD input.""" + query_blocks = tl.cdiv(seqlen_q, BLOCK_Q) + batch = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + query_offsets = query_block * BLOCK_Q + row_offsets // BLOCK_H + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + flat_query_rows = query_offsets * batch_size + batch + row_mask = (query_offsets < seqlen_q) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + flat_query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + flat_query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = (query_offsets + 1) // ratio + + for key_start in range(0, seqlen_k, BLOCK_K): + key_offsets = key_start + tl.arange(0, BLOCK_K) + key_mask = key_offsets < seqlen_k + k_offsets = ( + batch * stride_k_batch + + dim_offsets[:, None] * stride_k_dim + + key_offsets[None, :] * stride_k_row + ) + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (key_offsets[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + output_offsets = ( + batch * stride_out_batch + + query_offsets * stride_out_row + + head_offsets * stride_out_head + ) + tl.store(output + output_offsets, lse, mask=row_mask) + + @triton.jit + def _csa_compressed_lse_thd_kernel( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + output, + stride_q_row: tl.constexpr, + stride_q_head: tl.constexpr, + stride_q_dim: tl.constexpr, + stride_k_row: tl.constexpr, + stride_k_dim: tl.constexpr, + stride_noncomp_row: tl.constexpr, + stride_noncomp_head: tl.constexpr, + stride_out_row: tl.constexpr, + stride_out_head: tl.constexpr, + softmax_scale, + num_sequences: tl.constexpr, + max_seqlen_q: tl.constexpr, + max_seqlen_k: tl.constexpr, + num_heads: tl.constexpr, + head_dim: tl.constexpr, + ratio: tl.constexpr, + HAS_Q_CAUSAL_OFFSETS: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Add the causal compressed-key mass for packed THD input.""" + query_blocks = tl.cdiv(max_seqlen_q, BLOCK_Q) + sequence = tl.program_id(0) // query_blocks + query_block = tl.program_id(0) % query_blocks + head_block = tl.program_id(1) + + query_start = tl.load(cu_seqlens_q + sequence) + query_end = tl.load(cu_seqlens_q + sequence + 1) + key_start_offset = tl.load(cu_seqlens_k + sequence) + key_end_offset = tl.load(cu_seqlens_k + sequence + 1) + query_length = query_end - query_start + key_length = key_end_offset - key_start_offset + causal_offset = 0 + if HAS_Q_CAUSAL_OFFSETS: + causal_offset = tl.load(q_causal_offsets + sequence) + + row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) + local_queries = query_block * BLOCK_Q + row_offsets // BLOCK_H + query_rows = query_start + local_queries + head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H + row_mask = (local_queries < query_length) & (head_offsets < num_heads) + dim_offsets = tl.arange(0, BLOCK_D) + dim_mask = dim_offsets < head_dim + + q_offsets = ( + query_rows[:, None] * stride_q_row + + head_offsets[:, None] * stride_q_head + + dim_offsets[None, :] * stride_q_dim + ) + q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) + + running_max = tl.load( + non_compressed_lse + + query_rows * stride_noncomp_row + + head_offsets * stride_noncomp_head, + mask=row_mask, + other=-float("inf"), + ).to(tl.float32) + running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) + visible_keys = tl.minimum((local_queries + causal_offset + 1) // ratio, key_length) + + for key_start in range(0, max_seqlen_k, BLOCK_K): + local_keys = key_start + tl.arange(0, BLOCK_K) + key_mask = local_keys < key_length + k_offsets = (key_start_offset + local_keys[None, :]) * stride_k_row + dim_offsets[ + :, None + ] * stride_k_dim + k = tl.load( + compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 + ) + logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale + score_mask = ( + row_mask[:, None] + & key_mask[None, :] + & (local_keys[None, :] < visible_keys[:, None]) + ) + logits = tl.where(score_mask, logits, -float("inf")) + + tile_max = tl.max(logits, axis=1) + new_max = tl.maximum(running_max, tile_max) + old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) + tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) + running_sum = running_sum * old_scale + tile_sum + running_max = new_max + + lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) + tl.store( + output + query_rows * stride_out_row + head_offsets * stride_out_head, + lse, + mask=row_mask, + ) + + +def can_use_fused_csa_teacher_lse( + query: Tensor, full_kv: Tensor, compressed_kv: Tensor, attn_sink: Tensor, window_indices: Tensor +) -> bool: + """Return whether the Triton teacher-LSE kernels support these tensors.""" + tensors = (query, full_kv, compressed_kv, attn_sink, window_indices) + if not _TRITON_AVAILABLE or not all(tensor.is_cuda for tensor in tensors): + return False + if query.dtype not in (torch.bfloat16, torch.float16): + return False + if full_kv.dtype != query.dtype or compressed_kv.dtype != query.dtype: + return False + if query.ndim != 3 or full_kv.ndim != 2 or compressed_kv.ndim not in (2, 3): + return False + if query.shape[-1] != full_kv.shape[-1] or query.shape[-1] != compressed_kv.shape[-1]: + return False + if query.shape[-1] < 16 or query.shape[-1] > 512: + return False + if query.stride(-1) != 1 or full_kv.stride(-1) != 1 or compressed_kv.stride(-1) != 1: + return False + if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: + return False + if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: + return False + return window_indices.dtype in (torch.int32, torch.int64) + + +@torch.no_grad() +def fused_csa_teacher_lse( + query: Tensor, + full_kv: Tensor, + compressed_kv: Tensor, + attn_sink: Tensor, + window_indices: Tensor, + softmax_scale: float, + ratio: int, + *, + batch_size: Optional[int] = None, + seqlen_q: Optional[int] = None, + cu_seqlens_q: Optional[Tensor] = None, + cu_seqlens_k: Optional[Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + q_causal_offsets: Optional[Tensor] = None, +) -> Tensor: + """Compute the full dense CSA teacher LSE without a score-matrix temporary. + + ``query`` and ``full_kv`` use FlashMLA's flat-global layout. Sliding-window + indices address ``full_kv`` directly. ``compressed_kv`` is B/K/D for SBHD + or packed K/D for THD. The result is B/S/H for SBHD and T/H for THD. + """ + if ratio <= 0: + raise ValueError(f"ratio must be positive, got {ratio}") + if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, attn_sink, window_indices): + raise ValueError("unsupported tensor layout or dtype for fused CSA teacher LSE") + + total_q, num_heads, head_dim = query.shape + block_d = max(16, triton.next_power_of_2(head_dim)) + window_block_h = min(128, max(16, triton.next_power_of_2(num_heads))) + window_block_k = min(64, max(16, triton.next_power_of_2(max(1, window_indices.shape[1])))) + window_num_stages = 1 if window_block_h == 128 and block_d == 512 else 2 + compressed_block_h = 16 + compressed_block_k = 32 + compressed_block_q = 8 + + non_compressed_lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + window_grid = (total_q, triton.cdiv(num_heads, window_block_h)) + with torch.cuda.device(query.device): + _csa_window_lse_kernel[window_grid]( + query, + full_kv, + window_indices, + attn_sink, + non_compressed_lse, + query.stride(0), + query.stride(1), + query.stride(2), + full_kv.stride(0), + full_kv.stride(1), + window_indices.stride(0), + window_indices.stride(1), + attn_sink.stride(0), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + softmax_scale, + num_heads, + head_dim, + full_kv.shape[0], + window_indices.shape[1], + BLOCK_H=window_block_h, + BLOCK_D=block_d, + BLOCK_K=window_block_k, + num_warps=8, + num_stages=window_num_stages, + ) + + if cu_seqlens_q is None: + if batch_size is None or seqlen_q is None: + raise ValueError("SBHD fused CSA teacher LSE requires batch_size and seqlen_q") + if compressed_kv.ndim != 3 or compressed_kv.shape[0] != batch_size: + raise ValueError("SBHD compressed_kv must have shape [batch, seqlen_k, dim]") + if total_q != batch_size * seqlen_q: + raise ValueError("flat query length must equal batch_size * seqlen_q") + + output = torch.empty( + (batch_size, seqlen_q, num_heads), device=query.device, dtype=torch.float32 + ) + compressed_grid = ( + batch_size * triton.cdiv(seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + _csa_compressed_lse_sbhd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + compressed_kv.stride(2), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + softmax_scale, + batch_size, + seqlen_q, + compressed_kv.shape[1], + num_heads, + head_dim, + ratio, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + + if any(value is None for value in (cu_seqlens_k, max_seqlen_q, max_seqlen_k)): + raise ValueError("THD fused CSA teacher LSE requires packed-sequence metadata") + if compressed_kv.ndim != 2: + raise ValueError("THD compressed_kv must have shape [total_k, dim]") + if cu_seqlens_q.shape != cu_seqlens_k.shape: + raise ValueError("THD query and compressed-KV cumulative lengths must match") + + output = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) + num_sequences = cu_seqlens_q.numel() - 1 + compressed_grid = ( + num_sequences * triton.cdiv(max_seqlen_q, compressed_block_q), + triton.cdiv(num_heads, compressed_block_h), + ) + q_offsets_arg = q_causal_offsets if q_causal_offsets is not None else cu_seqlens_q + _csa_compressed_lse_thd_kernel[compressed_grid]( + query, + compressed_kv, + non_compressed_lse, + cu_seqlens_q, + cu_seqlens_k, + q_offsets_arg, + output, + query.stride(0), + query.stride(1), + query.stride(2), + compressed_kv.stride(0), + compressed_kv.stride(1), + non_compressed_lse.stride(0), + non_compressed_lse.stride(1), + output.stride(0), + output.stride(1), + softmax_scale, + num_sequences, + max_seqlen_q, + max_seqlen_k, + num_heads, + head_dim, + ratio, + HAS_Q_CAUSAL_OFFSETS=q_causal_offsets is not None, + BLOCK_Q=compressed_block_q, + BLOCK_H=compressed_block_h, + BLOCK_D=block_d, + BLOCK_K=compressed_block_k, + num_warps=8, + num_stages=2, + ) + return output + + +__all__ = ["can_use_fused_csa_teacher_lse", "fused_csa_teacher_lse"] 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 fd57c46be3b..1ca4f5cb32e 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 @@ -28,487 +28,7 @@ import torch from torch import Tensor -try: - import triton - import triton.language as tl - - _TRITON_AVAILABLE = True -except ImportError: - triton = None - tl = None - _TRITON_AVAILABLE = False - - -if _TRITON_AVAILABLE: - - @triton.jit - def _csa_window_lse_kernel( - query, - full_kv, - window_indices, - attn_sink, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_kv_row: tl.constexpr, - stride_kv_dim: tl.constexpr, - stride_idx_row: tl.constexpr, - stride_idx_col: tl.constexpr, - stride_sink: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - total_kv: tl.constexpr, - window_width: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Compute ``log(exp(sink) + sum_window(exp(q @ k * scale)))``.""" - query_row = tl.program_id(0) - head_block = tl.program_id(1) - - head_offsets = head_block * BLOCK_H + tl.arange(0, BLOCK_H) - dim_offsets = tl.arange(0, BLOCK_D) - head_mask = head_offsets < num_heads - dim_mask = dim_offsets < head_dim - - q_offsets = ( - query_row * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=head_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - attn_sink + head_offsets * stride_sink, mask=head_mask, other=-float("inf") - ).to(tl.float32) - running_sum = tl.where(head_mask & (running_max > -float("inf")), 1.0, 0.0) - - for key_start in range(0, window_width, BLOCK_K): - key_offsets = key_start + tl.arange(0, BLOCK_K) - index_mask = key_offsets < window_width - global_indices = tl.load( - window_indices + query_row * stride_idx_row + key_offsets * stride_idx_col, - mask=index_mask, - other=-1, - ) - valid_keys = index_mask & (global_indices >= 0) & (global_indices < total_kv) - safe_indices = tl.where(valid_keys, global_indices, 0) - - k_offsets = dim_offsets[:, None] * stride_kv_dim + safe_indices[None, :] * stride_kv_row - k = tl.load( - full_kv + k_offsets, mask=dim_mask[:, None] & valid_keys[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = head_mask[:, None] & valid_keys[None, :] - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - tl.store( - output + query_row * stride_out_row + head_offsets * stride_out_head, - lse, - mask=head_mask, - ) - - @triton.jit - def _csa_compressed_lse_sbhd_kernel( - query, - compressed_kv, - non_compressed_lse, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_k_batch: tl.constexpr, - stride_k_row: tl.constexpr, - stride_k_dim: tl.constexpr, - stride_noncomp_row: tl.constexpr, - stride_noncomp_head: tl.constexpr, - stride_out_batch: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - batch_size: tl.constexpr, - seqlen_q: tl.constexpr, - seqlen_k: tl.constexpr, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - ratio: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Add the causal compressed-key mass for fixed-shape SBHD input.""" - query_blocks = tl.cdiv(seqlen_q, BLOCK_Q) - batch = tl.program_id(0) // query_blocks - query_block = tl.program_id(0) % query_blocks - head_block = tl.program_id(1) - - row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) - query_offsets = query_block * BLOCK_Q + row_offsets // BLOCK_H - head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H - flat_query_rows = query_offsets * batch_size + batch - row_mask = (query_offsets < seqlen_q) & (head_offsets < num_heads) - dim_offsets = tl.arange(0, BLOCK_D) - dim_mask = dim_offsets < head_dim - - q_offsets = ( - flat_query_rows[:, None] * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - non_compressed_lse - + flat_query_rows * stride_noncomp_row - + head_offsets * stride_noncomp_head, - mask=row_mask, - other=-float("inf"), - ).to(tl.float32) - running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) - visible_keys = (query_offsets + 1) // ratio - - for key_start in range(0, seqlen_k, BLOCK_K): - key_offsets = key_start + tl.arange(0, BLOCK_K) - key_mask = key_offsets < seqlen_k - k_offsets = ( - batch * stride_k_batch - + dim_offsets[:, None] * stride_k_dim - + key_offsets[None, :] * stride_k_row - ) - k = tl.load( - compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = ( - row_mask[:, None] - & key_mask[None, :] - & (key_offsets[None, :] < visible_keys[:, None]) - ) - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - output_offsets = ( - batch * stride_out_batch - + query_offsets * stride_out_row - + head_offsets * stride_out_head - ) - tl.store(output + output_offsets, lse, mask=row_mask) - - @triton.jit - def _csa_compressed_lse_thd_kernel( - query, - compressed_kv, - non_compressed_lse, - cu_seqlens_q, - cu_seqlens_k, - q_causal_offsets, - output, - stride_q_row: tl.constexpr, - stride_q_head: tl.constexpr, - stride_q_dim: tl.constexpr, - stride_k_row: tl.constexpr, - stride_k_dim: tl.constexpr, - stride_noncomp_row: tl.constexpr, - stride_noncomp_head: tl.constexpr, - stride_out_row: tl.constexpr, - stride_out_head: tl.constexpr, - softmax_scale, - num_sequences: tl.constexpr, - max_seqlen_q: tl.constexpr, - max_seqlen_k: tl.constexpr, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - ratio: tl.constexpr, - HAS_Q_CAUSAL_OFFSETS: tl.constexpr, - BLOCK_Q: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, - BLOCK_K: tl.constexpr, - ): - """Add the causal compressed-key mass for packed THD input.""" - query_blocks = tl.cdiv(max_seqlen_q, BLOCK_Q) - sequence = tl.program_id(0) // query_blocks - query_block = tl.program_id(0) % query_blocks - head_block = tl.program_id(1) - - query_start = tl.load(cu_seqlens_q + sequence) - query_end = tl.load(cu_seqlens_q + sequence + 1) - key_start_offset = tl.load(cu_seqlens_k + sequence) - key_end_offset = tl.load(cu_seqlens_k + sequence + 1) - query_length = query_end - query_start - key_length = key_end_offset - key_start_offset - causal_offset = 0 - if HAS_Q_CAUSAL_OFFSETS: - causal_offset = tl.load(q_causal_offsets + sequence) - - row_offsets = tl.arange(0, BLOCK_Q * BLOCK_H) - local_queries = query_block * BLOCK_Q + row_offsets // BLOCK_H - query_rows = query_start + local_queries - head_offsets = head_block * BLOCK_H + row_offsets % BLOCK_H - row_mask = (local_queries < query_length) & (head_offsets < num_heads) - dim_offsets = tl.arange(0, BLOCK_D) - dim_mask = dim_offsets < head_dim - - q_offsets = ( - query_rows[:, None] * stride_q_row - + head_offsets[:, None] * stride_q_head - + dim_offsets[None, :] * stride_q_dim - ) - q = tl.load(query + q_offsets, mask=row_mask[:, None] & dim_mask[None, :], other=0.0) - - running_max = tl.load( - non_compressed_lse - + query_rows * stride_noncomp_row - + head_offsets * stride_noncomp_head, - mask=row_mask, - other=-float("inf"), - ).to(tl.float32) - running_sum = tl.where(row_mask & (running_max > -float("inf")), 1.0, 0.0) - visible_keys = tl.minimum((local_queries + causal_offset + 1) // ratio, key_length) - - for key_start in range(0, max_seqlen_k, BLOCK_K): - local_keys = key_start + tl.arange(0, BLOCK_K) - key_mask = local_keys < key_length - k_offsets = (key_start_offset + local_keys[None, :]) * stride_k_row + dim_offsets[ - :, None - ] * stride_k_dim - k = tl.load( - compressed_kv + k_offsets, mask=dim_mask[:, None] & key_mask[None, :], other=0.0 - ) - logits = tl.dot(q, k, out_dtype=tl.float32) * softmax_scale - score_mask = ( - row_mask[:, None] - & key_mask[None, :] - & (local_keys[None, :] < visible_keys[:, None]) - ) - logits = tl.where(score_mask, logits, -float("inf")) - - tile_max = tl.max(logits, axis=1) - new_max = tl.maximum(running_max, tile_max) - old_scale = tl.where(running_max > -float("inf"), tl.exp(running_max - new_max), 0.0) - tile_sum = tl.sum(tl.where(score_mask, tl.exp(logits - new_max[:, None]), 0.0), axis=1) - running_sum = running_sum * old_scale + tile_sum - running_max = new_max - - lse = tl.where(running_sum > 0.0, running_max + tl.log(running_sum), -float("inf")) - tl.store( - output + query_rows * stride_out_row + head_offsets * stride_out_head, - lse, - mask=row_mask, - ) - - -def can_use_fused_csa_teacher_lse( - query: Tensor, full_kv: Tensor, compressed_kv: Tensor, attn_sink: Tensor, window_indices: Tensor -) -> bool: - """Return whether the Triton teacher-LSE kernels support these tensors.""" - tensors = (query, full_kv, compressed_kv, attn_sink, window_indices) - if not _TRITON_AVAILABLE or not all(tensor.is_cuda for tensor in tensors): - return False - if query.dtype not in (torch.bfloat16, torch.float16): - return False - if full_kv.dtype != query.dtype or compressed_kv.dtype != query.dtype: - return False - if query.ndim != 3 or full_kv.ndim != 2 or compressed_kv.ndim not in (2, 3): - return False - if query.shape[-1] != full_kv.shape[-1] or query.shape[-1] != compressed_kv.shape[-1]: - return False - if query.shape[-1] < 16 or query.shape[-1] > 512: - return False - if query.stride(-1) != 1 or full_kv.stride(-1) != 1 or compressed_kv.stride(-1) != 1: - return False - if attn_sink.ndim != 1 or attn_sink.numel() != query.shape[1]: - return False - if window_indices.ndim != 2 or window_indices.shape[0] != query.shape[0]: - return False - return window_indices.dtype in (torch.int32, torch.int64) - - -@torch.no_grad() -def fused_csa_teacher_lse( - query: Tensor, - full_kv: Tensor, - compressed_kv: Tensor, - attn_sink: Tensor, - window_indices: Tensor, - softmax_scale: float, - ratio: int, - *, - batch_size: Optional[int] = None, - seqlen_q: Optional[int] = None, - cu_seqlens_q: Optional[Tensor] = None, - cu_seqlens_k: Optional[Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_k: Optional[int] = None, - q_causal_offsets: Optional[Tensor] = None, -) -> Tensor: - """Compute the full dense CSA teacher LSE without a score-matrix temporary. - - ``query`` and ``full_kv`` use FlashMLA's flat-global layout. Sliding-window - indices address ``full_kv`` directly. ``compressed_kv`` is B/K/D for SBHD - or packed K/D for THD. The result is B/S/H for SBHD and T/H for THD. - """ - if ratio <= 0: - raise ValueError(f"ratio must be positive, got {ratio}") - if not can_use_fused_csa_teacher_lse(query, full_kv, compressed_kv, attn_sink, window_indices): - raise ValueError("unsupported tensor layout or dtype for fused CSA teacher LSE") - - total_q, num_heads, head_dim = query.shape - block_d = max(16, triton.next_power_of_2(head_dim)) - window_block_h = min(128, max(16, triton.next_power_of_2(num_heads))) - window_block_k = min(64, max(16, triton.next_power_of_2(max(1, window_indices.shape[1])))) - window_num_stages = 1 if window_block_h == 128 and block_d == 512 else 2 - compressed_block_h = 16 - compressed_block_k = 32 - compressed_block_q = 8 - - non_compressed_lse = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) - window_grid = (total_q, triton.cdiv(num_heads, window_block_h)) - with torch.cuda.device(query.device): - _csa_window_lse_kernel[window_grid]( - query, - full_kv, - window_indices, - attn_sink, - non_compressed_lse, - query.stride(0), - query.stride(1), - query.stride(2), - full_kv.stride(0), - full_kv.stride(1), - window_indices.stride(0), - window_indices.stride(1), - attn_sink.stride(0), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - softmax_scale, - num_heads, - head_dim, - full_kv.shape[0], - window_indices.shape[1], - BLOCK_H=window_block_h, - BLOCK_D=block_d, - BLOCK_K=window_block_k, - num_warps=8, - num_stages=window_num_stages, - ) - - if cu_seqlens_q is None: - if batch_size is None or seqlen_q is None: - raise ValueError("SBHD fused CSA teacher LSE requires batch_size and seqlen_q") - if compressed_kv.ndim != 3 or compressed_kv.shape[0] != batch_size: - raise ValueError("SBHD compressed_kv must have shape [batch, seqlen_k, dim]") - if total_q != batch_size * seqlen_q: - raise ValueError("flat query length must equal batch_size * seqlen_q") - - output = torch.empty( - (batch_size, seqlen_q, num_heads), device=query.device, dtype=torch.float32 - ) - compressed_grid = ( - batch_size * triton.cdiv(seqlen_q, compressed_block_q), - triton.cdiv(num_heads, compressed_block_h), - ) - _csa_compressed_lse_sbhd_kernel[compressed_grid]( - query, - compressed_kv, - non_compressed_lse, - output, - query.stride(0), - query.stride(1), - query.stride(2), - compressed_kv.stride(0), - compressed_kv.stride(1), - compressed_kv.stride(2), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - output.stride(0), - output.stride(1), - output.stride(2), - softmax_scale, - batch_size, - seqlen_q, - compressed_kv.shape[1], - num_heads, - head_dim, - ratio, - BLOCK_Q=compressed_block_q, - BLOCK_H=compressed_block_h, - BLOCK_D=block_d, - BLOCK_K=compressed_block_k, - num_warps=8, - num_stages=2, - ) - return output - - if any(value is None for value in (cu_seqlens_k, max_seqlen_q, max_seqlen_k)): - raise ValueError("THD fused CSA teacher LSE requires packed-sequence metadata") - if compressed_kv.ndim != 2: - raise ValueError("THD compressed_kv must have shape [total_k, dim]") - if cu_seqlens_q.shape != cu_seqlens_k.shape: - raise ValueError("THD query and compressed-KV cumulative lengths must match") - - output = torch.empty((total_q, num_heads), device=query.device, dtype=torch.float32) - num_sequences = cu_seqlens_q.numel() - 1 - compressed_grid = ( - num_sequences * triton.cdiv(max_seqlen_q, compressed_block_q), - triton.cdiv(num_heads, compressed_block_h), - ) - q_offsets_arg = q_causal_offsets if q_causal_offsets is not None else cu_seqlens_q - _csa_compressed_lse_thd_kernel[compressed_grid]( - query, - compressed_kv, - non_compressed_lse, - cu_seqlens_q, - cu_seqlens_k, - q_offsets_arg, - output, - query.stride(0), - query.stride(1), - query.stride(2), - compressed_kv.stride(0), - compressed_kv.stride(1), - non_compressed_lse.stride(0), - non_compressed_lse.stride(1), - output.stride(0), - output.stride(1), - softmax_scale, - num_sequences, - max_seqlen_q, - max_seqlen_k, - num_heads, - head_dim, - ratio, - HAS_Q_CAUSAL_OFFSETS=q_causal_offsets is not None, - BLOCK_Q=compressed_block_q, - BLOCK_H=compressed_block_h, - BLOCK_D=block_d, - BLOCK_K=compressed_block_k, - num_warps=8, - num_stages=2, - ) - return output - +from .csa_teacher_lse import can_use_fused_csa_teacher_lse, fused_csa_teacher_lse # --------------------------------------------------------------------------- # Lazy kernel imports 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 df718cad99c..ed6bbb3bea8 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 @@ -31,6 +31,10 @@ from megatron.core.transformer.experimental_attention_variant.csa_utils import ( fused_sparse_attention as dk, ) +from megatron.core.transformer.experimental_attention_variant.csa_utils.csa_teacher_lse import ( + can_use_fused_csa_teacher_lse, + fused_csa_teacher_lse, +) from megatron.core.transformer.experimental_attention_variant.csa_utils.fused_sparse_attention import ( CSASparseAttnFunc, FusedCSAIndexerSparseAttnFromTopkFunc, @@ -45,10 +49,8 @@ _kl_loss_from_target_predict, batch_of_row, build_flat_topk_idxs, - can_use_fused_csa_teacher_lse, csa_sparse_attn, fused_csa_indexer_sparse_attn, - fused_csa_teacher_lse, indexer_topk, local_to_global_flat, ) From 293308deb9b0d4e2238375fdc3698d035c26cde3 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Mon, 10 Aug 2026 05:53:47 -0700 Subject: [PATCH 5/5] Use compact indices for fused CSA attention Signed-off-by: Hongxiao Bai --- .../csa_utils/fused_sparse_attention.py | 95 +++++++++++---- .../test_csa_fused_sparse_attention.py | 114 +++++++++++++++++- 2 files changed, 186 insertions(+), 23 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 1ca4f5cb32e..761172a3ba3 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 @@ -572,6 +572,30 @@ def local_to_global_flat( return global_idxs.int() +def _compact_flat_topk_idxs(global_idxs: Tensor) -> Tuple[Tensor, Tensor]: + """Pack valid global indices into a per-row prefix. + + The returned ``topk_length`` selects that prefix in FlashMLA forward and + cuDNN DSA backward. Invalid suffix entries remain ``-1`` until forward has + consumed them; callers may then replace the ignored suffix with a safe + non-negative placeholder before backward. + """ + if global_idxs.ndim != 2: + raise ValueError(f"global_idxs must be 2-D (rows, topk), got {tuple(global_idxs.shape)}") + + if global_idxs.is_cuda: + _ensure_dsa_namespace() + res = _DSA.compactify_wrapper(global_idxs) + compact_idxs, topk_length = res["indices"], res["topk_length"] + else: + valid_mask = global_idxs >= 0 + sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) + compact_idxs = global_idxs.gather(-1, sorted_indices) + topk_length = valid_mask.sum(dim=-1).int() + + return compact_idxs.int().contiguous(), topk_length.int().contiguous() + + def build_flat_topk_idxs( *idx_groups: Tensor, batch_size: int, @@ -620,21 +644,9 @@ def build_flat_topk_idxs( topk_length_flat = None if compact: - if global_idxs.is_cuda: - # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA - # namespace. Replaces a stable argsort + gather + sum + permute - # chain with one global-load + global-store per element. - _ensure_dsa_namespace() - res = _DSA.compactify_wrapper(global_idxs) - global_idxs, topk_length_flat = res["indices"], res["topk_length"] - else: - # CPU fallback so the unit tests that exercise this helper without - # CUDA still work. Production callers always go through the CUDA - # path above. - valid_mask = global_idxs >= 0 - sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) - global_idxs = global_idxs.gather(-1, sorted_indices) - topk_length_flat = valid_mask.sum(dim=-1).int() + # The CUDA path is a single warp-per-row CuTe DSL kernel; the helper + # retains a stable PyTorch fallback for CPU-only unit tests. + global_idxs, topk_length_flat = _compact_flat_topk_idxs(global_idxs) return global_idxs, topk_length_flat @@ -1447,6 +1459,13 @@ def forward( combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1) global_idxs = local_to_global_flat(combined_local, b) + # The dense teacher still needs the fixed window suffix. Attention no + # 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. + window_global_idxs = global_idxs[:, indexer_topk:] + global_idxs, topk_length = _compact_flat_topk_idxs(global_idxs) + # ---- 4. FlashMLA forward (flat layout for both SBHD and THD). -------- if is_thd: q_flat = query @@ -1464,9 +1483,14 @@ def forward( global_idxs, softmax_scale, attn_sink=attn_sink, - topk_length=None, + topk_length=topk_length, indexer_topk=0, ) + # cuDNN DSA backward assumes every slot is a non-negative address when + # ``topk_length`` is supplied, including ignored suffix slots. Forward + # has already consumed the ``-1`` sentinels, so sanitize in-place and + # avoid keeping a second TopK-sized tensor for backward. + global_idxs.clamp_min_(0) # ---- 4b. Derive padding-row mask for loss exclusion. ----------------- # When CUDA-graph padding makes cu_seqlens_q cover all total_q rows @@ -1488,6 +1512,12 @@ def forward( real_len_per_row = real_seg_lens[row_batch_ids].to(torch.int32) padding_row_mask = pos_in_seg >= real_len_per_row + # FlashMLA correctly treats a zero-length row as sink-only, but + # cuDNN DSA backward requires at least one tile. The indices were + # sanitized above, so make padding rows consume one harmless + # placeholder during backward; dO and LSE are masked below. + topk_length.masked_fill_(padding_row_mask, 1) + # ---- 5. Derive predict from indexer_scores, compute target. ---------- # Layout-specific attn tensors (detached — loss is not differentiable # through them). @@ -1574,7 +1604,7 @@ def forward( kv_flat, k_attn_compressed_det, attn_sink, - global_idxs[:, indexer_topk:], + window_global_idxs, softmax_scale, ratio, **dense_attn_kwargs, @@ -1721,6 +1751,7 @@ def forward( kv_flat, attn_sink, global_idxs, + topk_length, out_flat, lse, precomputed_grad_q_indexer, @@ -1729,6 +1760,7 @@ def forward( ) ctx.softmax_scale = softmax_scale ctx.is_thd = is_thd + ctx.padding_row_mask = padding_row_mask ctx.np_ = np_ ctx.d = d if is_thd: @@ -1754,6 +1786,7 @@ def backward(ctx, grad_output, grad_loss): kv_flat, attn_sink, global_idxs, + topk_length, out_flat, lse, precomputed_grad_q_indexer, @@ -1772,6 +1805,10 @@ def backward(ctx, grad_output, grad_loss): sq, b, skv = ctx.sq, ctx.b, ctx.skv dO_flat = grad_output.reshape(sq * b, np_, d_v) + if ctx.padding_row_mask is not None: + dO_flat = dO_flat.masked_fill(ctx.padding_row_mask[:, None, None], 0) + lse = lse.masked_fill(ctx.padding_row_mask[:, None], 0) + attn_bwd = _DSA.sparse_attention_backward_wrapper( q_flat, kv_flat, @@ -1781,7 +1818,7 @@ def backward(ctx, grad_output, grad_loss): attn_sink, global_idxs, softmax_scale=ctx.softmax_scale, - topk_length=None, + topk_length=topk_length, ) if is_thd: grad_query = attn_bwd["dq"] @@ -1867,6 +1904,11 @@ def forward( total_comp = k_indexer.shape[0] indexer_topk = indexer_topk_idxs.shape[-1] + # Preserve the fixed window suffix for the dense teacher before + # compacting the complete attention index set. + window_topk_idxs = topk_idxs[:, indexer_topk:] + topk_idxs, topk_length = _compact_flat_topk_idxs(topk_idxs) + # Do not request FlashMLA's partial indexer LSE: it omits both the # window and sink masses required by the CSA teacher. out_flat, lse, _ = _csa_fwd_flash_mla( @@ -1875,9 +1917,14 @@ def forward( topk_idxs, softmax_scale, attn_sink=attn_sink, - topk_length=None, + topk_length=topk_length, indexer_topk=0, ) + topk_idxs.clamp_min_(0) + if q_padding_mask is not None: + # Keep padded sink-only rows out of cuDNN DSA's zero-tile path. + # Backward masks their dO and LSE before using this placeholder. + topk_length.masked_fill_(q_padding_mask, 1) bwd_loss_coeff = loss_coeff * total_q / loss_divisor unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) @@ -1952,7 +1999,7 @@ def forward( kv_full, compressed_kv.detach(), attn_sink, - topk_idxs[:, indexer_topk:], + window_topk_idxs, softmax_scale, ratio, cu_seqlens_q=cu_seqlens_q, @@ -2029,6 +2076,7 @@ def forward( kv_full, attn_sink, topk_idxs, + topk_length, out_flat, lse, saved_grad_q_indexer, @@ -2036,6 +2084,7 @@ def forward( saved_grad_weights, ) ctx.softmax_scale = softmax_scale + ctx.q_padding_mask = q_padding_mask return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss @@ -2048,6 +2097,7 @@ def backward(ctx, grad_output, grad_loss): kv_full, attn_sink, topk_idxs, + topk_length, out_flat, lse, saved_grad_q_indexer, @@ -2056,6 +2106,9 @@ def backward(ctx, grad_output, grad_loss): ) = ctx.saved_tensors dO_flat = grad_output.reshape(query.shape[0], query.shape[1], out_flat.shape[-1]) + if ctx.q_padding_mask is not None: + dO_flat = dO_flat.masked_fill(ctx.q_padding_mask[:, None, None], 0) + lse = lse.masked_fill(ctx.q_padding_mask[:, None], 0) attn_bwd = _DSA.sparse_attention_backward_wrapper( query, kv_full, @@ -2065,7 +2118,7 @@ def backward(ctx, grad_output, grad_loss): attn_sink, topk_idxs, softmax_scale=ctx.softmax_scale, - topk_length=None, + topk_length=topk_length, ) return ( attn_bwd["dq"], 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 ed6bbb3bea8..8892ec06b70 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 @@ -95,6 +95,13 @@ def _make_local_idxs(b: int, sq: int, topk: int, *, with_invalid: bool = False) return base +def _mock_compactify(global_idxs: torch.Tensor): + """Stable PyTorch reference for the mocked cuDNN compactify wrapper.""" + valid = global_idxs >= 0 + order = valid.int().argsort(dim=-1, descending=True, stable=True) + return {'indices': global_idxs.gather(-1, order), 'topk_length': valid.sum(dim=-1).int()} + + def _uniform_dist(B, S, K, dev): """Uniform ``1/K`` distribution of shape ``(B, S, K)``.""" return torch.full((B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev) @@ -1165,6 +1172,7 @@ def _install_full_dsa_mock( target_fn = predict_fn fake_dsa = MagicMock(name='_DSA_full_stub') + fake_dsa.compactify_wrapper.side_effect = _mock_compactify def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} @@ -1282,6 +1290,7 @@ def _install_full_dsa_mock_dense( ) fake_dsa = MagicMock(name='_DSA_full_dense_stub') + fake_dsa.compactify_wrapper.side_effect = _mock_compactify def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} @@ -1469,6 +1478,7 @@ def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel # ---- (a) forward pass-through (no grads needed) ------------------ inputs = self._make_inputs() + inputs['window_idxs'][..., ::2] = -1 fake_dsa_a, flash_stub_a = _install_full_dsa_mock( b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] ) @@ -1498,12 +1508,22 @@ def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel assert ( flash_stub_a.call_args.kwargs['indexer_topk'] == 0 ), "(a) partial indexer LSE should not be requested" + expected_topk_length = torch.full( + (s['sq'] * s['b'],), + 2 + inputs['window_idxs'].shape[-1] // 2, + dtype=torch.int32, + device='cuda', + ) + torch.testing.assert_close( + flash_stub_a.call_args.kwargs['topk_length'], expected_topk_length + ) # ---- (b) backward grad propagation ------------------------------- dk._DSA = None # fresh mocks dk._flash_mla_sparse_fwd = None inputs_b = self._make_inputs(requires_grad=True) - _install_full_dsa_mock( + inputs_b['window_idxs'][..., ::2] = -1 + fake_dsa_b, _ = _install_full_dsa_mock( b=s['b'], sq=s['sq'], np_=s['np_'], @@ -1542,6 +1562,11 @@ def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel f"(b) {name}: grad does not equal full({value}); " f"got first elem = {grad.float().flatten()[0].item()}" ) + attn_bwd_call = fake_dsa_b.sparse_attention_backward_wrapper.call_args + torch.testing.assert_close(attn_bwd_call.kwargs['topk_length'], expected_topk_length) + assert torch.all( + attn_bwd_call.args[6] >= 0 + ), "(b) ignored backward slots were not sanitized" # ---- (c) indexer_topk > n_comp clamp ----------------------------- dk._DSA = None @@ -1707,8 +1732,10 @@ def test_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state expected_teacher_lse = torch.full( (s['b'], s['sq'], s['np_']), 37.0, dtype=torch.float32, device='cuda' ) + teacher_window_indices = [] def fake_dense_teacher_lse(*args, **kwargs): + teacher_window_indices.append(args[5].detach().clone()) del args, kwargs return expected_teacher_lse @@ -1716,6 +1743,7 @@ def fake_dense_teacher_lse(*args, **kwargs): # ---- (a) forward kernel selection + arg shapes ------------------- inputs_a = self._make_inputs() + inputs_a['window_idxs'][..., ::2] = -1 fake_dsa_a, _ = _install_full_dsa_mock_dense( b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] ) @@ -1752,6 +1780,9 @@ def fake_dense_teacher_lse(*args, **kwargs): assert sm_arg == softmax_scale, "(a) dense attn score: positional softmax_scale" assert attn_call.kwargs['qhead_per_kv_head'] == s['np_'] assert attn_call.kwargs['ratio'] == ratio + torch.testing.assert_close( + teacher_window_indices[0], local_to_global_flat(inputs_a['window_idxs'], s['b']) + ) # ---- (b) forward-eager indexer backward + grad propagation -------- dk._DSA = None @@ -1827,6 +1858,7 @@ def _inputs(): def test_sparse_loss_uses_full_flash_lse_plus_sink(self, monkeypatch): inputs = self._inputs() + inputs['topk_idxs'][1, 0] = -1 total_q, num_heads, _ = inputs['query'].shape full_lse = torch.arange(total_q * num_heads, dtype=torch.float32).reshape( total_q, num_heads @@ -1844,8 +1876,10 @@ def fake_flash( topk_length=None, indexer_topk=0, ): - del kv_full, topk_idxs, softmax_scale, d_v, attn_sink, topk_length + del kv_full, softmax_scale, d_v, attn_sink seen['indexer_topk'] = indexer_topk + seen['attention_topk'] = topk_idxs.detach().clone() + seen['topk_length'] = topk_length.detach().clone() return torch.zeros_like(query), full_lse, partial_lse class FakeDSA: @@ -1893,6 +1927,11 @@ def indexer_backward_wrapper(q, w, k, *args, **kwargs): expected = torch.logaddexp(full_lse, inputs['attn_sink'].view(1, num_heads)).unsqueeze(0) torch.testing.assert_close(seen['teacher_lse'], expected) assert seen['indexer_topk'] == 0 + expected_attention_topk = torch.tensor([[4, 0], [1, -1], [5, 2], [5, 3]], dtype=torch.int32) + torch.testing.assert_close(seen['attention_topk'], expected_attention_topk) + torch.testing.assert_close( + seen['topk_length'], torch.tensor([2, 1, 2, 2], dtype=torch.int32) + ) def test_dense_loss_passes_recomputed_full_teacher_lse(self, monkeypatch): inputs = self._inputs() @@ -1969,6 +2008,77 @@ def dense_indexer_backward_wrapper(q, w, k, *args, **kwargs): assert seen['dense_teacher_called'] torch.testing.assert_close(seen['teacher_lse'], sentinel_lse) + 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) + + total_q, num_heads, _ = inputs['query'].shape + q_padding_mask = torch.tensor([False, True, False, False]) + seen = {} + + def fake_flash(query, *args, **kwargs): + del args, kwargs + return torch.zeros_like(query), torch.full((total_q, num_heads), 3.0), None + + class FakeDSA: + @staticmethod + def sparse_indexer_score_recompute_wrapper(q, k, w, topk, **kwargs): + del q, k, w, kwargs + return {'predict': torch.ones_like(topk, dtype=torch.float32)} + + @staticmethod + def sparse_attn_score_recompute_wrapper(q, k, lse, topk, scale, **kwargs): + del q, k, lse, scale, kwargs + return {'target': torch.ones_like(topk, dtype=torch.float32)} + + @staticmethod + def sparse_attention_backward_wrapper( + q, kv, out, dO, lse, sink, topk, *, softmax_scale, topk_length + ): + del out, softmax_scale + seen['topk'] = topk.detach().clone() + seen['topk_length'] = topk_length.detach().clone() + seen['dO'] = dO.detach().clone() + seen['lse'] = lse.detach().clone() + return { + 'dq': torch.zeros_like(q), + 'dkv': torch.zeros_like(kv), + 'd_sink': torch.zeros_like(sink), + } + + monkeypatch.setattr(dk, '_ensure_dsa_namespace', lambda: None) + monkeypatch.setattr(dk, '_csa_fwd_flash_mla', fake_flash) + monkeypatch.setattr(dk, '_DSA', FakeDSA) + + output, loss = FusedCSAIndexerSparseAttnFromTopkFunc.apply( + *inputs.values(), + 1.0, + 1.0, + 0.0, + float(total_q), + True, + 2, + total_q, + ( + torch.tensor([0, total_q], dtype=torch.int32), + torch.tensor([0, inputs['k_indexer'].shape[0]], dtype=torch.int32), + torch.tensor([0], dtype=torch.int32), + ), + q_padding_mask, + ) + (output.sum() + loss).backward() + + torch.testing.assert_close( + seen['topk'], torch.tensor([[4, 0], [0, 0], [5, 2], [5, 3]], dtype=torch.int32) + ) + torch.testing.assert_close( + seen['topk_length'], torch.tensor([2, 1, 2, 2], dtype=torch.int32) + ) + assert torch.count_nonzero(seen['dO'][1]) == 0 + assert torch.count_nonzero(seen['lse'][1]) == 0 + # --------------------------------------------------------------------------- # Real-kernel parity tests (cuDNN + optional FlashMLA)