From 9e8146323381d60e9850e5dee116cbdb0f7e280e Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Wed, 22 Jul 2026 16:17:17 +0800 Subject: [PATCH 1/8] Fix DSv4 unfused CSA indexer teacher Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../experimental_attention_variant/csa.py | 375 ++++++++++++++---- .../test_attention_variant_csa.py | 190 +++++++++ .../test_dsv4_hybrid_native_parity.py | 102 +++-- 3 files changed, 548 insertions(+), 119 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index a4404624678..dc35a7a8e2a 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -20,7 +20,6 @@ from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, - FusedDSAIndexerLoss, fused_qk_topk_naive, fused_qk_topk_naive_thd, rotate_activation, @@ -622,6 +621,226 @@ def unfused_compressed_sparse_attn( return output.reshape(b, sq, np_ * hn).permute(1, 0, 2).contiguous() +def _compute_unfused_csa_indexer_loss( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: 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, + causal_mask: torch.Tensor, + pg_collection: Optional[ProcessGroupCollection], + calculate_per_token_loss: bool = False, +) -> torch.Tensor: + """Compute the unfused CSA indexer loss from the real core-attention teacher. + + The Lightning Indexer predicts only compressed-KV positions, but its teacher + comes from CSA core attention. Therefore each attention head must normalize + its selected compressed logits together with the sliding-window logits and + attention sink before heads are summed and the compressed target is L1 + normalized. Normalizing compressed logits independently would give every + head unit compressed probability mass and discard the head weighting induced + by the window and sink. + + This helper is intentionally specific to the PyTorch CSA fallback. The + generic DSA loss remains unchanged because DSA has neither a sliding-window + branch nor an attention sink. + """ + batch_size, seqlen_q, n_compressed = index_scores.shape + if n_compressed == 0: + return index_scores.float().sum() * 0.0 + + if causal_mask.dim() == 2: + causal_mask = causal_mask.unsqueeze(0) + causal_valid = causal_mask > float("-inf") + if causal_valid.shape[0] == 1 and batch_size > 1: + causal_valid = causal_valid.expand(batch_size, -1, -1) + + if sparse_loss: + selected_compressed = topk_indices + safe_compressed = selected_compressed.clamp(min=0, max=n_compressed - 1).long() + selected_valid = (selected_compressed >= 0) & torch.gather( + causal_valid, dim=-1, index=safe_compressed + ) + predict_logits = torch.gather(index_scores, dim=-1, index=safe_compressed) + else: + selected_compressed = torch.arange( + n_compressed, device=index_scores.device, dtype=topk_indices.dtype + ).view(1, 1, -1) + selected_compressed = selected_compressed.expand(batch_size, seqlen_q, -1) + safe_compressed = selected_compressed.long() + selected_valid = causal_valid + predict_logits = index_scores + + row_valid = selected_valid.any(dim=-1, keepdim=True) + predict_logits = predict_logits.masked_fill(~selected_valid, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + compressed_attention_indices = torch.where( + selected_valid, safe_compressed + compressed_kv_offset, torch.full_like(safe_compressed, -1) + ) + attention_indices = torch.cat( + [window_indices.to(compressed_attention_indices.dtype), compressed_attention_indices], + dim=-1, + ) + safe_attention_indices = attention_indices.clamp(min=0).long() + head_dim = kv_full.shape[-1] + gathered_kv = torch.gather( + kv_full.detach().permute(1, 0, 2).unsqueeze(1).expand(-1, seqlen_q, -1, -1), + dim=2, + index=safe_attention_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), + ) + + query_float = query.detach().permute(1, 2, 0, 3).float() + attention_logits = torch.einsum("bhsd,bskd->bhsk", query_float, gathered_kv.float()) + attention_logits = attention_logits * softmax_scale + attention_logits = attention_logits.masked_fill( + (attention_indices < 0).unsqueeze(1), float("-inf") + ) + + sink = attn_sink.detach().view(1, -1, 1, 1).float() + logits_max = torch.maximum(attention_logits.max(dim=-1, keepdim=True).values, sink) + exp_logits = torch.exp(attention_logits - logits_max) + exp_sink = torch.exp(sink - logits_max) + attention_probs = exp_logits / (exp_logits.sum(dim=-1, keepdim=True) + exp_sink) + + compressed_width = compressed_attention_indices.shape[-1] + target = attention_probs[..., -compressed_width:].sum(dim=1) + if pg_collection is not None and pg_collection.tp.size() > 1: + target = target.contiguous() + torch.distributed.all_reduce(target, group=pg_collection.tp) + eps = torch.finfo(torch.float32).tiny + target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + target = target * row_valid.float() + + target_safe = target.clamp(min=eps) + predict_safe = predict.clamp(min=eps) + kl_per_element = target_safe * (torch.log(target_safe) - torch.log(predict_safe)) + kl_per_row = kl_per_element.sum(dim=-1) + kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) + kl_div = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return kl_div * loss_coeff + + +def _unfused_csa_indexer_topk_and_loss_thd( + q_indexer: torch.Tensor, + k_indexer: torch.Tensor, + weights_indexer: torch.Tensor, + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + window_indices: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_q_unpadded: Optional[torch.Tensor], + cu_seqlens_compressed: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + cu_seqlens_kv_full: torch.Tensor, + topk: int, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + ratio: int, + pg_collection: Optional[ProcessGroupCollection], + calculate_per_token_loss: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run the THD unfused CSA indexer and full-denominator loss per segment.""" + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "_unfused_csa_indexer_topk_and_loss_thd uses GPU-to-CPU synchronizations " + "and cannot run during CUDA graph capture. Use the fused DSA path instead." + ) + + total_q = q_indexer.shape[0] + topk_indices = torch.full((total_q, topk), -1, dtype=torch.int32, device=q_indexer.device) + weighted_losses = [] + + q_offsets = cu_seqlens_q.tolist() + q_unpadded_offsets = ( + cu_seqlens_q_unpadded.tolist() if cu_seqlens_q_unpadded is not None else q_offsets + ) + compressed_offsets = cu_seqlens_compressed.tolist() + kv_offsets = cu_seqlens_kv.tolist() + kv_full_offsets = cu_seqlens_kv_full.tolist() + if not ( + len(q_offsets) + == len(q_unpadded_offsets) + == len(compressed_offsets) + == len(kv_offsets) + == len(kv_full_offsets) + ): + raise ValueError("THD CSA indexer-loss cu_seqlens must have matching segment counts.") + + for batch_idx in range(len(q_offsets) - 1): + q_start, q_padded_end = q_offsets[batch_idx : batch_idx + 2] + q_unpadded_start, q_unpadded_end = q_unpadded_offsets[batch_idx : batch_idx + 2] + compressed_start, compressed_padded_end = compressed_offsets[batch_idx : batch_idx + 2] + kv_start, kv_end = kv_offsets[batch_idx : batch_idx + 2] + kv_full_start, kv_full_end = kv_full_offsets[batch_idx : batch_idx + 2] + seqlen_q = q_unpadded_end - q_unpadded_start + n_compressed = seqlen_q // ratio + if seqlen_q > q_padded_end - q_start: + raise ValueError("THD CSA unpadded query length exceeds its physical segment length.") + if n_compressed > compressed_padded_end - compressed_start: + raise ValueError("THD CSA valid compressed length exceeds its physical segment length.") + if seqlen_q == 0 or n_compressed == 0: + continue + q_end = q_start + seqlen_q + compressed_end = compressed_start + n_compressed + + causal_mask = torch.where( + torch.arange(n_compressed, device=q_indexer.device).unsqueeze(0) + >= torch.arange(1, seqlen_q + 1, device=q_indexer.device).unsqueeze(1) // ratio, + float("-inf"), + 0.0, + ).unsqueeze(0) + + segment_topk = min(topk, n_compressed) + index_scores, segment_topk_indices = fused_qk_topk_naive( + q_indexer[q_start:q_end].unsqueeze(1), + k_indexer[compressed_start:compressed_end].unsqueeze(1), + weights_indexer[q_start:q_end].unsqueeze(1), + segment_topk, + causal_mask, + ) + topk_indices[q_start:q_end, :segment_topk] = segment_topk_indices.squeeze(0).int() + + segment_loss = _compute_unfused_csa_indexer_loss( + index_scores, + segment_topk_indices, + query[q_start:q_end].unsqueeze(1), + kv_full[kv_full_start:kv_full_end].unsqueeze(1), + attn_sink, + window_indices[q_start:q_end].unsqueeze(0), + compressed_kv_offset=kv_end - kv_start, + softmax_scale=softmax_scale, + loss_coeff=loss_coeff, + sparse_loss=sparse_loss, + causal_mask=causal_mask, + pg_collection=pg_collection, + calculate_per_token_loss=calculate_per_token_loss, + ) + weighted_losses.append( + segment_loss if calculate_per_token_loss else segment_loss * seqlen_q + ) + + if not weighted_losses: + zero_loss = ( + q_indexer.float().sum() + k_indexer.float().sum() + weights_indexer.float().sum() + ) * 0.0 + return topk_indices, zero_loss + + indexer_loss = torch.stack(weighted_losses).sum() + if not calculate_per_token_loss: + indexer_loss = indexer_loss / float(max(total_q, 1)) + return topk_indices, indexer_loss + + def _unfused_indexer_sparse_attn_from_topk( query: torch.Tensor, kv_full: torch.Tensor, @@ -654,6 +873,10 @@ def _unfused_indexer_sparse_attn_from_topk( total_q, np_, hn = query.shape cu_seqlens_q, cu_seqlens_k, q_causal_offsets = indexer_layout + indexer_topk = indexer_topk_indices.shape[-1] + attention_indices = topk_indices + if q_padding_mask is not None: + attention_indices = attention_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) if not sparse_loss: q_rows = torch.arange(total_q, dtype=cu_seqlens_q.dtype, device=query.device) @@ -701,15 +924,37 @@ def _unfused_indexer_sparse_attn_from_topk( target_logits = (target_logits * softmax_scale).masked_fill( ~valid.unsqueeze(1), float("-inf") ) - target_logits = torch.where(row_valid.unsqueeze(1), target_logits, 0.0) + + window_indices = attention_indices[start:end, indexer_topk:] + safe_window_indices = window_indices.clamp(min=0).long() + window_kv = kv_full.detach().index_select(0, safe_window_indices.reshape(-1)) + window_kv = window_kv.reshape(end - start, window_indices.shape[-1], hn) + window_logits = torch.einsum( + "rhd,rkd->rhk", query[start:end].detach().float(), window_kv.float() + ) + window_logits = (window_logits * softmax_scale).masked_fill( + ~(window_indices >= 0).unsqueeze(1), float("-inf") + ) + sink = attn_sink.detach().view(1, np_, 1).float() - scores_max = torch.maximum(target_logits.max(dim=-1, keepdim=True).values, sink) + scores_max = target_logits.max(dim=-1, keepdim=True).values + if window_logits.shape[-1] > 0: + scores_max = torch.maximum( + scores_max, window_logits.max(dim=-1, keepdim=True).values + ) + scores_max = torch.maximum(scores_max, sink) exp_scores = torch.exp(target_logits - scores_max) + exp_window = torch.exp(window_logits - scores_max) exp_sink = torch.exp(sink - scores_max) - target = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - target = (target * row_valid.unsqueeze(1).float()).sum(dim=1) + target = exp_scores / ( + exp_scores.sum(dim=-1, keepdim=True) + + exp_window.sum(dim=-1, keepdim=True) + + exp_sink + ) + target = target.sum(dim=1) eps = torch.finfo(torch.float32).tiny target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + target = target * row_valid.float() target = target.clamp(min=eps) predict = predict.clamp(min=eps) @@ -721,8 +966,6 @@ def _unfused_indexer_sparse_attn_from_topk( return output, raw_local_loss * float(loss_coeff) / float(loss_divisor) - indexer_topk = indexer_topk_indices.shape[-1] - if q_padding_mask is not None: indexer_topk_indices = indexer_topk_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) @@ -750,25 +993,26 @@ def _unfused_indexer_sparse_attn_from_topk( predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) predict = predict * row_valid.float() - selected_kv = compressed_kv.detach().index_select(0, safe_indices.reshape(-1)) - selected_kv = selected_kv.reshape(total_q, indexer_topk, hn) + safe_attention_indices = attention_indices.clamp(min=0).long() + selected_kv = kv_full.detach().index_select(0, safe_attention_indices.reshape(-1)) + selected_kv = selected_kv.reshape(total_q, attention_indices.shape[-1], hn) - # This fallback does not run FlashMLA, so it cannot reuse the fused path's - # lse_indexer-based target helper. Recompute the same sink-aware target and - # KL terms explicitly to keep the no-fusion path numerically aligned. + # Recompute the actual CSA denominator over compressed Top-K, sliding + # window, and sink, then retain only the compressed probabilities for the + # Lightning Indexer teacher. attn_scores = torch.einsum("rhd,rkd->rhk", query.detach().float(), selected_kv.float()) attn_scores = attn_scores * softmax_scale - attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + attn_scores = attn_scores.masked_fill(~(attention_indices >= 0).unsqueeze(1), float("-inf")) sink = attn_sink.detach().view(1, np_, 1).float() scores_max = torch.maximum(attn_scores.max(dim=-1, keepdim=True).values, sink) exp_scores = torch.exp(attn_scores - scores_max) exp_sink = torch.exp(sink - scores_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 = attn_probs[..., :indexer_topk].sum(dim=1) + eps = torch.finfo(torch.float32).tiny + target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) 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) @@ -1618,23 +1862,26 @@ def _forward_unfused_csa( x_det, qr_det, packed_seq_params ) indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) - key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) - # ``FusedDSAIndexerLoss`` does not accept a separate - # indexer_softmax_scale; apply it here via the - # weights-scaling trick so the effective weights match - # the pre-scale-split behaviour. weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale - topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + index_scores, topk_indices_compressed = fused_qk_topk_naive( q_indexer, - weights_for_unfused, k_indexer, - query.detach(), - key_for_loss.detach(), - self.softmax_scale, + weights_for_unfused, min(self.indexer.index_topk, n_compressed), - indexer_loss_coeff, causal_mask, + ) + indexer_loss = _compute_unfused_csa_indexer_loss( + index_scores, + topk_indices_compressed, + query, + kv_full, + self.attn_sink, + window_idxs, + offset, + self.softmax_scale, + indexer_loss_coeff, getattr(self.config, "dsa_indexer_use_sparse_loss", True), + causal_mask, self.indexer.pg_collection, self.config.calculate_per_token_loss, ) @@ -1897,16 +2144,16 @@ def _forward_unfused_csa_thd( x: torch.Tensor, qr: torch.Tensor, kv_full_thd: torch.Tensor, - compressed_kv: Optional[torch.Tensor], + _compressed_kv: Optional[torch.Tensor], n_compressed_total: int, - np_: int, + _np: int, total_q: int, cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor, cu_seqlens_kv_full: torch.Tensor, cu_seqlens_compressed: torch.Tensor, window_idxs: torch.Tensor, - max_seqlen_q: int, + _max_seqlen_q: int, max_seqlen_compressed_idx: int, packed_seq_params: PackedSeqParams, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -1937,66 +2184,32 @@ def _forward_unfused_csa_thd( w_thd = weights_indexer.squeeze(1) k_thd = k_indexer.squeeze(1) - key_for_loss_thd = compressed_kv.unsqueeze(1).expand(-1, np_, -1) - weights_for_unfused = w_thd * self.indexer.softmax_scale + weights_for_unfused = w_thd.float() * self.indexer.softmax_scale indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) - # ``_forward_thd`` (caller) absorbs trailing padded - # tokens into the last ``cu_seqlens_q[-1]`` bucket - # so ``batch_of_row`` doesn't OOB; the Compressor - # does the same to ``cu_seqlens_compressed_idx[-1]``. - # Both are correct for the sparse-attention path, - # but the per-segment indexer-loss loop in - # ``fwd/bwd_fused_indexer_loss_naive_thd`` would - # then iterate a "fake" absorbed-padding segment - # with ``seqlen_k_b < topk`` — triggering a write - # shape mismatch and a downstream ``scatter_`` OOB - # on its ``-1`` entries. - # - # Restore the original (pre-absorption) cu_seqlens - # for the loss path so the segment loop's - # ``if seqlen_k_b == 0: continue`` guard skips the - # padding-only iteration. When no padding exists, - # ``packed_seq_params`` already equals the absorbed - # version and this is a no-op. - # Rebuild compressed cu_seqlens from *unpadded* Q lengths. - # This may disagree with the Indexer's cu_seqlens_compressed_idx - # (which uses padded lengths) for the last segment — that's - # intentional: the extra compressed tokens from padding sit at - # the tail of k_thd and are simply never visited by the loss - # loop, which is correct since they don't represent real data. - # Non-last segments are unaffected (padding absorption only - # extends the final segment). - cu_seqlens_q_for_loss = packed_seq_params.cu_seqlens_q - seg_lens_q = cu_seqlens_q_for_loss[1:] - cu_seqlens_q_for_loss[:-1] - cu_seqlens_compressed_idx_for_loss = torch.cat( - [ - torch.zeros( - 1, - dtype=cu_seqlens_q_for_loss.dtype, - device=cu_seqlens_q_for_loss.device, - ), - (seg_lens_q // self.compress_ratio) - .cumsum(0) - .to(cu_seqlens_q_for_loss.dtype), - ] - ) - topk_indices_cmp, indexer_loss = FusedDSAIndexerLoss.apply( + # Keep physical padded offsets for tensor slicing, while using + # the original cu_seqlens to exclude each segment's padding + # from the indexer loss and top-k output. + topk_indices_cmp, indexer_loss = _unfused_csa_indexer_topk_and_loss_thd( q_thd, - weights_for_unfused, k_thd, - query.detach(), - key_for_loss_thd.detach(), - self.softmax_scale, + weights_for_unfused, + query, + kv_full_thd, + self.attn_sink, + window_idxs, + cu_seqlens_q, + packed_seq_params.cu_seqlens_q, + cu_seqlens_compressed_idx, + cu_seqlens_kv, + cu_seqlens_kv_full, min(self.indexer.index_topk, max_seqlen_compressed_idx), + self.softmax_scale, indexer_loss_coeff, - None, getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.compress_ratio, self.indexer.pg_collection, self.config.calculate_per_token_loss, - cu_seqlens_q_for_loss, - cu_seqlens_compressed_idx_for_loss, - self.compress_ratio, ) if indexer_loss_coeff > 0: diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index 29b5a9361c7..e8292b6602b 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -17,6 +17,8 @@ CSAIndexer, CSAIndexerSubmodules, _apply_rope, + _compute_unfused_csa_indexer_loss, + _unfused_csa_indexer_topk_and_loss_thd, build_cu_seqlens_kv_full, cat_per_segment, get_compress_topk_idxs, @@ -66,6 +68,194 @@ def patch_hadamard_if_needed(): # =========================================================================== +@pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) +def test_unfused_csa_indexer_loss_uses_full_attention_denominator(sparse_loss): + """Window and sink logits must affect the compressed-key teacher across heads.""" + query = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) + compressed_kv = torch.tensor([[2.0, 0.0], [0.0, 2.0]])[:, None, :] + index_scores = torch.tensor([[[0.2, -0.1]]], requires_grad=True) + topk_indices = torch.tensor([[[0, 1]]]) + causal_mask = torch.zeros((1, 1, 2)) + window_indices = torch.tensor([[[0]]]) + + def compute_loss(window_kv, attn_sink, scores=index_scores): + kv_full = torch.cat([window_kv.view(1, 1, 2), compressed_kv]) + return _compute_unfused_csa_indexer_loss( + scores, + topk_indices, + query, + kv_full, + attn_sink, + window_indices, + compressed_kv_offset=1, + softmax_scale=1.0, + loss_coeff=1.0, + sparse_loss=sparse_loss, + causal_mask=causal_mask, + pg_collection=None, + ) + + baseline_loss = compute_loss(torch.tensor([0.0, 0.0]), torch.tensor([0.0, 0.0])) + window_perturbed_loss = compute_loss(torch.tensor([6.0, 0.0]), torch.tensor([0.0, 0.0])) + sink_perturbed_loss = compute_loss(torch.tensor([0.0, 0.0]), torch.tensor([6.0, 0.0])) + balanced_extreme_window_loss = compute_loss( + torch.tensor([30.0, 30.0]), torch.tensor([0.0, 0.0]) + ) + predictor_tail_loss = compute_loss( + torch.tensor([30.0, 30.0]), torch.tensor([0.0, 0.0]), torch.tensor([[[0.0, -40.0]]]) + ) + + torch.testing.assert_close(baseline_loss, torch.tensor(0.0112081543)) + torch.testing.assert_close(window_perturbed_loss, torch.tensor(0.4118546844)) + torch.testing.assert_close(sink_perturbed_loss, torch.tensor(0.4118546844)) + torch.testing.assert_close(balanced_extreme_window_loss, baseline_loss) + torch.testing.assert_close(predictor_tail_loss, torch.tensor(19.3068528)) + + +def test_unfused_csa_indexer_loss_ignores_rows_without_compressed_keys(): + """Early rows with no causal compressed key must have finite zero loss and gradient.""" + index_scores = torch.tensor([[[0.2, -0.1]]], requires_grad=True) + loss = _compute_unfused_csa_indexer_loss( + index_scores, + torch.tensor([[[0, 1]]]), + torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]), + torch.tensor([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]])[:, None, :], + torch.tensor([0.0, 0.0]), + torch.tensor([[[0]]]), + compressed_kv_offset=1, + softmax_scale=1.0, + loss_coeff=1.0, + sparse_loss=True, + causal_mask=torch.full((1, 1, 2), float("-inf")), + pg_collection=None, + ) + + loss.backward() + torch.testing.assert_close(loss, torch.tensor(0.0)) + torch.testing.assert_close(index_scores.grad, torch.zeros_like(index_scores)) + + +def test_unfused_csa_indexer_loss_preserves_sparse_and_dense_support(): + """Sparse loss uses selected compressed keys; dense loss uses every causal key.""" + query = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) + index_scores = torch.tensor([[[0.2, -0.1, -0.3]]]) + topk_indices = torch.tensor([[[0, 1]]]) + causal_mask = torch.zeros((1, 1, 3)) + + def compute_loss(third_compressed_kv, sparse_loss): + kv_full = torch.tensor([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0], third_compressed_kv])[ + :, None, : + ] + return _compute_unfused_csa_indexer_loss( + index_scores, + topk_indices, + query, + kv_full, + torch.tensor([0.0, 0.0]), + torch.tensor([[[0]]]), + compressed_kv_offset=1, + softmax_scale=1.0, + loss_coeff=1.0, + sparse_loss=sparse_loss, + causal_mask=causal_mask, + pg_collection=None, + ) + + sparse_baseline = compute_loss([0.0, 0.0], sparse_loss=True) + sparse_perturbed = compute_loss([10.0, 0.0], sparse_loss=True) + dense_baseline = compute_loss([0.0, 0.0], sparse_loss=False) + dense_perturbed = compute_loss([10.0, 0.0], sparse_loss=False) + + torch.testing.assert_close(sparse_perturbed, sparse_baseline) + assert torch.abs(dense_perturbed - dense_baseline) > 1e-3 + + +def test_unfused_csa_indexer_loss_thd_handles_per_segment_padding(): + """The THD wrapper must separate physical offsets from each segment's real length.""" + generator = torch.Generator().manual_seed(7) + q_indexer = torch.randn(24, 2, 2, generator=generator, requires_grad=True) + k_indexer = torch.randn(6, 2, generator=generator, requires_grad=True) + weights = torch.randn(24, 2, generator=generator, requires_grad=True) + query = torch.randn(24, 2, 2, generator=generator) + kv_full = torch.randn(30, 2, generator=generator) + window_indices = torch.stack( + [torch.tensor([max(0, position - 1), position]) for position in list(range(12)) * 2] + ) + + topk_indices, loss = _unfused_csa_indexer_topk_and_loss_thd( + q_indexer, + k_indexer, + weights, + query, + kv_full, + torch.tensor([0.0, 1.0]), + window_indices, + torch.tensor([0, 12, 24], dtype=torch.int32), + torch.tensor([0, 9, 20], dtype=torch.int32), + torch.tensor([0, 3, 6], dtype=torch.int32), + torch.tensor([0, 12, 24], dtype=torch.int32), + torch.tensor([0, 15, 30], dtype=torch.int32), + topk=2, + softmax_scale=1.0, + loss_coeff=0.1, + sparse_loss=True, + ratio=4, + pg_collection=None, + ) + + def segment_oracle(q_start, seqlen_q, compressed_start, kv_full_start): + n_compressed = seqlen_q // 4 + causal_mask = torch.where( + torch.arange(n_compressed).unsqueeze(0) + >= torch.arange(1, seqlen_q + 1).unsqueeze(1) // 4, + float("-inf"), + 0.0, + ).unsqueeze(0) + segment_scores = torch.einsum( + "shd,td->sht", + q_indexer[q_start : q_start + seqlen_q].float(), + k_indexer[compressed_start : compressed_start + n_compressed].float(), + ) + segment_scores = ( + torch.relu(segment_scores) * weights[q_start : q_start + seqlen_q].float().unsqueeze(-1) + ).sum(dim=1) + segment_scores = segment_scores.unsqueeze(0) + causal_mask + segment_topk = segment_scores.topk(2, dim=-1).indices + return _compute_unfused_csa_indexer_loss( + segment_scores, + segment_topk, + query[q_start : q_start + seqlen_q].unsqueeze(1), + kv_full[kv_full_start : kv_full_start + 15].unsqueeze(1), + torch.tensor([0.0, 1.0]), + window_indices[q_start : q_start + seqlen_q].unsqueeze(0), + compressed_kv_offset=12, + softmax_scale=1.0, + loss_coeff=0.1, + sparse_loss=True, + causal_mask=causal_mask, + pg_collection=None, + ) + + with torch.no_grad(): + expected_loss = (segment_oracle(0, 9, 0, 0) * 9 + segment_oracle(12, 11, 3, 15) * 11) / 24 + + assert topk_indices.shape == (24, 2) + torch.testing.assert_close( + topk_indices[[9, 10, 11, 23]], torch.full((4, 2), -1, dtype=torch.int32) + ) + assert torch.isfinite(loss) + assert loss > 0 + torch.testing.assert_close(loss, expected_loss) + loss.backward() + for grad in (q_indexer.grad, k_indexer.grad, weights.grad): + assert grad is not None + assert torch.isfinite(grad).all() + assert torch.count_nonzero(grad) > 0 + torch.testing.assert_close(q_indexer.grad[[9, 10, 11, 23]], torch.zeros(4, 2, 2)) + torch.testing.assert_close(weights.grad[[9, 10, 11, 23]], torch.zeros(4, 2)) + torch.testing.assert_close(k_indexer.grad[[2, 5]], torch.zeros(2, 2)) + + class TestGetWindowTopkIdxs: """Test get_window_topk_idxs helper.""" 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 0831424c88d..c9f3269388f 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 @@ -423,50 +423,74 @@ def _native_unfused_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, causal_mask: torch.Tensor, calculate_per_token_loss: bool, ) -> torch.Tensor: - sq, batch_size, num_heads, _ = query.size() - sk = compressed_kv.size(0) - mask = causal_mask.to(dtype=torch.float32) - + sq, batch_size, _, head_dim = query.size() + n_compressed = index_scores.shape[-1] + causal_valid = causal_mask > float("-inf") + if sparse_loss: + selected_compressed = topk_indices + safe_compressed = selected_compressed.clamp(min=0, max=n_compressed - 1).long() + selected_valid = (selected_compressed >= 0) & torch.gather( + causal_valid, dim=-1, index=safe_compressed + ) + predict_logits = torch.gather(index_scores, dim=-1, index=safe_compressed) + else: + selected_compressed = torch.arange( + n_compressed, device=index_scores.device, dtype=topk_indices.dtype + ).view(1, 1, -1) + selected_compressed = selected_compressed.expand(batch_size, sq, -1) + safe_compressed = selected_compressed.long() + selected_valid = causal_valid + predict_logits = index_scores + + row_valid = selected_valid.any(dim=-1, keepdim=True) + predict_logits = predict_logits.masked_fill(~selected_valid, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) * row_valid.float() + + compressed_attention_indices = torch.where( + selected_valid, safe_compressed + compressed_kv_offset, torch.full_like(safe_compressed, -1) + ) + attention_indices = torch.cat( + [window_indices.to(compressed_attention_indices.dtype), compressed_attention_indices], + dim=-1, + ) + safe_attention_indices = attention_indices.clamp(min=0).long() + gathered_kv = torch.gather( + kv_full.detach().permute(1, 0, 2).unsqueeze(1).expand(-1, sq, -1, -1), + dim=2, + index=safe_attention_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), + ) attention_scores = torch.einsum( - "sbhd,tbd->bhst", query.detach().float(), compressed_kv.detach().float() + "bhsd,bskd->bhsk", query.detach().permute(1, 2, 0, 3).float(), gathered_kv.float() + ) + attention_scores = (attention_scores * softmax_scale).masked_fill( + (attention_indices < 0).unsqueeze(1), float("-inf") ) - attention_scores = attention_scores * softmax_scale - attention_scores = attention_scores + mask.view(batch_size, 1, sq, sk) - index_scores = index_scores + mask + sink = attn_sink.detach().view(1, -1, 1, 1).float() + score_max = torch.maximum(attention_scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(attention_scores - score_max) + exp_sink = torch.exp(sink - score_max) + attention_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - if sparse_loss: - index_mask = torch.full( - (batch_size, sq, sk), float("-inf"), dtype=torch.float32, device=index_scores.device - ).scatter_(-1, topk_indices.clamp(min=0), 0) - attention_scores = attention_scores + index_mask.view(batch_size, 1, sq, sk) - index_scores = index_scores + index_mask - - row_valid = (mask > float("-inf")).any(dim=-1) - attn_row_mask = row_valid.view(batch_size, 1, sq, 1) - idx_row_mask = row_valid.view(batch_size, sq, 1) - - attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) - index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) - - attention_probs = F.softmax(attention_scores, dim=-1, dtype=torch.float32) - predict = F.softmax(index_scores, dim=-1, dtype=torch.float32) - attention_probs = attention_probs * attn_row_mask.float() - predict = predict * idx_row_mask.float() - - target = attention_probs.sum(dim=1) - target = target / target.sum(dim=-1, keepdim=True) + compressed_width = compressed_attention_indices.shape[-1] + target = attention_probs[..., -compressed_width:].sum(dim=1) eps = torch.finfo(torch.float32).tiny + target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + target = target * row_valid.float() 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, kl_per_row, torch.zeros_like(kl_per_row)) + 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 * loss_coeff @@ -700,7 +724,10 @@ def forward( index_scores, topk_compressed, query.detach(), - compressed_kv.detach(), + kv_full.detach(), + self.attn_sink, + window_idxs, + offset, self.softmax_scale, self.indexer_loss_coeff, self.indexer_use_sparse_loss, @@ -1376,13 +1403,12 @@ def test_thd_padded_attention_matches_unpadded( pad_offset += pl hidden_padded = hidden_padded.clone().requires_grad_(True) - # cu_seqlens_q: unpadded real boundaries (cumsum of real lengths - # within the padded physical layout). + # cu_seqlens_q: cumulative unpadded lengths, matching the metadata + # produced by data_schedule. Physical boundaries live separately in + # cu_seqlens_q_padded. cu_seqlens_q_vals = [0] - offset = 0 - for rl, pl in zip(seg_lens, padded_seg_lens): - cu_seqlens_q_vals.append(offset + rl) - offset += pl + for real_length in seg_lens: + cu_seqlens_q_vals.append(cu_seqlens_q_vals[-1] + real_length) cu_seqlens_unpadded = torch.tensor(cu_seqlens_q_vals, dtype=torch.int32, device='cuda') # cu_seqlens_q_padded: padded boundaries (cumsum of padded lengths From f7a277af9af6d9721ee473f5a267d0874c29248b Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 16:09:06 +0800 Subject: [PATCH 2/8] Preserve custom autograd for the CSA indexer loss Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../experimental_attention_variant/csa.py | 404 ++++++------------ .../experimental_attention_variant/dsa.py | 69 ++- .../test_attention_variant_csa.py | 374 ++++++++-------- .../test_dsv4_hybrid_native_parity.py | 19 +- 4 files changed, 393 insertions(+), 473 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index ca168c31474..2ee614ef1bd 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -28,6 +28,7 @@ from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, fused_qk_topk_naive, rotate_activation, ) @@ -684,224 +685,55 @@ def unfused_compressed_sparse_attn( return output.reshape(b, sq, np_ * hn).permute(1, 0, 2).contiguous() -def _compute_unfused_csa_indexer_loss( - index_scores: torch.Tensor, - topk_indices: torch.Tensor, +@torch.no_grad() +def _compute_unfused_csa_non_compressed_lse( query: 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, - causal_mask: torch.Tensor, - pg_collection: Optional[ProcessGroupCollection], - calculate_per_token_loss: bool = False, ) -> torch.Tensor: - """Compute the unfused CSA indexer loss from the real core-attention teacher. - - The Lightning Indexer predicts only compressed-KV positions, but its teacher - comes from CSA core attention. Therefore each attention head must normalize - its selected compressed logits together with the sliding-window logits and - attention sink before heads are summed and the compressed target is L1 - normalized. Normalizing compressed logits independently would give every - head unit compressed probability mass and discard the head weighting induced - by the window and sink. - - This helper is intentionally specific to the PyTorch CSA fallback. The - generic DSA loss remains unchanged because DSA has neither a sliding-window - branch nor an attention sink. + """Return detached sliding-window-plus-sink log mass for the CSA teacher. + + The returned canonical layout is ``[batch, local_heads, seqlen_q]`` for + both SBHD and packed THD inputs. THD callers must pass flat-global window + indices into ``kv_full``. """ - batch_size, seqlen_q, n_compressed = index_scores.shape - if n_compressed == 0: - return index_scores.float().sum() * 0.0 - - if causal_mask.dim() == 2: - causal_mask = causal_mask.unsqueeze(0) - causal_valid = causal_mask > float("-inf") - if causal_valid.shape[0] == 1 and batch_size > 1: - causal_valid = causal_valid.expand(batch_size, -1, -1) - - if sparse_loss: - selected_compressed = topk_indices - safe_compressed = selected_compressed.clamp(min=0, max=n_compressed - 1).long() - selected_valid = (selected_compressed >= 0) & torch.gather( - causal_valid, dim=-1, index=safe_compressed - ) - predict_logits = torch.gather(index_scores, dim=-1, index=safe_compressed) + is_thd = query.ndim == 3 + if is_thd: + q_flat = query + kv_flat = kv_full + global_indices = window_indices + batch_size = 1 + seqlen_q = query.shape[0] else: - selected_compressed = torch.arange( - n_compressed, device=index_scores.device, dtype=topk_indices.dtype - ).view(1, 1, -1) - selected_compressed = selected_compressed.expand(batch_size, seqlen_q, -1) - safe_compressed = selected_compressed.long() - selected_valid = causal_valid - predict_logits = index_scores - - row_valid = selected_valid.any(dim=-1, keepdim=True) - predict_logits = predict_logits.masked_fill(~selected_valid, float("-inf")) - predict_logits = predict_logits.masked_fill(~row_valid, 0.0) - predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) - predict = predict * row_valid.float() + seqlen_q, batch_size, num_heads, head_dim = query.shape + n_kv = kv_full.shape[0] + q_flat = query.permute(1, 0, 2, 3).reshape(batch_size * seqlen_q, num_heads, head_dim) + kv_flat = kv_full.permute(1, 0, 2).reshape(batch_size * n_kv, head_dim) + valid = window_indices >= 0 + batch_offsets = torch.arange(batch_size, device=query.device).view(batch_size, 1, 1) * n_kv + global_indices = torch.where(valid, window_indices + batch_offsets, window_indices).reshape( + batch_size * seqlen_q, -1 + ) - compressed_attention_indices = torch.where( - selected_valid, safe_compressed + compressed_kv_offset, torch.full_like(safe_compressed, -1) - ) - attention_indices = torch.cat( - [window_indices.to(compressed_attention_indices.dtype), compressed_attention_indices], - dim=-1, - ) - safe_attention_indices = attention_indices.clamp(min=0).long() - head_dim = kv_full.shape[-1] + rows, num_heads, head_dim = q_flat.shape + safe_indices = global_indices.clamp(min=0).long() gathered_kv = torch.gather( - kv_full.detach().permute(1, 0, 2).unsqueeze(1).expand(-1, seqlen_q, -1, -1), - dim=2, - index=safe_attention_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), + kv_flat.unsqueeze(0).expand(rows, -1, -1), + dim=1, + index=safe_indices.unsqueeze(-1).expand(-1, -1, head_dim), ) - - query_float = query.detach().permute(1, 2, 0, 3).float() - attention_logits = torch.einsum("bhsd,bskd->bhsk", query_float, gathered_kv.float()) - attention_logits = attention_logits * softmax_scale - attention_logits = attention_logits.masked_fill( - (attention_indices < 0).unsqueeze(1), float("-inf") + window_logits = torch.einsum("rhd,rkd->rhk", q_flat.float(), gathered_kv.float()) + window_logits = (window_logits * softmax_scale).masked_fill( + ~(global_indices >= 0).unsqueeze(1), float("-inf") ) + window_lse = torch.logsumexp(window_logits, dim=-1) + non_compressed_lse = torch.logaddexp(window_lse, attn_sink.detach().view(1, num_heads).float()) - sink = attn_sink.detach().view(1, -1, 1, 1).float() - logits_max = torch.maximum(attention_logits.max(dim=-1, keepdim=True).values, sink) - exp_logits = torch.exp(attention_logits - logits_max) - exp_sink = torch.exp(sink - logits_max) - attention_probs = exp_logits / (exp_logits.sum(dim=-1, keepdim=True) + exp_sink) - - compressed_width = compressed_attention_indices.shape[-1] - target = attention_probs[..., -compressed_width:].sum(dim=1) - if pg_collection is not None and pg_collection.tp.size() > 1: - target = target.contiguous() - torch.distributed.all_reduce(target, group=pg_collection.tp) - eps = torch.finfo(torch.float32).tiny - target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) - target = target * row_valid.float() - - target_safe = target.clamp(min=eps) - predict_safe = predict.clamp(min=eps) - kl_per_element = target_safe * (torch.log(target_safe) - torch.log(predict_safe)) - kl_per_row = kl_per_element.sum(dim=-1) - kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) - kl_div = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() - return kl_div * loss_coeff - - -def _unfused_csa_indexer_topk_and_loss_thd( - q_indexer: torch.Tensor, - k_indexer: torch.Tensor, - weights_indexer: torch.Tensor, - query: torch.Tensor, - kv_full: torch.Tensor, - attn_sink: torch.Tensor, - window_indices: torch.Tensor, - cu_seqlens_q: torch.Tensor, - cu_seqlens_q_unpadded: Optional[torch.Tensor], - cu_seqlens_compressed: torch.Tensor, - cu_seqlens_kv: torch.Tensor, - cu_seqlens_kv_full: torch.Tensor, - topk: int, - softmax_scale: float, - loss_coeff: float, - sparse_loss: bool, - ratio: int, - pg_collection: Optional[ProcessGroupCollection], - calculate_per_token_loss: bool = False, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Run the THD unfused CSA indexer and full-denominator loss per segment.""" - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "_unfused_csa_indexer_topk_and_loss_thd uses GPU-to-CPU synchronizations " - "and cannot run during CUDA graph capture. Use the fused DSA path instead." - ) - - total_q = q_indexer.shape[0] - topk_indices = torch.full((total_q, topk), -1, dtype=torch.int32, device=q_indexer.device) - weighted_losses = [] - - q_offsets = cu_seqlens_q.tolist() - q_unpadded_offsets = ( - cu_seqlens_q_unpadded.tolist() if cu_seqlens_q_unpadded is not None else q_offsets - ) - compressed_offsets = cu_seqlens_compressed.tolist() - kv_offsets = cu_seqlens_kv.tolist() - kv_full_offsets = cu_seqlens_kv_full.tolist() - if not ( - len(q_offsets) - == len(q_unpadded_offsets) - == len(compressed_offsets) - == len(kv_offsets) - == len(kv_full_offsets) - ): - raise ValueError("THD CSA indexer-loss cu_seqlens must have matching segment counts.") - - for batch_idx in range(len(q_offsets) - 1): - q_start, q_padded_end = q_offsets[batch_idx : batch_idx + 2] - q_unpadded_start, q_unpadded_end = q_unpadded_offsets[batch_idx : batch_idx + 2] - compressed_start, compressed_padded_end = compressed_offsets[batch_idx : batch_idx + 2] - kv_start, kv_end = kv_offsets[batch_idx : batch_idx + 2] - kv_full_start, kv_full_end = kv_full_offsets[batch_idx : batch_idx + 2] - seqlen_q = q_unpadded_end - q_unpadded_start - n_compressed = seqlen_q // ratio - if seqlen_q > q_padded_end - q_start: - raise ValueError("THD CSA unpadded query length exceeds its physical segment length.") - if n_compressed > compressed_padded_end - compressed_start: - raise ValueError("THD CSA valid compressed length exceeds its physical segment length.") - if seqlen_q == 0 or n_compressed == 0: - continue - q_end = q_start + seqlen_q - compressed_end = compressed_start + n_compressed - - causal_mask = torch.where( - torch.arange(n_compressed, device=q_indexer.device).unsqueeze(0) - >= torch.arange(1, seqlen_q + 1, device=q_indexer.device).unsqueeze(1) // ratio, - float("-inf"), - 0.0, - ).unsqueeze(0) - - segment_topk = min(topk, n_compressed) - index_scores, segment_topk_indices = fused_qk_topk_naive( - q_indexer[q_start:q_end].unsqueeze(1), - k_indexer[compressed_start:compressed_end].unsqueeze(1), - weights_indexer[q_start:q_end].unsqueeze(1), - segment_topk, - causal_mask, - ) - topk_indices[q_start:q_end, :segment_topk] = segment_topk_indices.squeeze(0).int() - - segment_loss = _compute_unfused_csa_indexer_loss( - index_scores, - segment_topk_indices, - query[q_start:q_end].unsqueeze(1), - kv_full[kv_full_start:kv_full_end].unsqueeze(1), - attn_sink, - window_indices[q_start:q_end].unsqueeze(0), - compressed_kv_offset=kv_end - kv_start, - softmax_scale=softmax_scale, - loss_coeff=loss_coeff, - sparse_loss=sparse_loss, - causal_mask=causal_mask, - pg_collection=pg_collection, - calculate_per_token_loss=calculate_per_token_loss, - ) - weighted_losses.append( - segment_loss if calculate_per_token_loss else segment_loss * seqlen_q - ) - - if not weighted_losses: - zero_loss = ( - q_indexer.float().sum() + k_indexer.float().sum() + weights_indexer.float().sum() - ) * 0.0 - return topk_indices, zero_loss - - indexer_loss = torch.stack(weighted_losses).sum() - if not calculate_per_token_loss: - indexer_loss = indexer_loss / float(max(total_q, 1)) - return topk_indices, indexer_loss + if is_thd: + return non_compressed_lse.transpose(0, 1).unsqueeze(0).contiguous() + return non_compressed_lse.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() def _unfused_indexer_sparse_attn_from_topk( @@ -936,10 +768,6 @@ def _unfused_indexer_sparse_attn_from_topk( total_q, np_, hn = query.shape cu_seqlens_q, cu_seqlens_k, q_causal_offsets = indexer_layout - indexer_topk = indexer_topk_indices.shape[-1] - attention_indices = topk_indices - if q_padding_mask is not None: - attention_indices = attention_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) if not sparse_loss: q_rows = torch.arange(total_q, dtype=cu_seqlens_q.dtype, device=query.device) @@ -987,37 +815,15 @@ def _unfused_indexer_sparse_attn_from_topk( target_logits = (target_logits * softmax_scale).masked_fill( ~valid.unsqueeze(1), float("-inf") ) - - window_indices = attention_indices[start:end, indexer_topk:] - safe_window_indices = window_indices.clamp(min=0).long() - window_kv = kv_full.detach().index_select(0, safe_window_indices.reshape(-1)) - window_kv = window_kv.reshape(end - start, window_indices.shape[-1], hn) - window_logits = torch.einsum( - "rhd,rkd->rhk", query[start:end].detach().float(), window_kv.float() - ) - window_logits = (window_logits * softmax_scale).masked_fill( - ~(window_indices >= 0).unsqueeze(1), float("-inf") - ) - + target_logits = torch.where(row_valid.unsqueeze(1), target_logits, 0.0) sink = attn_sink.detach().view(1, np_, 1).float() - scores_max = target_logits.max(dim=-1, keepdim=True).values - if window_logits.shape[-1] > 0: - scores_max = torch.maximum( - scores_max, window_logits.max(dim=-1, keepdim=True).values - ) - scores_max = torch.maximum(scores_max, sink) + scores_max = torch.maximum(target_logits.max(dim=-1, keepdim=True).values, sink) exp_scores = torch.exp(target_logits - scores_max) - exp_window = torch.exp(window_logits - scores_max) exp_sink = torch.exp(sink - scores_max) - target = exp_scores / ( - exp_scores.sum(dim=-1, keepdim=True) - + exp_window.sum(dim=-1, keepdim=True) - + exp_sink - ) - target = target.sum(dim=1) + target = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + target = (target * row_valid.unsqueeze(1).float()).sum(dim=1) eps = torch.finfo(torch.float32).tiny target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) - target = target * row_valid.float() target = target.clamp(min=eps) predict = predict.clamp(min=eps) @@ -1029,6 +835,8 @@ def _unfused_indexer_sparse_attn_from_topk( return output, raw_local_loss * float(loss_coeff) / float(loss_divisor) + indexer_topk = indexer_topk_indices.shape[-1] + if q_padding_mask is not None: indexer_topk_indices = indexer_topk_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) @@ -1056,26 +864,25 @@ def _unfused_indexer_sparse_attn_from_topk( predict = torch.softmax(predict_logits, dim=-1, dtype=torch.float32) predict = predict * row_valid.float() - safe_attention_indices = attention_indices.clamp(min=0).long() - selected_kv = kv_full.detach().index_select(0, safe_attention_indices.reshape(-1)) - selected_kv = selected_kv.reshape(total_q, attention_indices.shape[-1], hn) + selected_kv = compressed_kv.detach().index_select(0, safe_indices.reshape(-1)) + selected_kv = selected_kv.reshape(total_q, indexer_topk, hn) - # Recompute the actual CSA denominator over compressed Top-K, sliding - # window, and sink, then retain only the compressed probabilities for the - # Lightning Indexer teacher. + # This fallback does not run FlashMLA, so it cannot reuse the fused path's + # lse_indexer-based target helper. Recompute the same sink-aware target and + # KL terms explicitly to keep the no-fusion path numerically aligned. attn_scores = torch.einsum("rhd,rkd->rhk", query.detach().float(), selected_kv.float()) attn_scores = attn_scores * softmax_scale - attn_scores = attn_scores.masked_fill(~(attention_indices >= 0).unsqueeze(1), float("-inf")) + attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) sink = attn_sink.detach().view(1, np_, 1).float() scores_max = torch.maximum(attn_scores.max(dim=-1, keepdim=True).values, sink) exp_scores = torch.exp(attn_scores - scores_max) exp_sink = torch.exp(sink - scores_max) attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - target = attn_probs[..., :indexer_topk].sum(dim=1) - eps = torch.finfo(torch.float32).tiny - target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + 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) @@ -1952,28 +1759,30 @@ def _forward_unfused_csa( x_det, qr_det, packed_seq_params ) indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale - index_scores, topk_indices_compressed = fused_qk_topk_naive( + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv_full, self.attn_sink, window_idxs, self.softmax_scale + ) + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( q_indexer, - k_indexer, weights_for_unfused, - min(self.indexer.index_topk, n_compressed), - causal_mask, - ) - indexer_loss = _compute_unfused_csa_indexer_loss( - index_scores, - topk_indices_compressed, - query, - kv_full, - self.attn_sink, - window_idxs, - offset, + k_indexer, + query.detach(), + key_for_loss.detach(), self.softmax_scale, + min(self.indexer.index_topk, n_compressed), indexer_loss_coeff, - getattr(self.config, "dsa_indexer_use_sparse_loss", True), causal_mask, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), self.indexer.pg_collection, + None, + None, + None, + None, self.config.calculate_per_token_loss, + True, + non_compressed_lse, ) if indexer_loss_coeff > 0: DSAIndexerLossLoggingHelper.save_loss_to_tracker( @@ -2234,16 +2043,16 @@ def _forward_unfused_csa_thd( x: torch.Tensor, qr: torch.Tensor, kv_full_thd: torch.Tensor, - _compressed_kv: Optional[torch.Tensor], + compressed_kv: Optional[torch.Tensor], n_compressed_total: int, - _np: int, + np_: int, total_q: int, cu_seqlens_q: torch.Tensor, cu_seqlens_kv: torch.Tensor, cu_seqlens_kv_full: torch.Tensor, cu_seqlens_compressed: torch.Tensor, window_idxs: torch.Tensor, - _max_seqlen_q: int, + max_seqlen_q: int, max_seqlen_compressed_idx: int, packed_seq_params: PackedSeqParams, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: @@ -2274,32 +2083,67 @@ def _forward_unfused_csa_thd( w_thd = weights_indexer.squeeze(1) k_thd = k_indexer.squeeze(1) - weights_for_unfused = w_thd.float() * self.indexer.softmax_scale + key_for_loss_thd = compressed_kv.unsqueeze(1).expand(-1, np_, -1) + weights_for_unfused = w_thd * self.indexer.softmax_scale indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) - # Keep physical padded offsets for tensor slicing, while using - # the original cu_seqlens to exclude each segment's padding - # from the indexer loss and top-k output. - topk_indices_cmp, indexer_loss = _unfused_csa_indexer_topk_and_loss_thd( - q_thd, - k_thd, - weights_for_unfused, - query, - kv_full_thd, - self.attn_sink, + # Build loss metadata from the original unpadded lengths. + # Indexer K may retain static-capacity tail rows, but DSA's + # row-wise bounds never admit those keys and + # ``query_valid_rows`` zeroes padded query gradients. + cu_seqlens_q_for_loss = packed_seq_params.cu_seqlens_q + seg_lens_q = cu_seqlens_q_for_loss[1:] - cu_seqlens_q_for_loss[:-1] + cu_seqlens_compressed_idx_for_loss = torch.cat( + [ + torch.zeros( + 1, + dtype=cu_seqlens_q_for_loss.dtype, + device=cu_seqlens_q_for_loss.device, + ), + (seg_lens_q // self.compress_ratio) + .cumsum(0) + .to(cu_seqlens_q_for_loss.dtype), + ] + ) + (varlen_starts, varlen_ends, query_valid_rows, compressed_offsets) = ( + _build_compressed_thd_indexer_metadata( + cu_seqlens_q_for_loss, + cu_seqlens_compressed_idx_for_loss, + total_q=q_thd.shape[0], + ratio=self.compress_ratio, + ) + ) + flat_window_idxs, _ = build_flat_topk_idxs( window_idxs, - cu_seqlens_q, - packed_seq_params.cu_seqlens_q, - cu_seqlens_compressed_idx, - cu_seqlens_kv, - cu_seqlens_kv_full, - min(self.indexer.index_topk, max_seqlen_compressed_idx), + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv_full, + ) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv_full_thd, self.attn_sink, flat_window_idxs, self.softmax_scale + ) + topk_indices_global, indexer_loss = FusedDSAIndexerLoss.apply( + q_thd.unsqueeze(1), + weights_for_unfused.unsqueeze(1), + k_thd.unsqueeze(1), + query.detach(), + key_for_loss_thd.detach(), self.softmax_scale, + min(self.indexer.index_topk, max_seqlen_compressed_idx), indexer_loss_coeff, + None, getattr(self.config, "dsa_indexer_use_sparse_loss", True), - self.compress_ratio, self.indexer.pg_collection, + varlen_starts, + varlen_ends, + None, + query_valid_rows, self.config.calculate_per_token_loss, + True, + non_compressed_lse, + ) + topk_indices_cmp = _compressed_thd_topk_to_local( + topk_indices_global, compressed_offsets ) if indexer_loss_coeff > 0: diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 69dea778065..76a7a531c95 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -457,6 +457,7 @@ def compute_dsa_indexer_loss( key_positions: Optional[torch.Tensor] = None, query_valid_rows: Optional[torch.Tensor] = None, calculate_per_token_loss: bool = False, + non_compressed_lse: torch.Tensor | None = None, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -482,6 +483,10 @@ def compute_dsa_indexer_loss( varlen_starts: Optional row-wise key start bounds [sq] for packed THD. varlen_ends: Optional row-wise key end bounds [sq] for packed THD. key_positions: Optional global key positions [sk] for packed THD. + non_compressed_lse: Optional detached FP32 log-sum-exp contribution + [batch, heads, seqlen_q] from teacher keys that are intentionally + omitted from ``key``. When provided, the selected ``key`` logits + are normalized with this external mass before heads are summed. Returns: index_loss: KL divergence loss (scalar). @@ -550,8 +555,8 @@ def compute_dsa_indexer_loss( attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask # [b, np, sq, sk] -> [b, np, sq, sk] - attention_scores = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # [b, sq, sk] -> [b, sq, sk] index_log_scores = dsa_masking.masked_log_softmax( @@ -575,6 +580,43 @@ def compute_dsa_indexer_loss( ) +def _compute_indexer_teacher_probabilities( + attention_scores: torch.Tensor, + attention_valid_mask: torch.Tensor, + non_compressed_lse: torch.Tensor | None = None, +) -> torch.Tensor: + """Normalize selected teacher logits, optionally with omitted attention mass. + + ``non_compressed_lse`` is a sufficient statistic for teacher logits that + must participate in the softmax denominator but must not appear in the + compressed-key target returned by this helper. + """ + b, np, sq, sk = attention_scores.shape + expanded_valid_mask = attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk) + if non_compressed_lse is None: + return dsa_masking.masked_softmax(attention_scores.float(), expanded_valid_mask, dim=-1) + + expected_shape = (b, np, sq) + if tuple(non_compressed_lse.shape) != expected_shape: + raise ValueError( + "non_compressed_lse must have shape [batch, heads, seqlen_q], " + f"got {tuple(non_compressed_lse.shape)}, expected {expected_shape}" + ) + if non_compressed_lse.device != attention_scores.device: + raise ValueError( + "non_compressed_lse and attention_scores must be on the same device, " + f"got {non_compressed_lse.device} and {attention_scores.device}" + ) + if non_compressed_lse.requires_grad: + raise ValueError("non_compressed_lse must be detached") + + masked_scores = attention_scores.float().masked_fill(~expanded_valid_mask, float("-inf")) + compressed_lse = torch.logsumexp(masked_scores, dim=-1) + full_lse = torch.logaddexp(non_compressed_lse.float(), compressed_lse) + probabilities = torch.exp(masked_scores - full_lse.unsqueeze(-1)) + return torch.where(expanded_valid_mask, probabilities, torch.zeros_like(probabilities)) + + def _compute_index_scores( q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True ) -> torch.Tensor: @@ -688,6 +730,7 @@ def fwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of forward pass for indexer loss.""" index_scores, topk_indices = fused_qk_topk_naive( @@ -717,6 +760,7 @@ def fwd_fused_indexer_loss_naive( key_positions=key_positions, query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, + non_compressed_lse=non_compressed_lse, ) return topk_indices, indexer_loss @@ -741,6 +785,7 @@ def bwd_fused_indexer_loss_naive( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """Naive implementation of backward pass for indexer loss.""" query, _ = dsa_layout.ensure_sbhd(query, "query") @@ -813,8 +858,8 @@ def bwd_fused_indexer_loss_naive( else: index_valid_mask = base_valid_mask attention_valid_mask = index_valid_mask if sparse_loss else base_valid_mask - attention_scores_softmax = dsa_masking.masked_softmax( - attention_scores.float(), attention_valid_mask.unsqueeze(1).expand(b, np, sq, sk), dim=-1 + attention_scores_softmax = _compute_indexer_teacher_probabilities( + attention_scores, attention_valid_mask, non_compressed_lse=non_compressed_lse ) # Free attention_scores immediately del attention_scores @@ -951,6 +996,7 @@ def bwd_fused_indexer_loss_naive( "query_valid_rows", "calculate_per_token_loss", "use_relu", + "non_compressed_lse", ) @@ -977,6 +1023,7 @@ def forward( query_valid_rows=None, calculate_per_token_loss: bool = False, use_relu: bool = True, + non_compressed_lse: torch.Tensor | None = None, ): """ Fused forward: index_scores never materialized in full. @@ -999,10 +1046,17 @@ def forward( query_valid_rows=query_valid_rows, calculate_per_token_loss=calculate_per_token_loss, use_relu=use_relu, + non_compressed_lse=non_compressed_lse, ) # Save for backward (recomputation strategy) - ctx.save_for_backward(q, weights, k, query, key, topk_indices) + saved_non_compressed_lse = ( + non_compressed_lse + if non_compressed_lse is not None + else q.new_empty(0, dtype=torch.float32) + ) + ctx.save_for_backward(q, weights, k, query, key, topk_indices, saved_non_compressed_lse) + ctx.has_non_compressed_lse = non_compressed_lse is not None ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss @@ -1022,7 +1076,8 @@ def backward(ctx, grad_topk_indices, grad_loss): """ Backward: Recompute what we need. """ - q, weights, k, query, key, topk_indices = ctx.saved_tensors + q, weights, k, query, key, topk_indices, saved_non_compressed_lse = ctx.saved_tensors + non_compressed_lse = saved_non_compressed_lse if ctx.has_non_compressed_lse else None grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( q, @@ -1043,6 +1098,7 @@ def backward(ctx, grad_topk_indices, grad_loss): query_valid_rows=ctx.query_valid_rows, calculate_per_token_loss=ctx.calculate_per_token_loss, use_relu=ctx.use_relu, + non_compressed_lse=non_compressed_lse, ) grad_by_name = { @@ -1052,6 +1108,7 @@ def backward(ctx, grad_topk_indices, grad_loss): # query and key are detached in forward, so return None for their gradients. "query": None, "key": None, + "non_compressed_lse": None, } return tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index 9fc6969dd75..fbe1f71552a 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -17,8 +17,7 @@ CSAIndexer, CSAIndexerSubmodules, _apply_rope, - _compute_unfused_csa_indexer_loss, - _unfused_csa_indexer_topk_and_loss_thd, + _compute_unfused_csa_non_compressed_lse, build_cu_seqlens_kv_full, cat_per_segment, get_compress_topk_idxs, @@ -27,6 +26,11 @@ get_window_topk_idxs_thd, unfused_compressed_sparse_attn, ) +from megatron.core.transformer.experimental_attention_variant.dsa import ( + FusedDSAIndexerLoss, + compute_dsa_indexer_loss, + fused_qk_topk_naive, +) from megatron.core.transformer.transformer_config import MLATransformerConfig from tests.unit_tests.test_utilities import Utils @@ -68,192 +72,208 @@ def patch_hadamard_if_needed(): # =========================================================================== -@pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) -def test_unfused_csa_indexer_loss_uses_full_attention_denominator(sparse_loss): - """Window and sink logits must affect the compressed-key teacher across heads.""" - query = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) - compressed_kv = torch.tensor([[2.0, 0.0], [0.0, 2.0]])[:, None, :] - index_scores = torch.tensor([[[0.2, -0.1]]], requires_grad=True) - topk_indices = torch.tensor([[[0, 1]]]) - causal_mask = torch.zeros((1, 1, 2)) - window_indices = torch.tensor([[[0]]]) - - def compute_loss(window_kv, attn_sink, scores=index_scores): - kv_full = torch.cat([window_kv.view(1, 1, 2), compressed_kv]) - return _compute_unfused_csa_indexer_loss( - scores, - topk_indices, - query, - kv_full, - attn_sink, - window_indices, - compressed_kv_offset=1, - softmax_scale=1.0, - loss_coeff=1.0, - sparse_loss=sparse_loss, - causal_mask=causal_mask, - pg_collection=None, - ) +class _SingleRankTP: + @staticmethod + def size(): + return 1 - baseline_loss = compute_loss(torch.tensor([0.0, 0.0]), torch.tensor([0.0, 0.0])) - window_perturbed_loss = compute_loss(torch.tensor([6.0, 0.0]), torch.tensor([0.0, 0.0])) - sink_perturbed_loss = compute_loss(torch.tensor([0.0, 0.0]), torch.tensor([6.0, 0.0])) - balanced_extreme_window_loss = compute_loss( - torch.tensor([30.0, 30.0]), torch.tensor([0.0, 0.0]) - ) - predictor_tail_loss = compute_loss( - torch.tensor([30.0, 30.0]), torch.tensor([0.0, 0.0]), torch.tensor([[[0.0, -40.0]]]) - ) - torch.testing.assert_close(baseline_loss, torch.tensor(0.0112081543)) - torch.testing.assert_close(window_perturbed_loss, torch.tensor(0.4118546844)) - torch.testing.assert_close(sink_perturbed_loss, torch.tensor(0.4118546844)) - torch.testing.assert_close(balanced_extreme_window_loss, baseline_loss) - torch.testing.assert_close(predictor_tail_loss, torch.tensor(19.3068528)) - - -def test_unfused_csa_indexer_loss_ignores_rows_without_compressed_keys(): - """Early rows with no causal compressed key must have finite zero loss and gradient.""" - index_scores = torch.tensor([[[0.2, -0.1]]], requires_grad=True) - loss = _compute_unfused_csa_indexer_loss( - index_scores, - torch.tensor([[[0, 1]]]), - torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]), - torch.tensor([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0]])[:, None, :], - torch.tensor([0.0, 0.0]), - torch.tensor([[[0]]]), - compressed_kv_offset=1, - softmax_scale=1.0, - loss_coeff=1.0, - sparse_loss=True, - causal_mask=torch.full((1, 1, 2), float("-inf")), - pg_collection=None, - ) +class _SingleRankPG: + tp = _SingleRankTP() - loss.backward() - torch.testing.assert_close(loss, torch.tensor(0.0)) - torch.testing.assert_close(index_scores.grad, torch.zeros_like(index_scores)) - - -def test_unfused_csa_indexer_loss_preserves_sparse_and_dense_support(): - """Sparse loss uses selected compressed keys; dense loss uses every causal key.""" - query = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) - index_scores = torch.tensor([[[0.2, -0.1, -0.3]]]) - topk_indices = torch.tensor([[[0, 1]]]) - causal_mask = torch.zeros((1, 1, 3)) - - def compute_loss(third_compressed_kv, sparse_loss): - kv_full = torch.tensor([[0.0, 0.0], [2.0, 0.0], [0.0, 2.0], third_compressed_kv])[ - :, None, : - ] - return _compute_unfused_csa_indexer_loss( - index_scores, - topk_indices, - query, - kv_full, - torch.tensor([0.0, 0.0]), - torch.tensor([[[0]]]), - compressed_kv_offset=1, - softmax_scale=1.0, - loss_coeff=1.0, - sparse_loss=sparse_loss, - causal_mask=causal_mask, - pg_collection=None, - ) - sparse_baseline = compute_loss([0.0, 0.0], sparse_loss=True) - sparse_perturbed = compute_loss([10.0, 0.0], sparse_loss=True) - dense_baseline = compute_loss([0.0, 0.0], sparse_loss=False) - dense_perturbed = compute_loss([10.0, 0.0], sparse_loss=False) - - torch.testing.assert_close(sparse_perturbed, sparse_baseline) - assert torch.abs(dense_perturbed - dense_baseline) > 1e-3 - - -def test_unfused_csa_indexer_loss_thd_handles_per_segment_padding(): - """The THD wrapper must separate physical offsets from each segment's real length.""" - generator = torch.Generator().manual_seed(7) - q_indexer = torch.randn(24, 2, 2, generator=generator, requires_grad=True) - k_indexer = torch.randn(6, 2, generator=generator, requires_grad=True) - weights = torch.randn(24, 2, generator=generator, requires_grad=True) - query = torch.randn(24, 2, 2, generator=generator) - kv_full = torch.randn(30, 2, generator=generator) - window_indices = torch.stack( - [torch.tensor([max(0, position - 1), position]) for position in list(range(12)) * 2] +@pytest.mark.parametrize("layout", ["sbhd", "thd"]) +def test_unfused_csa_non_compressed_lse_matches_window_and_sink_oracle(layout): + torch.manual_seed(17) + num_heads, head_dim = 2, 4 + sink = torch.randn(num_heads, requires_grad=True) + + if layout == "sbhd": + seqlen_q, batch_size, n_kv = 3, 2, 5 + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + kv_full = torch.randn(n_kv, batch_size, head_dim, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2]], [[-1, 0], [0, 2], [2, 4]]]) + expected = torch.empty(batch_size, num_heads, seqlen_q) + with torch.no_grad(): + for batch in range(batch_size): + for row in range(seqlen_q): + for head in range(num_heads): + logits = [sink[head]] + for key_index in window_indices[batch, row]: + if key_index >= 0: + logits.append( + torch.dot(query[row, batch, head], kv_full[key_index, batch]) + ) + expected[batch, head, row] = torch.logsumexp(torch.stack(logits), dim=0) + else: + seqlen_q, batch_size, n_kv = 4, 1, 6 + query = torch.randn(seqlen_q, num_heads, head_dim, requires_grad=True) + kv_full = torch.randn(n_kv, head_dim, requires_grad=True) + window_indices = torch.tensor([[-1, 0], [0, 1], [2, 3], [4, 5]]) + expected = torch.empty(batch_size, num_heads, seqlen_q) + with torch.no_grad(): + for row in range(seqlen_q): + for head in range(num_heads): + logits = [sink[head]] + for key_index in window_indices[row]: + if key_index >= 0: + logits.append(torch.dot(query[row, head], kv_full[key_index])) + expected[0, head, row] = torch.logsumexp(torch.stack(logits), dim=0) + + actual = _compute_unfused_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=1.0 ) - topk_indices, loss = _unfused_csa_indexer_topk_and_loss_thd( - q_indexer, - k_indexer, - weights, - query, - kv_full, - torch.tensor([0.0, 1.0]), - window_indices, - torch.tensor([0, 12, 24], dtype=torch.int32), - torch.tensor([0, 9, 20], dtype=torch.int32), - torch.tensor([0, 3, 6], dtype=torch.int32), - torch.tensor([0, 12, 24], dtype=torch.int32), - torch.tensor([0, 15, 30], dtype=torch.int32), - topk=2, - softmax_scale=1.0, - loss_coeff=0.1, - sparse_loss=True, - ratio=4, - pg_collection=None, + assert actual.shape == (batch_size, num_heads, seqlen_q) + assert actual.dtype == torch.float32 + assert not actual.requires_grad + torch.testing.assert_close(actual, expected) + + for teacher_tensor in (query, kv_full, sink): + assert teacher_tensor.grad is None + + +def _independent_csa_indexer_loss( + index_scores, + topk_indices, + query, + compressed_kv, + window_kv, + window_indices, + sink, + *, + sparse_loss, + loss_coeff, +): + batch_size, seqlen_q, n_compressed = index_scores.shape + num_heads = query.shape[2] + losses = [] + for batch in range(batch_size): + for row in range(seqlen_q): + selected = ( + topk_indices[batch, row].tolist() if sparse_loss else list(range(n_compressed)) + ) + target = [] + for compressed_index in selected: + head_mass = 0.0 + for head in range(num_heads): + non_compressed_logits = [sink[head]] + for window_index in window_indices[batch, row]: + if window_index >= 0: + non_compressed_logits.append( + torch.dot(query[row, batch, head], window_kv[window_index, batch]) + ) + compressed_logits = [ + torch.dot(query[row, batch, head], compressed_kv[key_index, batch]) + for key_index in selected + ] + denominator = torch.logsumexp( + torch.stack(non_compressed_logits + compressed_logits), dim=0 + ) + head_mass = head_mass + torch.exp( + compressed_logits[selected.index(compressed_index)] - denominator + ) + target.append(head_mass) + target = torch.stack(target) + target = target / target.sum() + predict_log = torch.log_softmax(index_scores[batch, row, selected], dim=-1) + losses.append((target * (torch.log(target) - predict_log)).sum()) + return torch.stack(losses).mean() * loss_coeff + + +@pytest.mark.parametrize("sparse_loss", [False, True], ids=["dense", "sparse"]) +def test_csa_lse_custom_autograd_matches_full_teacher_and_reference(sparse_loss): + torch.manual_seed(29) + seqlen_q, batch_size, num_heads, head_dim = 4, 1, 2, 3 + n_compressed, index_heads, index_dim = 3, 2, 2 + index_topk, loss_coeff = 2, 0.7 + + q = torch.randn(seqlen_q, batch_size, index_heads, index_dim, requires_grad=True) + weights = torch.randn(seqlen_q, batch_size, index_heads, requires_grad=True) + k = torch.randn(n_compressed, batch_size, index_dim, requires_grad=True) + query = torch.randn(seqlen_q, batch_size, num_heads, head_dim, requires_grad=True) + window_kv = torch.randn(seqlen_q, batch_size, head_dim, requires_grad=True) + compressed_kv = torch.randn(n_compressed, batch_size, head_dim, requires_grad=True) + sink = torch.randn(num_heads, requires_grad=True) + window_indices = torch.tensor([[[-1, 0], [0, 1], [1, 2], [2, 3]]]) + kv_full = torch.cat([window_kv, compressed_kv], dim=0) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=1.0 + ) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, num_heads, -1) + compressed_mask = torch.zeros(seqlen_q, n_compressed) + + q_ref = q.detach().clone().requires_grad_(True) + weights_ref = weights.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + index_scores_ref, topk_ref = fused_qk_topk_naive(q_ref, k_ref, weights_ref, index_topk) + index_scores_oracle = index_scores_ref.detach().clone() + loss_ref = compute_dsa_indexer_loss( + index_scores_ref, + topk_ref, + query.detach(), + key_for_loss.detach(), + 1.0, + loss_coeff, + sparse_loss, + _SingleRankPG(), + mask=compressed_mask, + non_compressed_lse=non_compressed_lse, ) + loss_ref.backward() - def segment_oracle(q_start, seqlen_q, compressed_start, kv_full_start): - n_compressed = seqlen_q // 4 - causal_mask = torch.where( - torch.arange(n_compressed).unsqueeze(0) - >= torch.arange(1, seqlen_q + 1).unsqueeze(1) // 4, - float("-inf"), - 0.0, - ).unsqueeze(0) - segment_scores = torch.einsum( - "shd,td->sht", - q_indexer[q_start : q_start + seqlen_q].float(), - k_indexer[compressed_start : compressed_start + n_compressed].float(), - ) - segment_scores = ( - torch.relu(segment_scores) * weights[q_start : q_start + seqlen_q].float().unsqueeze(-1) - ).sum(dim=1) - segment_scores = segment_scores.unsqueeze(0) + causal_mask - segment_topk = segment_scores.topk(2, dim=-1).indices - return _compute_unfused_csa_indexer_loss( - segment_scores, - segment_topk, - query[q_start : q_start + seqlen_q].unsqueeze(1), - kv_full[kv_full_start : kv_full_start + 15].unsqueeze(1), - torch.tensor([0.0, 1.0]), - window_indices[q_start : q_start + seqlen_q].unsqueeze(0), - compressed_kv_offset=12, - softmax_scale=1.0, - loss_coeff=0.1, - sparse_loss=True, - causal_mask=causal_mask, - pg_collection=None, - ) + saved_shapes = [] - with torch.no_grad(): - expected_loss = (segment_oracle(0, 9, 0, 0) * 9 + segment_oracle(12, 11, 3, 15) * 11) / 24 + def pack_hook(tensor): + saved_shapes.append(tuple(tensor.shape)) + return tensor - assert topk_indices.shape == (24, 2) - torch.testing.assert_close( - topk_indices[[9, 10, 11, 23]], torch.full((4, 2), -1, dtype=torch.int32) + with torch.autograd.graph.saved_tensors_hooks(pack_hook, lambda tensor: tensor): + topk_actual, loss_actual = FusedDSAIndexerLoss.apply( + q, + weights, + k, + query, + key_for_loss, + 1.0, + index_topk, + loss_coeff, + compressed_mask, + sparse_loss, + _SingleRankPG(), + None, + None, + None, + None, + False, + True, + non_compressed_lse, + ) + loss_actual.backward() + + independent_loss = _independent_csa_indexer_loss( + index_scores_oracle, + topk_ref, + query.detach(), + compressed_kv.detach(), + window_kv.detach(), + window_indices, + sink.detach(), + sparse_loss=sparse_loss, + loss_coeff=loss_coeff, ) - assert torch.isfinite(loss) - assert loss > 0 - torch.testing.assert_close(loss, expected_loss) - loss.backward() - for grad in (q_indexer.grad, k_indexer.grad, weights.grad): - assert grad is not None - assert torch.isfinite(grad).all() - assert torch.count_nonzero(grad) > 0 - torch.testing.assert_close(q_indexer.grad[[9, 10, 11, 23]], torch.zeros(4, 2, 2)) - torch.testing.assert_close(weights.grad[[9, 10, 11, 23]], torch.zeros(4, 2)) - torch.testing.assert_close(k_indexer.grad[[2, 5]], torch.zeros(2, 2)) + + torch.testing.assert_close(loss_actual, independent_loss) + torch.testing.assert_close(loss_actual, loss_ref) + torch.testing.assert_close(topk_actual, topk_ref) + torch.testing.assert_close(q.grad, q_ref.grad) + torch.testing.assert_close(weights.grad, weights_ref.grad) + torch.testing.assert_close(k.grad, k_ref.grad) + assert (batch_size, seqlen_q, n_compressed) not in saved_shapes + assert tuple(non_compressed_lse.shape) in saved_shapes + + for teacher_tensor in (query, window_kv, compressed_kv, sink): + assert teacher_tensor.grad is None class TestGetWindowTopkIdxs: 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 cb1f9b15b0d..28445701845 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 @@ -467,7 +467,7 @@ def _native_unfused_sparse_indexer_loss( row_valid = selected_valid.any(dim=-1, keepdim=True) predict_logits = predict_logits.masked_fill(~selected_valid, float("-inf")) predict_logits = predict_logits.masked_fill(~row_valid, 0.0) - predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) * row_valid.float() + predict_log = F.log_softmax(predict_logits, dim=-1, dtype=torch.float32) compressed_attention_indices = torch.where( selected_valid, safe_compressed + compressed_kv_offset, torch.full_like(safe_compressed, -1) @@ -496,12 +496,10 @@ def _native_unfused_sparse_indexer_loss( compressed_width = compressed_attention_indices.shape[-1] target = attention_probs[..., -compressed_width:].sum(dim=1) - eps = torch.finfo(torch.float32).tiny + eps = 1e-10 target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) target = target * row_valid.float() - 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 = (target * (torch.log(target.clamp(min=eps)) - predict_log)).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 * loss_coeff @@ -1481,12 +1479,13 @@ def test_thd_padded_attention_matches_unpadded( pad_offset += pl hidden_padded = hidden_padded.clone().requires_grad_(True) - # cu_seqlens_q: cumulative unpadded lengths, matching the metadata - # produced by data_schedule. Physical boundaries live separately in - # cu_seqlens_q_padded. + # cu_seqlens_q: unpadded real boundaries (cumsum of real lengths + # within the padded physical layout). cu_seqlens_q_vals = [0] - for real_length in seg_lens: - cu_seqlens_q_vals.append(cu_seqlens_q_vals[-1] + real_length) + offset = 0 + for rl, pl in zip(seg_lens, padded_seg_lens): + cu_seqlens_q_vals.append(offset + rl) + offset += pl cu_seqlens_unpadded = torch.tensor(cu_seqlens_q_vals, dtype=torch.int32, device='cuda') # cu_seqlens_q_padded: padded boundaries (cumsum of padded lengths From 3ae4e2bb49fee19b6cefa1d562af438e9340c5a9 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Thu, 30 Jul 2026 18:50:15 +0800 Subject: [PATCH 3/8] Harden CSA indexer loss edge cases Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../experimental_attention_variant/csa.py | 159 ++++++++++------ .../experimental_attention_variant/dsa.py | 19 +- .../test_csa_indexer_loss.py | 175 ++++++++++++++++++ 3 files changed, 295 insertions(+), 58 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 2ee614ef1bd..2803a457145 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -118,15 +118,20 @@ def _get_csa_compressed_capacity( def _build_compressed_thd_indexer_metadata( - cu_seqlens_q: torch.Tensor, cu_seqlens_compressed: torch.Tensor, total_q: int, ratio: int + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, + total_q: int, + ratio: int, + cu_seqlens_q_unpadded: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Lower packed compressed-THD metadata to DSA row-wise varlen bounds. Compressed keys are concatenated segment-major. For query row ``i`` in segment ``b``, the legal global compressed-key interval is ``[cu_comp[b], cu_comp[b] + floor((pos_i + 1) / ratio))``. Rows beyond - ``cu_seqlens_q[-1]`` are static-capacity padding and receive an empty - interval plus ``query_valid_rows=False``. + their segment's unpadded length are static-capacity padding and receive + an empty interval plus ``query_valid_rows=False``. Physical query and + compressed offsets always come from the padded cumulative lengths. Returns ``(starts, ends, query_valid_rows, compressed_offset_per_row)``. The final tensor converts DSA's global compressed indices back to the @@ -142,6 +147,13 @@ def _build_compressed_thd_indexer_metadata( ) if cu_seqlens_q.device != cu_seqlens_compressed.device: raise ValueError("query and compressed cumulative lengths must be on the same device") + if cu_seqlens_q_unpadded is not None: + if cu_seqlens_q_unpadded.ndim != 1 or cu_seqlens_q_unpadded.shape != cu_seqlens_q.shape: + raise ValueError( + "unpadded query cumulative lengths must match the physical query segment shape" + ) + if cu_seqlens_q_unpadded.device != cu_seqlens_q.device: + raise ValueError("physical and unpadded query lengths must be on the same device") device = cu_seqlens_q.device row_idx = torch.arange(total_q, device=device, dtype=torch.int64) @@ -150,8 +162,12 @@ def _build_compressed_thd_indexer_metadata( compressed_starts = cu_seqlens_compressed[segment_ids].to(dtype=torch.int64) compressed_limits = cu_seqlens_compressed[segment_ids + 1].to(dtype=torch.int64) - query_valid_rows = row_idx < cu_seqlens_q[-1].to(dtype=torch.int64) position_in_segment = row_idx - query_starts + if cu_seqlens_q_unpadded is None: + query_valid_rows = row_idx < cu_seqlens_q[-1].to(dtype=torch.int64) + else: + real_lens = (cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1]).to(dtype=torch.int64) + query_valid_rows = position_in_segment < real_lens[segment_ids] causal_ends = compressed_starts + torch.div( position_in_segment + 1, ratio, rounding_mode="floor" ) @@ -692,6 +708,7 @@ def _compute_unfused_csa_non_compressed_lse( attn_sink: torch.Tensor, window_indices: torch.Tensor, softmax_scale: float, + chunk_size: int = 512, ) -> torch.Tensor: """Return detached sliding-window-plus-sink log mass for the CSA teacher. @@ -699,41 +716,88 @@ def _compute_unfused_csa_non_compressed_lse( both SBHD and packed THD inputs. THD callers must pass flat-global window indices into ``kv_full``. """ - is_thd = query.ndim == 3 - if is_thd: - q_flat = query - kv_flat = kv_full - global_indices = window_indices - batch_size = 1 - seqlen_q = query.shape[0] - else: + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if attn_sink.ndim != 1: + raise ValueError(f"attn_sink must be 1D, got shape {tuple(attn_sink.shape)}") + + if query.ndim == 4: seqlen_q, batch_size, num_heads, head_dim = query.shape + if kv_full.ndim != 3 or kv_full.shape[1:] != (batch_size, head_dim): + raise ValueError( + "SBHD non-compressed LSE expects kv_full shape " + f"[sk, {batch_size}, {head_dim}], got {tuple(kv_full.shape)}" + ) + if window_indices.ndim != 3 or window_indices.shape[:2] != (batch_size, seqlen_q): + raise ValueError( + "SBHD non-compressed LSE expects window_indices shape " + f"[{batch_size}, {seqlen_q}, window], got {tuple(window_indices.shape)}" + ) + n_kv = kv_full.shape[0] - q_flat = query.permute(1, 0, 2, 3).reshape(batch_size * seqlen_q, num_heads, head_dim) - kv_flat = kv_full.permute(1, 0, 2).reshape(batch_size * n_kv, head_dim) - valid = window_indices >= 0 - batch_offsets = torch.arange(batch_size, device=query.device).view(batch_size, 1, 1) * n_kv - global_indices = torch.where(valid, window_indices + batch_offsets, window_indices).reshape( - batch_size * seqlen_q, -1 + q_flat = query.detach().permute(1, 0, 2, 3).reshape(-1, num_heads, head_dim) + kv_flat = kv_full.detach().permute(1, 0, 2).reshape(-1, head_dim) + batch_offsets = ( + torch.arange(batch_size, device=window_indices.device, dtype=torch.int64) * n_kv + ).view(batch_size, 1, 1) + window_indices_i64 = window_indices.to(dtype=torch.int64) + global_indices = torch.where( + window_indices_i64 >= 0, window_indices_i64 + batch_offsets, window_indices_i64 + ).reshape(batch_size * seqlen_q, -1) + output_shape = (batch_size, num_heads, seqlen_q) + elif query.ndim == 3: + total_q, num_heads, head_dim = query.shape + if kv_full.ndim != 2 or kv_full.shape[1] != head_dim: + raise ValueError( + "THD non-compressed LSE expects kv_full shape " + f"[total_kv, {head_dim}], got {tuple(kv_full.shape)}" + ) + if window_indices.ndim != 2 or window_indices.shape[0] != total_q: + raise ValueError( + "THD non-compressed LSE expects window_indices shape " + f"[{total_q}, window], got {tuple(window_indices.shape)}" + ) + + q_flat = query.detach() + kv_flat = kv_full.detach() + global_indices = window_indices.to(dtype=torch.int64) + output_shape = (1, num_heads, total_q) + else: + raise ValueError( + "non-compressed LSE query must be SBHD [sq, b, h, d] or " + f"THD [total_q, h, d], got {tuple(query.shape)}" ) - rows, num_heads, head_dim = q_flat.shape - safe_indices = global_indices.clamp(min=0).long() - gathered_kv = torch.gather( - kv_flat.unsqueeze(0).expand(rows, -1, -1), - dim=1, - index=safe_indices.unsqueeze(-1).expand(-1, -1, head_dim), - ) - window_logits = torch.einsum("rhd,rkd->rhk", q_flat.float(), gathered_kv.float()) - window_logits = (window_logits * softmax_scale).masked_fill( - ~(global_indices >= 0).unsqueeze(1), float("-inf") - ) - window_lse = torch.logsumexp(window_logits, dim=-1) - non_compressed_lse = torch.logaddexp(window_lse, attn_sink.detach().view(1, num_heads).float()) + if attn_sink.numel() != num_heads: + raise ValueError(f"attn_sink must contain {num_heads} head values, got {attn_sink.numel()}") + if not (query.device == kv_full.device == attn_sink.device == global_indices.device): + raise ValueError("query, kv_full, attn_sink, and window_indices must be on the same device") + + sink = attn_sink.detach().to(dtype=torch.float32).view(1, num_heads) + lse_chunks = [] + num_rows = q_flat.shape[0] + for start in range(0, num_rows, chunk_size): + end = min(start + chunk_size, num_rows) + indices = global_indices[start:end] + safe_indices = indices.clamp(min=0) + gathered_kv = kv_flat.index_select(0, safe_indices.reshape(-1)).reshape( + end - start, indices.shape[-1], head_dim + ) + window_logits = torch.einsum("rhd,rkd->rhk", q_flat[start:end].float(), gathered_kv.float()) + window_logits = (window_logits * softmax_scale).masked_fill( + (indices < 0).unsqueeze(1), float("-inf") + ) + window_lse = torch.logsumexp(window_logits, dim=-1) + lse_chunks.append(torch.logaddexp(window_lse, sink)) - if is_thd: - return non_compressed_lse.transpose(0, 1).unsqueeze(0).contiguous() - return non_compressed_lse.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() + if lse_chunks: + lse_flat = torch.cat(lse_chunks, dim=0) + else: + lse_flat = torch.empty((0, num_heads), dtype=torch.float32, device=query.device) + + if query.ndim == 4: + return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() + return lse_flat.transpose(0, 1).unsqueeze(0).reshape(output_shape).contiguous() def _unfused_indexer_sparse_attn_from_topk( @@ -1762,7 +1826,7 @@ def _forward_unfused_csa( key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale non_compressed_lse = _compute_unfused_csa_non_compressed_lse( - query, kv_full, self.attn_sink, window_idxs, self.softmax_scale + query, kv_full[:offset], self.attn_sink, window_idxs, self.softmax_scale ) topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( q_indexer, @@ -2087,30 +2151,15 @@ def _forward_unfused_csa_thd( weights_for_unfused = w_thd * self.indexer.softmax_scale indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) - # Build loss metadata from the original unpadded lengths. - # Indexer K may retain static-capacity tail rows, but DSA's - # row-wise bounds never admit those keys and - # ``query_valid_rows`` zeroes padded query gradients. - cu_seqlens_q_for_loss = packed_seq_params.cu_seqlens_q - seg_lens_q = cu_seqlens_q_for_loss[1:] - cu_seqlens_q_for_loss[:-1] - cu_seqlens_compressed_idx_for_loss = torch.cat( - [ - torch.zeros( - 1, - dtype=cu_seqlens_q_for_loss.dtype, - device=cu_seqlens_q_for_loss.device, - ), - (seg_lens_q // self.compress_ratio) - .cumsum(0) - .to(cu_seqlens_q_for_loss.dtype), - ] - ) + # Physical padded offsets define the packed address space; + # unpadded lengths identify the real rows within each segment. (varlen_starts, varlen_ends, query_valid_rows, compressed_offsets) = ( _build_compressed_thd_indexer_metadata( - cu_seqlens_q_for_loss, - cu_seqlens_compressed_idx_for_loss, + cu_seqlens_q, + cu_seqlens_compressed_idx, total_q=q_thd.shape[0], ratio=self.compress_ratio, + cu_seqlens_q_unpadded=packed_seq_params.cu_seqlens_q, ) ) flat_window_idxs, _ = build_flat_topk_idxs( diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 76a7a531c95..0ae46478ce0 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -570,7 +570,7 @@ def compute_dsa_indexer_loss( # attention scores are scattered to TP ranks in head dimension. torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) # The target is already non-negative because it is a sum of softmax probabilities. - attention_scores = dsa_indexer_loss.normalize_indexer_target(attention_scores) + attention_scores = _normalize_indexer_teacher_target(attention_scores, non_compressed_lse) return dsa_indexer_loss.indexer_loss_from_target( attention_scores, index_log_scores, @@ -617,6 +617,15 @@ def _compute_indexer_teacher_probabilities( return torch.where(expanded_valid_mask, probabilities, torch.zeros_like(probabilities)) +def _normalize_indexer_teacher_target( + target: torch.Tensor, non_compressed_lse: torch.Tensor | None +) -> torch.Tensor: + """L1-normalize teacher mass without changing the legacy DSA path.""" + if non_compressed_lse is None: + return dsa_indexer_loss.normalize_indexer_target(target) + return target / target.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) + + def _compute_index_scores( q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True ) -> torch.Tensor: @@ -882,7 +891,9 @@ def bwd_fused_indexer_loss_naive( # L1 normalize. Fully masked packed/varlen rows can have zero summed # attention mass; clamp the denominator so those rows stay finite and are # later zeroed by the row-valid loss mask. - attention_scores_normalized = dsa_indexer_loss.normalize_indexer_target(attention_scores_sum) + attention_scores_normalized = _normalize_indexer_teacher_target( + attention_scores_sum, non_compressed_lse + ) # Free attention_scores_sum - no longer needed after normalization del attention_scores_sum @@ -1068,6 +1079,7 @@ def forward( ctx.query_valid_rows = query_valid_rows ctx.calculate_per_token_loss = calculate_per_token_loss ctx.use_relu = use_relu + ctx.num_inputs = len(ctx.needs_input_grad) return topk_indices, loss @@ -1110,7 +1122,8 @@ def backward(ctx, grad_topk_indices, grad_loss): "key": None, "non_compressed_lse": None, } - return tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + gradients = tuple(grad_by_name.get(name) for name in _FUSED_DSA_INDEXER_LOSS_INPUT_NAMES) + return gradients[: ctx.num_inputs] class DSAIndexerLossAutoScaler(torch.autograd.Function): diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_loss.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_loss.py index bc2cac80674..9021a1b9481 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_loss.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_indexer_loss.py @@ -9,9 +9,14 @@ from megatron.core.transformer.experimental_attention_variant.csa import ( _build_compressed_thd_indexer_metadata, _compressed_thd_topk_to_local, + _compute_unfused_csa_non_compressed_lse, +) +from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + build_flat_topk_idxs, ) from megatron.core.transformer.experimental_attention_variant.dsa import ( FusedDSAIndexerLoss, + _normalize_indexer_teacher_target, fused_qk_topk_naive, ) from tests.unit_tests.test_utilities import Utils @@ -23,6 +28,60 @@ def _causal_compressed_mask(sq, sk, ratio, device): return torch.where(columns >= positions // ratio, float('-inf'), 0.0).unsqueeze(0) +class _SingleRankGroup: + @staticmethod + def size(): + return 1 + + +class _SingleRankPGCollection: + tp = _SingleRankGroup() + + +def _make_legacy_abi_case(seed=23): + generator = torch.Generator().manual_seed(seed) + sq, sk, batch = 8, 2, 1 + index_heads, index_dim = 2, 4 + attention_heads, attention_dim = 2, 5 + q = torch.randn(sq, batch, index_heads, index_dim, generator=generator) + weights = torch.randn(sq, batch, index_heads, generator=generator) + k = torch.randn(sk, batch, index_dim, generator=generator) + query = torch.randn(sq, batch, attention_heads, attention_dim, generator=generator) + key = torch.randn(sk, batch, attention_heads, attention_dim, generator=generator) + mask = _causal_compressed_mask(sq, sk, ratio=4, device='cpu') + return q, weights, k, query, key, mask + + +def _run_legacy_abi_case(num_inputs): + q, weights, k, query, key, mask = _make_legacy_abi_case() + q = q.requires_grad_(True) + weights = weights.requires_grad_(True) + k = k.requires_grad_(True) + args = [ + q, + weights, + k, + query, + key, + key.shape[-1] ** -0.5, + 2, + 0.7, + mask, + True, + _SingleRankPGCollection(), + None, + None, + None, + None, + False, + True, + None, + ] + topk_indices, loss = FusedDSAIndexerLoss.apply(*args[:num_inputs]) + loss.backward() + return topk_indices, loss.detach(), q.grad, weights.grad, k.grad + + def _run_compressed_thd_topk(q, k, weights, topk, cu_q, cu_comp, ratio): starts, ends, valid_rows, offsets = _build_compressed_thd_indexer_metadata( cu_q, cu_comp, total_q=q.shape[0], ratio=ratio @@ -80,6 +139,101 @@ def _run_compressed_thd_loss( return _compressed_thd_topk_to_local(topk_global, offsets), loss +@pytest.mark.parametrize("num_inputs", [14, 17, 18]) +def test_custom_autograd_backward_matches_legacy_arity(num_inputs): + reference = _run_legacy_abi_case(18) + actual = _run_legacy_abi_case(num_inputs) + for actual_tensor, reference_tensor in zip(actual, reference): + torch.testing.assert_close(actual_tensor, reference_tensor, rtol=0, atol=0) + + +def test_corrected_teacher_normalizes_tiny_positive_mass(): + target = torch.tensor([[[1.0e-30, 2.0e-30]]], dtype=torch.float32) + + corrected = _normalize_indexer_teacher_target( + target, torch.zeros((1, 1, 1), dtype=torch.float32) + ) + legacy = _normalize_indexer_teacher_target(target, None) + + torch.testing.assert_close( + corrected, torch.tensor([[[1.0 / 3.0, 2.0 / 3.0]]]), rtol=1e-6, atol=0 + ) + assert corrected.sum().item() == pytest.approx(1.0) + torch.testing.assert_close( + legacy, target / torch.tensor(1.0e-10, dtype=torch.float32), rtol=0, atol=0 + ) + + +def test_non_compressed_lse_chunking_and_all_invalid_window(): + generator = torch.Generator().manual_seed(31) + query = torch.randn(5, 2, 3, 4, generator=generator) + kv = torch.randn(5, 2, 4, generator=generator) + sink = torch.randn(3, generator=generator) + window_indices = torch.tensor( + [[[0, -1], [0, 1], [1, 2], [2, 3], [3, 4]], [[0, -1], [0, 1], [1, 2], [2, 3], [3, 4]]] + ) + + row_chunks = _compute_unfused_csa_non_compressed_lse( + query, kv, sink, window_indices, softmax_scale=0.5, chunk_size=1 + ) + full_chunk = _compute_unfused_csa_non_compressed_lse( + query, kv, sink, window_indices, softmax_scale=0.5, chunk_size=10 + ) + torch.testing.assert_close(row_chunks, full_chunk, rtol=0, atol=0) + + all_invalid = torch.full_like(window_indices, -1) + sink_only = _compute_unfused_csa_non_compressed_lse( + query, kv, sink, all_invalid, softmax_scale=0.5 + ) + expected = sink.view(1, -1, 1).expand(2, -1, 5) + torch.testing.assert_close(sink_only, expected, rtol=0, atol=0) + + +def test_thd_non_compressed_lse_uses_global_interleaved_kv_offsets(): + query = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.5, 0.5], [1.0, -1.0]], + [[1.0, 1.0], [-1.0, 0.5]], + [[0.25, -0.5], [0.75, 1.0]], + ] + ) + kv_full_thd = torch.tensor( + [ + [1.0, 0.0], + [0.0, 2.0], + [100.0, 100.0], + [3.0, 0.0], + [0.0, 4.0], + [200.0, 200.0], + [300.0, 300.0], + ] + ) + local_window = torch.tensor([[0, -1], [0, 1], [0, -1], [0, 1]]) + cu_q = torch.tensor([0, 2, 4], dtype=torch.int32) + cu_kv_full = torch.tensor([0, 3, 7], dtype=torch.int32) + global_window, _ = build_flat_topk_idxs( + local_window, batch_size=-1, cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv_full + ) + assert torch.equal( + global_window, torch.tensor([[0, -1], [0, 1], [3, -1], [3, 4]], dtype=torch.int32) + ) + + sink = torch.tensor([0.25, -0.5]) + actual = _compute_unfused_csa_non_compressed_lse( + query, kv_full_thd, sink, global_window, softmax_scale=0.5, chunk_size=1 + ) + gathered = kv_full_thd.index_select(0, global_window.clamp_min(0).reshape(-1)).reshape(4, 2, 2) + logits = torch.einsum('rhd,rkd->rhk', query.float(), gathered.float()) * 0.5 + logits = logits.masked_fill((global_window < 0).unsqueeze(1), float('-inf')) + expected = ( + torch.logaddexp(torch.logsumexp(logits, dim=-1), sink.view(1, -1)) + .transpose(0, 1) + .unsqueeze(0) + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + class TestCompressedThdMetadata: def test_two_segments_and_padding_map_exactly(self): cu_q = torch.tensor([0, 8, 13], dtype=torch.int32) @@ -94,6 +248,27 @@ def test_two_segments_and_padding_map_exactly(self): assert torch.equal(valid_rows, torch.tensor([True] * 13 + [False, False])) assert torch.equal(offsets, starts) + def test_per_segment_padding_uses_physical_offsets_and_real_lengths(self): + physical_cu_q = torch.tensor([0, 12, 24], dtype=torch.int32) + real_cu_q = torch.tensor([0, 9, 20], dtype=torch.int32) + physical_cu_comp = torch.tensor([0, 3, 6], dtype=torch.int32) + + starts, ends, valid_rows, offsets = _build_compressed_thd_indexer_metadata( + physical_cu_q, physical_cu_comp, total_q=24, ratio=4, cu_seqlens_q_unpadded=real_cu_q + ) + + expected_valid = torch.tensor([True] * 9 + [False] * 3 + [True] * 11 + [False]) + assert torch.equal(valid_rows, expected_valid) + assert torch.equal(starts[:9], torch.zeros(9, dtype=torch.int64)) + assert torch.equal(starts[9:12], torch.zeros(3, dtype=torch.int64)) + assert torch.equal(starts[12:23], torch.full((11,), 3, dtype=torch.int64)) + assert starts[23] == 0 + assert torch.equal(offsets, starts) + assert ends[8] == 2 + assert torch.equal(ends[9:12], torch.zeros(3, dtype=torch.int64)) + assert ends[22] == 5 + assert ends[23] == 0 + @pytest.mark.parametrize( "cu_q,cu_comp,ratio,match", [ From 40e1c3c5f4514928725f887e2697f2e05a3b139a Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 1 Aug 2026 18:28:48 +0800 Subject: [PATCH 4/8] Add CP2 unfused indexer precision regression Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../test_dsv4_hybrid_attention_cp.py | 145 +++++++++++++++++- 1 file changed, 144 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 7f2378df896..5e8e1d88c0f 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -11,13 +11,17 @@ import torch import torch.distributed as dist import torch.nn.functional as F +from tests.unit_tests.test_utilities import Utils import megatron.core.parallel_state as parallel_state from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from tests.unit_tests.test_utilities import Utils +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, +) from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_attention import ( _SEED, _build_attention, @@ -549,6 +553,145 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): self._cp1_graph_time_cache[layer_number] = measured return measured + @pytest.mark.parametrize("sparse_loss", [True, False], ids=["sparse", "dense"]) + def test_cp2_unfused_indexer_loss_and_grad_matches_cp1(self, sparse_loss): + """CP2 unfused indexer loss and gradients must match the CP1 teacher.""" + if self.cp_size != 2: + pytest.skip("The exact unfused indexer regression is a focused CP2 gate") + + seq_len = 64 + packed = _make_thd_packed_seq_params((seq_len,)) + local_rows = seq_len // self.cp_size + local_idx = torch.arange( + self.cp_rank * local_rows, + (self.cp_rank + 1) * local_rows, + device='cuda', + ) + + common_config = dict( + num_layers=1, + csa_compress_ratios=[4], + csa_window_size=8, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=sparse_loss, + csa_dense_mode=False, + csa_compress_rotary_base=40000, + layernorm_epsilon=1e-6, + normalization="RMSNorm", + qk_layernorm=True, + layernorm_zero_centered_gamma=False, + expert_model_parallel_size=1, + dsa_kernel_backend="none", + apply_rope_fusion=False, + ) + config_cp = _make_config( + **common_config, + context_parallel_size=self.cp_size, + cp_partition_mode="contiguous", + sequence_packing_scheduler="dp_balanced", + ) + config_ref = _make_config( + **common_config, + context_parallel_size=1, + cp_partition_mode="zigzag", + sequence_packing_scheduler=None, + ) + + torch.manual_seed(_SEED + 5960) + model_parallel_cuda_manual_seed(_SEED + 5960) + cp_attn = _build_attention(config_cp, layer_number=1, pg_collection=self.pg).cuda() + ref_attn = _build_attention(config_ref, layer_number=1, pg_collection=self.ref_pg).cuda() + _copy_module_parameters(cp_attn, ref_attn) + cp_attn.train() + ref_attn.train() + + full_hidden = torch.randn( + seq_len, 1, config_cp.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + local_hidden = full_hidden.index_select(0, local_idx).detach().clone().requires_grad_(True) + ref_hidden = full_hidden.detach().clone().requires_grad_(True) + + DSAIndexerLossAutoScaler.set_loss_scale(torch.ones((), device='cuda')) + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + local_out, _ = cp_attn( + hidden_states=local_hidden, attention_mask=None, packed_seq_params=packed + ) + local_loss = DSAIndexerLossLoggingHelper.tracker["values"][0].detach().float().clone() + + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + ref_out, _ = ref_attn( + hidden_states=ref_hidden, attention_mask=None, packed_seq_params=packed + ) + ref_loss = DSAIndexerLossLoggingHelper.tracker["values"][0].detach().float().clone() + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + + cp_loss = local_loss.clone() + dist.all_reduce(cp_loss, group=self.pg.cp) + + # A zero upstream gradient isolates the indexer-loss backward injected + # by DSAIndexerLossAutoScaler from the main attention objective. + local_out.backward(torch.zeros_like(local_out)) + ref_out.backward(torch.zeros_like(ref_out)) + + cp_params = dict(cp_attn.core_attention.indexer.named_parameters()) + ref_params = dict(ref_attn.core_attention.indexer.named_parameters()) + reduced_cp_grads = {} + missing_cp_grads = [] + for name, param in sorted(cp_params.items()): + if param.grad is None: + missing_cp_grads.append(name) + grad = torch.zeros_like(param, dtype=torch.float32) + else: + grad = param.grad.detach().float().clone() + dist.all_reduce(grad, group=self.pg.cp) + reduced_cp_grads[name] = grad + + # Complete every collective before evaluating failures so one rank + # cannot leave its peers blocked in an all-reduce. + errors = [] + loss_diff = (cp_loss - ref_loss).abs() + if not torch.isfinite(cp_loss) or not torch.isfinite(ref_loss): + errors.append(f"non-finite loss: cp2={cp_loss.item()}, cp1={ref_loss.item()}") + elif cp_loss <= 0 or ref_loss <= 0: + errors.append(f"non-positive loss: cp2={cp_loss.item()}, cp1={ref_loss.item()}") + if loss_diff > 1e-4: + errors.append( + f"indexer loss diff={loss_diff.item():.10e} exceeds 1e-4 " + f"(cp2={cp_loss.item():.10f}, cp1={ref_loss.item():.10f})" + ) + + if set(cp_params) != set(ref_params): + errors.append( + f"indexer parameter names differ: cp2={sorted(cp_params)}, cp1={sorted(ref_params)}" + ) + if missing_cp_grads: + errors.append(f"missing CP2 indexer gradients: {missing_cp_grads}") + + nonzero_grad = False + for name, cp_grad in reduced_cp_grads.items(): + ref_param = ref_params.get(name) + if ref_param is None or ref_param.grad is None: + errors.append(f"missing CP1 indexer gradient: {name}") + continue + ref_grad = ref_param.grad.detach().float() + nonzero_grad = nonzero_grad or bool(torch.count_nonzero(ref_grad).item()) + close = torch.isclose(cp_grad, ref_grad, rtol=1e-3, atol=1e-4) + if not bool(close.all().item()): + max_abs = (cp_grad - ref_grad).abs().max().item() + errors.append(f"indexer gradient {name} max_abs={max_abs:.10e}") + if not nonzero_grad: + errors.append("all CP1 indexer gradients are zero") + + failure_flag = torch.tensor(bool(errors), dtype=torch.int32, device='cuda') + dist.all_reduce(failure_flag, op=dist.ReduceOp.MAX, group=self.pg.cp) + if failure_flag.item(): + pytest.fail("; ".join(errors) if errors else "peer rank reported an alignment failure") + + del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out + DSAIndexerLossAutoScaler.main_loss_backward_scale = None + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + _clear_cuda_test_state() + @pytest.mark.parametrize( ("layer_number", "sparse_loss"), [(1, True), (2, True), (2, False), (3, True)], From 4a9222a537e41c2df1a32aaf3e76372b01a0dc2a Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Sat, 1 Aug 2026 19:39:57 +0800 Subject: [PATCH 5/8] Fix CP unfused CSA indexer loss alignment Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../experimental_attention_variant/csa.py | 89 +++++++++++++------ .../test_dsv4_hybrid_attention_cp.py | 58 ++++++++---- 2 files changed, 102 insertions(+), 45 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 2803a457145..6d98333b11f 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -29,6 +29,8 @@ DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, FusedDSAIndexerLoss, + _compute_indexer_teacher_probabilities, + _normalize_indexer_teacher_target, fused_qk_topk_naive, rotate_activation, ) @@ -813,12 +815,13 @@ def _unfused_indexer_sparse_attn_from_topk( softmax_scale: float, indexer_softmax_scale: float, loss_coeff: float, - loss_divisor: float, + loss_divisor: Union[float, torch.Tensor], sparse_loss: bool, ratio: int, _max_seqlen_q: int, indexer_layout: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], q_padding_mask: Optional[torch.Tensor] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """PyTorch sparse attention plus caller-supplied indexer loss for THD CP. @@ -832,6 +835,36 @@ def _unfused_indexer_sparse_attn_from_topk( total_q, np_, hn = query.shape cu_seqlens_q, cu_seqlens_k, q_causal_offsets = indexer_layout + indexer_topk = indexer_topk_indices.shape[-1] + attention_indices = topk_indices + if q_padding_mask is not None: + attention_indices = attention_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, kv_full, attn_sink, attention_indices[:, indexer_topk:], softmax_scale + ) + + def compressed_teacher_target( + logits: torch.Tensor, valid_mask: torch.Tensor, row_slice: slice + ) -> torch.Tensor: + lse_slice = non_compressed_lse[:, :, row_slice] + probabilities = _compute_indexer_teacher_probabilities( + logits.permute(1, 0, 2).unsqueeze(0), + valid_mask.unsqueeze(0), + non_compressed_lse=lse_slice, + ) + target = probabilities.sum(dim=1).squeeze(0) + if tp_group is not None and tp_group.size() > 1: + torch.distributed.all_reduce(target, group=tp_group) + return _normalize_indexer_teacher_target(target, lse_slice) + + def scale_local_loss(raw_local_loss: torch.Tensor) -> torch.Tensor: + if torch.is_tensor(loss_divisor): + divisor = loss_divisor.to(device=raw_local_loss.device, dtype=torch.float32).clamp_min( + 1.0 + ) + else: + divisor = max(float(loss_divisor), 1.0) + return raw_local_loss * float(loss_coeff) / divisor if not sparse_loss: q_rows = torch.arange(total_q, dtype=cu_seqlens_q.dtype, device=query.device) @@ -879,15 +912,9 @@ def _unfused_indexer_sparse_attn_from_topk( target_logits = (target_logits * softmax_scale).masked_fill( ~valid.unsqueeze(1), float("-inf") ) - target_logits = torch.where(row_valid.unsqueeze(1), target_logits, 0.0) - sink = attn_sink.detach().view(1, np_, 1).float() - scores_max = torch.maximum(target_logits.max(dim=-1, keepdim=True).values, sink) - exp_scores = torch.exp(target_logits - scores_max) - exp_sink = torch.exp(sink - scores_max) - target = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) - target = (target * row_valid.unsqueeze(1).float()).sum(dim=1) + target = compressed_teacher_target(target_logits, valid, slice(start, end)) eps = torch.finfo(torch.float32).tiny - target = target / target.sum(dim=-1, keepdim=True).clamp(min=eps) + target = target * row_valid.float() target = target.clamp(min=eps) predict = predict.clamp(min=eps) @@ -897,9 +924,7 @@ def _unfused_indexer_sparse_attn_from_topk( ) raw_local_loss = raw_local_loss + kl_per_row.sum() - return output, raw_local_loss * float(loss_coeff) / float(loss_divisor) - - indexer_topk = indexer_topk_indices.shape[-1] + return output, scale_local_loss(raw_local_loss) if q_padding_mask is not None: indexer_topk_indices = indexer_topk_indices.masked_fill(q_padding_mask.unsqueeze(-1), -1) @@ -931,19 +956,12 @@ def _unfused_indexer_sparse_attn_from_topk( selected_kv = compressed_kv.detach().index_select(0, safe_indices.reshape(-1)) selected_kv = selected_kv.reshape(total_q, indexer_topk, hn) - # This fallback does not run FlashMLA, so it cannot reuse the fused path's - # lse_indexer-based target helper. Recompute the same sink-aware target and - # KL terms explicitly to keep the no-fusion path numerically aligned. + # Normalize selected compressed logits with the same sliding-window and + # sink mass used by the real CSA attention teacher. attn_scores = torch.einsum("rhd,rkd->rhk", query.detach().float(), selected_kv.float()) attn_scores = attn_scores * softmax_scale attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) - sink = attn_sink.detach().view(1, np_, 1).float() - scores_max = torch.maximum(attn_scores.max(dim=-1, keepdim=True).values, sink) - exp_scores = torch.exp(attn_scores - scores_max) - exp_sink = torch.exp(sink - scores_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 = compressed_teacher_target(attn_scores, valid, slice(None)) target = target * row_valid.float() eps = torch.finfo(torch.float32).tiny @@ -953,7 +971,7 @@ def _unfused_indexer_sparse_attn_from_topk( kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) raw_local_loss = kl_per_row.sum() - indexer_loss = raw_local_loss * float(loss_coeff) / float(loss_divisor) + indexer_loss = scale_local_loss(raw_local_loss) return output, indexer_loss @@ -2694,11 +2712,16 @@ def _forward_thd_cp( real_seqlens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] positions = global_rows - cu_seqlens[batch_ids] q_padding_mask = positions >= real_seqlens[batch_ids] - output, indexer_loss = ( - FusedCSAIndexerSparseAttnFromTopkFunc.apply - if self.use_fused_kernels - else _unfused_indexer_sparse_attn_from_topk - )( + loss_divisor = 1 if self.config.calculate_per_token_loss else l_local * cp_size + if ( + not self.use_fused_kernels + and not self.config.calculate_per_token_loss + and cu_seqlens_q_unpadded is not None + ): + # Match the reference DSA mean over real query rows. Keep the + # scalar on device so padded THD remains CUDA-graph safe. + loss_divisor = cu_seqlens_q_unpadded[-1] + indexer_loss_args = ( query, kv_full_thd, self.attn_sink.float(), @@ -2711,13 +2734,21 @@ def _forward_thd_cp( self.softmax_scale, indexer.softmax_scale, indexer_loss_coeff, - 1 if self.config.calculate_per_token_loss else l_local * cp_size, + loss_divisor, sparse_indexer_loss, ratio, max_seqlen_q, indexer_layout, q_padding_mask, ) + if self.use_fused_kernels: + output, indexer_loss = FusedCSAIndexerSparseAttnFromTopkFunc.apply( + *indexer_loss_args + ) + else: + output, indexer_loss = _unfused_indexer_sparse_attn_from_topk( + *indexer_loss_args, tp_group=indexer.pg_collection.tp + ) if indexer_loss_coeff > 0: DSAIndexerLossLoggingHelper.save_loss_to_tracker( loss=indexer_loss, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 5e8e1d88c0f..09e40725bf7 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -11,7 +11,6 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from tests.unit_tests.test_utilities import Utils import megatron.core.parallel_state as parallel_state from megatron.core.extensions.transformer_engine import HAVE_TE @@ -22,6 +21,7 @@ DSAIndexerLossAutoScaler, DSAIndexerLossLoggingHelper, ) +from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_attention import ( _SEED, _build_attention, @@ -554,18 +554,30 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): return measured @pytest.mark.parametrize("sparse_loss", [True, False], ids=["sparse", "dense"]) - def test_cp2_unfused_indexer_loss_and_grad_matches_cp1(self, sparse_loss): + @pytest.mark.parametrize( + ("seg_lens", "padded_seg_lens"), + [((64,), None), ((5, 11, 17, 19), (8, 12, 20, 24))], + ids=["single", "ragged_padded"], + ) + def test_cp2_unfused_indexer_loss_and_grad_matches_cp1( + self, sparse_loss, seg_lens, padded_seg_lens, request + ): """CP2 unfused indexer loss and gradients must match the CP1 teacher.""" if self.cp_size != 2: pytest.skip("The exact unfused indexer regression is a focused CP2 gate") - seq_len = 64 - packed = _make_thd_packed_seq_params((seq_len,)) + def cleanup_indexer_loss_state(): + DSAIndexerLossAutoScaler.main_loss_backward_scale = None + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + _clear_cuda_test_state() + + request.addfinalizer(cleanup_indexer_loss_state) + + packed = _make_thd_packed_seq_params(seg_lens, padded_seg_lens) + seq_len = sum(padded_seg_lens or seg_lens) local_rows = seq_len // self.cp_size local_idx = torch.arange( - self.cp_rank * local_rows, - (self.cp_rank + 1) * local_rows, - device='cuda', + self.cp_rank * local_rows, (self.cp_rank + 1) * local_rows, device='cuda' ) common_config = dict( @@ -667,30 +679,44 @@ def test_cp2_unfused_indexer_loss_and_grad_matches_cp1(self, sparse_loss): if missing_cp_grads: errors.append(f"missing CP2 indexer gradients: {missing_cp_grads}") - nonzero_grad = False + nonzero_cp_grad = False + nonzero_ref_grad = False + grad_metrics = [] for name, cp_grad in reduced_cp_grads.items(): ref_param = ref_params.get(name) if ref_param is None or ref_param.grad is None: errors.append(f"missing CP1 indexer gradient: {name}") continue ref_grad = ref_param.grad.detach().float() - nonzero_grad = nonzero_grad or bool(torch.count_nonzero(ref_grad).item()) - close = torch.isclose(cp_grad, ref_grad, rtol=1e-3, atol=1e-4) - if not bool(close.all().item()): - max_abs = (cp_grad - ref_grad).abs().max().item() + nonzero_cp_grad = nonzero_cp_grad or bool(torch.count_nonzero(cp_grad).item()) + nonzero_ref_grad = nonzero_ref_grad or bool(torch.count_nonzero(ref_grad).item()) + max_abs = (cp_grad - ref_grad).abs().max().item() + grad_metrics.append( + f"{name}:max_abs={max_abs:.10e},cp_norm={cp_grad.norm().item():.10e}," + f"cp1_norm={ref_grad.norm().item():.10e}" + ) + if max_abs > 1e-4: errors.append(f"indexer gradient {name} max_abs={max_abs:.10e}") - if not nonzero_grad: + if not nonzero_cp_grad: + errors.append("all CP2 indexer gradients are zero") + if not nonzero_ref_grad: errors.append("all CP1 indexer gradients are zero") + if self.cp_rank == 0: + print( + f"[CP2-INDEXER] sparse_loss={sparse_loss} seg_lens={seg_lens} " + f"padded_seg_lens={padded_seg_lens} cp2={cp_loss.item():.10f} " + f"cp1={ref_loss.item():.10f} loss_abs={loss_diff.item():.10e}; " + + "; ".join(grad_metrics), + flush=True, + ) + failure_flag = torch.tensor(bool(errors), dtype=torch.int32, device='cuda') dist.all_reduce(failure_flag, op=dist.ReduceOp.MAX, group=self.pg.cp) if failure_flag.item(): pytest.fail("; ".join(errors) if errors else "peer rank reported an alignment failure") del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out - DSAIndexerLossAutoScaler.main_loss_backward_scale = None - DSAIndexerLossLoggingHelper.clean_loss_in_tracker() - _clear_cuda_test_state() @pytest.mark.parametrize( ("layer_number", "sparse_loss"), From 243b82680cc8ec10e54069037730d34cd83e3a15 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Mon, 3 Aug 2026 17:27:29 +0800 Subject: [PATCH 6/8] Remove unused CSA output shape Signed-off-by: JavaZeroo <2487163254@qq.com> --- megatron/core/transformer/experimental_attention_variant/csa.py | 1 - 1 file changed, 1 deletion(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 6d98333b11f..03b27b10067 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -746,7 +746,6 @@ def _compute_unfused_csa_non_compressed_lse( global_indices = torch.where( window_indices_i64 >= 0, window_indices_i64 + batch_offsets, window_indices_i64 ).reshape(batch_size * seqlen_q, -1) - output_shape = (batch_size, num_heads, seqlen_q) elif query.ndim == 3: total_q, num_heads, head_dim = query.shape if kv_full.ndim != 2 or kv_full.shape[1] != head_dim: From 29b2cd7b9aef70642efa76de35ecefb0e0986669 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 4 Aug 2026 16:07:34 +0800 Subject: [PATCH 7/8] Fix CSA LSE return lint error Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../core/transformer/experimental_attention_variant/csa.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 03b27b10067..3cf43968c3f 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -762,7 +762,6 @@ def _compute_unfused_csa_non_compressed_lse( q_flat = query.detach() kv_flat = kv_full.detach() global_indices = window_indices.to(dtype=torch.int64) - output_shape = (1, num_heads, total_q) else: raise ValueError( "non-compressed LSE query must be SBHD [sq, b, h, d] or " @@ -798,7 +797,7 @@ def _compute_unfused_csa_non_compressed_lse( if query.ndim == 4: return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() - return lse_flat.transpose(0, 1).unsqueeze(0).reshape(output_shape).contiguous() + return lse_flat.transpose(0, 1).unsqueeze(0).contiguous() def _unfused_indexer_sparse_attn_from_topk( From 241d0c7b4fe31de480bd9a10c708c23500e0f477 Mon Sep 17 00:00:00 2001 From: JavaZeroo <2487163254@qq.com> Date: Tue, 4 Aug 2026 16:47:38 +0800 Subject: [PATCH 8/8] Preserve CSA THD output layout Signed-off-by: JavaZeroo <2487163254@qq.com> --- .../core/transformer/experimental_attention_variant/csa.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 3cf43968c3f..46aa0c4312b 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -797,7 +797,9 @@ def _compute_unfused_csa_non_compressed_lse( if query.ndim == 4: return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() - return lse_flat.transpose(0, 1).unsqueeze(0).contiguous() + return ( + lse_flat.transpose(0, 1).unsqueeze(0).reshape(1, num_heads, lse_flat.shape[0]).contiguous() + ) def _unfused_indexer_sparse_attn_from_topk(