From bb891cb2ba853d2283ecc25352d7ec95e4f1abd4 Mon Sep 17 00:00:00 2001 From: Deyu Fu Date: Tue, 21 Jul 2026 23:31:41 +0800 Subject: [PATCH] feat(attention): add unfused compressed sparse attention Add the unfused compressed sparse attention algorithm, inert configuration fields, and focused algorithm coverage without enabling a model variant. Reconstructed from #4458 and its corrections as part of the frozen feature set in #5795. Signed-off-by: Deyu Fu --- .../experimental_attention_variant/csa.py | 780 ++++++++++++ .../core/transformer/transformer_config.py | 16 + .../models/test_hybrid_moe_model.py | 4 + .../test_attention_variant_csa.py | 1096 +++++++++++++++++ 4 files changed, 1896 insertions(+) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa.py create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py new file mode 100644 index 00000000000..0bc6c9d114a --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,780 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import copy +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +from megatron.core.models.common.embeddings import RotaryEmbedding, apply_rotary_pos_emb +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import nvtx_range_pop, nvtx_range_push + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=8) +def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: + """Compute sliding-window indices for a single sequence (cached). + + Returns: + indices: [seqlen, window_size] int tensor, -1 for invalid positions. + """ + base = torch.arange(seqlen, device=device_str).unsqueeze(1) + offsets = torch.arange(window_size, device=device_str) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix + + +def get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + """Sliding-window indices [batch, seqlen, window_size].""" + matrix = _get_window_topk_idxs_cached(window_size, seqlen, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +@lru_cache(maxsize=8) +def _get_compress_topk_idxs_cached( + ratio: int, seqlen: int, offset: int, device_str: str +) -> torch.Tensor: + """Compute all-compressed-positions indices for a single sequence (cached). + + Returns: + indices: [seqlen, seqlen // ratio] int tensor, -1 for future positions. + """ + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device_str).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix + + +def get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + """All-compressed-position indices [batch, seqlen, seqlen // ratio].""" + matrix = _get_compress_topk_idxs_cached(ratio, seqlen, offset, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_rope( + x: torch.Tensor, + nope_dim: int, + pos_dim: int, + rotary_pos_emb_module: RotaryEmbedding, + config: TransformerConfig, + rotary_seq_len: int, + ratio: int = 1, + cp_group: torch.distributed.ProcessGroup = None, +) -> torch.Tensor: + """Apply RoPE to the last ``qk_pos_emb_head_dim`` dims, leaving the rest unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs. When the input is 3-D a temporary head dimension is inserted for + ``apply_rotary_pos_emb`` and removed before returning. + """ + if ratio == 1: + total_seq_len = rotary_seq_len + else: + total_seq_len = rotary_seq_len * ratio + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0 + # regardless of which rotary class is in use. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if config.apply_rope_fusion: + # ``mscale=1.0`` keeps the cached cos/sin free of yarn's + # concentration factor so the fused kernel sees the same + # rotation as the unfused split-rotate path (DSv4 "pure + # rotation" contract). + rotary_pos_cos, rotary_pos_sin = rotary_pos_emb_module.get_cached_cos_sin( + total_seq_len, dtype=x.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + # Compressed-attention callers instantiate ``YarnRotaryEmbedding`` + # whenever ``compress_ratio > 1`` (regardless of ``config.rope_type``); + # its ``forward`` returns ``(emb, mscale)``. Base ``RotaryEmbedding`` + # returns a single tensor. Unpack either form uniformly; the + # caller-side ``mscale=1.0`` keeps the yarn concentration factor + # out of the rotation. + result = rotary_pos_emb_module(total_seq_len, packed_seq=False) + if isinstance(result, tuple): + rotary_pos_emb = result[0] + else: + rotary_pos_emb = result + if rotary_pos_emb is not None and ratio > 1: + rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_cos is not None and ratio > 1: + rotary_pos_cos = rotary_pos_cos[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_sin is not None and ratio > 1: + rotary_pos_sin = rotary_pos_sin[:total_seq_len:ratio][:rotary_seq_len] + + squeeze_head = x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + if config.apply_rope_fusion: + out = fused_mla_rope_inplace( + x, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + cp_group.rank(), + cp_group.size(), + remove_interleaving=True, + ) + else: + x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=config, + cu_seqlens=None, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + out = torch.cat([x_nope, x_pe], dim=-1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Differentiable sparse attention with MQA and attention sink. + + Args: + query: [sq, b, np, hn] multi-head query. + kv_full: [n_kv, b, hn] single-head KV (original + compressed). + attn_sink: [np] per-head learnable bias. + topk_indices: [b, sq, topk] indices into kv_full (int32, -1 = invalid). + softmax_scale: float + + Returns: + output: [sq, b, np * hn] + """ + sq, b, np_, hn = query.size() + + # --- Gather KV at topk positions --- + # kv_full: [n_kv, b, hn] -> [b, n_kv, hn] + kv_t = kv_full.permute(1, 0, 2) + + safe_indices = topk_indices.clamp(min=0).long() # [b, sq, topk] + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) # [b, sq, topk, hn] + # [b, n_kv, hn] -> [b, 1, n_kv, hn] -> gather -> [b, sq, topk, hn] + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp + ) + + # --- Attention scores --- + # query: [sq, b, np, hn] -> [b, np, sq, hn] + q = query.permute(1, 2, 0, 3).float() + kv_g = kv_gathered.float() # [b, sq, topk, hn] + + # [b, np, sq, topk] + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_g) * softmax_scale + + # Mask invalid + invalid_mask = (topk_indices < 0).unsqueeze(1) # [b, 1, sq, topk] + scores = scores.masked_fill(invalid_mask, float("-inf")) + + # --- Softmax with attention sink --- + sink = attn_sink.view(1, np_, 1, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values # [b, np, sq, 1] + scores_max = torch.max(scores_max, sink) + + exp_scores = torch.exp(scores - scores_max) # [b, np, sq, topk] + exp_sink = torch.exp(sink - scores_max) # [1, np, 1, 1] + + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + attn_weights = exp_scores / sum_exp # [b, np, sq, topk] + + # --- Weighted sum --- + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_g) + output = output.to(query.dtype) + + # [b, np, sq, hn] -> [sq, b, np, hn] -> [sq, b, np * hn] + output = output.permute(2, 0, 1, 3).contiguous() + output = output.reshape(sq, b, np_ * hn) + return output + + +# --------------------------------------------------------------------------- +# Compressor +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for CSA and HCA Compressor.""" + + linear_wkv: Union[ModuleSpec, type] = None + linear_wgate: Union[ModuleSpec, type] = None + norm: Union[ModuleSpec, type] = None + + +class Compressor(MegatronModule): + """Gated pooling compressor for CSA and HCA sparse attention. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens using learned gated weights. + + For ``compress_ratio == 4``, overlapping compression is used (``coff = 2``). + For ``compress_ratio == 128``, non-overlapping compression is used (``coff = 1``). + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressorSubmodules, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + + self.rotary_pos_emb = rotary_pos_emb + + proj_out_dim = self.coff * head_dim + + self.linear_wkv = build_module( + submodules.linear_wkv, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + self.linear_wgate = build_module( + submodules.linear_wgate, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # keep to high precision + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = nn.Parameter(_ape) + + norm_config = copy.copy(config) + norm_config.normalization = "RMSNorm" + self.norm = build_module( + submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon + ) + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + """Apply overlapping window transform for 4x compression. + + Input shape: [n_groups, ratio, b, coff * head_dim] + Output shape: [n_groups, 2 * ratio, b, head_dim] + """ + n_groups, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n_groups, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + new_tensor[1:, :ratio] = tensor[:-1, :, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states into shorter KV sequence. + + Args: + x: [sq, b, hidden_size] + + Returns: + compressed_kv [sq // ratio, b, head_dim] or None if too short. + """ + nvtx_range_push("compressor") + + sq, b, _ = x.size() + ratio = self.compress_ratio + + if sq < ratio: + nvtx_range_pop("compressor") + return None + + kv, _ = self.linear_wkv(x) # [sq, b, coff * head_dim] + score, _ = self.linear_wgate(x) # [sq, b, coff * head_dim] + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + + n_compressed = cutoff // ratio + + # Reshape: [n_compressed, ratio, b, coff * head_dim] + kv = kv.view(n_compressed, ratio, b, -1) + score = score.view(n_compressed, ratio, b, -1) + + # APE: [ratio, coff * head_dim] -> [1, ratio, 1, coff * head_dim] + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + kv = (kv * torch.softmax(score, dim=1)).sum(dim=1) # [n_compressed, b, head_dim] + + kv = self.norm(kv.to(x.dtype)) + + kv = _apply_rope( + kv, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + n_compressed, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + if self.rotate: + kv = rotate_activation(kv) + + nvtx_range_pop("compressor") + return kv # [n_compressed, b, head_dim] + + +# --------------------------------------------------------------------------- +# CSAIndexer +# --------------------------------------------------------------------------- + + +@dataclass +class CSAIndexerSubmodules: + """Submodule specs for CSAIndexer.""" + + linear_wq_b: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + compressor: Union[ModuleSpec, type] = None + + +class CSAIndexer(MegatronModule): + """Learned top-k retrieval over compressed positions for CSA sparse attention. + + Computes index scores to select the most relevant compressed KV positions for each + query. Reuses the scoring logic from ``DSAIndexer`` (einsum -> relu -> weight -> sum + -> topk) and ``rotate_activation`` (Hadamard transform) from ``dsa.py``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CSAIndexerSubmodules, + compress_ratio: int, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.hidden_size = config.hidden_size + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.q_lora_rank = ( + config.q_lora_rank if config.q_lora_rank is not None else config.hidden_size + ) + + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim**-0.5 + + self.rotary_pos_emb = rotary_pos_emb + + # Q projection + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # Weights projection + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # Own compressor (smaller head_dim, with Hadamard rotation) + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=compress_ratio, + head_dim=self.index_head_dim, + rotate=True, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute Q, compressed K, and weights before top-k selection.""" + nvtx_range_push("indexer_before_topk") + + sq, bsz, _ = x.size() + + # Q path + q, _ = self.linear_wq_b(qr) # [sq, b, n_heads * head_dim] + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + sq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + q = rotate_activation(q) + + # K path: own compressor + k = self.compressor(x) # [sq//ratio, b, index_head_dim] + + weights, _ = self.linear_weights_proj(x) # [sq, b, n_heads] + weights = weights * (self.index_n_heads**-0.5) + + nvtx_range_pop("indexer_before_topk") + return q, k, weights + + def forward( + self, x: torch.Tensor, qr: torch.Tensor, mask: Optional[torch.Tensor] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (index_scores, topk_indices).""" + nvtx_range_push("indexer") + q, k, weights = self.forward_before_topk(x, qr) + nvtx_range_push("indexer_qk_topk") + effective_topk = min(self.index_topk, k.size(0)) + index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, effective_topk, mask) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + +# --------------------------------------------------------------------------- +# CompressedSparseAttention (core attention) +# --------------------------------------------------------------------------- + + +@dataclass +class CompressedSparseAttentionSubmodules: + """Submodule specs for CompressedSparseAttention.""" + + compressor: Union[ModuleSpec, type] = None + indexer: Union[ModuleSpec, type] = None + + +class CompressedSparseAttention(MegatronModule): + """Sparse core attention for CompressedSparseAttention. + + Combines sliding window attention with compressed KV attention. The spec always + provides compressor and indexer submodule specs; this ``__init__`` inspects + ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: + + * ``ratio == 0``: window-only (compressor and indexer NOT built) + * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) + * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressedSparseAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + rotary_pos_emb: nn.Module = None, + compress_ratio: int = 0, + is_mtp_layer: bool = False, + ): + super().__init__(config=config) + + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + self.layer_number = layer_number + self.config.num_layers if is_mtp_layer else layer_number + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + self.n_local_heads = config.num_attention_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim**-0.5 + self.softmax_scale = softmax_scale + + self.force_unfused_dsa = getattr(config, 'force_unfused_dsa', True) + + # Learnable attention sink per head + self.attn_sink = nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) + + # Conditionally build Compressor (ratio > 1) + if self.compress_ratio > 1 and submodules.compressor is not None: + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=self.compress_ratio, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + else: + self.compressor = None + + # Conditionally build Indexer (ratio == 4) + if ( + self.compress_ratio == 4 + and not config.csa_dense_mode + and submodules.indexer is not None + ): + self.indexer = build_module( + submodules.indexer, + config=config, + compress_ratio=self.compress_ratio, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + else: + self.indexer = None + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor = None, + qr: torch.Tensor = None, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params=None, + ) -> torch.Tensor: + """Forward pass for CompressedSparseAttention. + + Args: + query: [sq, b, np, v_head_dim] + key: [sq, b, 1, v_head_dim] (single-head MQA; head dim squeezed internally) + value: unused (key == value in MQA) + attention_mask: attention mask (may be None for causal). + x: [sq, b, hidden_size] original hidden states. + qr: [sq, b, q_lora_rank] compressed query representation. + + Returns: + output: [sq, b, np * v_head_dim] + """ + nvtx_range_push("compressed_sparse_attn") + assert ( + packed_seq_params is None + ), "Packed sequence not supported for CompressedSparseAttention" + + sq, b, np, hn = query.size() + + # --- Step 1: Prepare single-head KV (squeeze singleton head dim) --- + kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] + + # --- Step 2: Compression --- + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) # [n_compressed, b, v_head_dim] + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + n_compressed = 0 + else: + kv_full = kv + n_compressed = 0 + + offset = sq # compressed indices start after original positions + + # --- Step 3: Window indices --- + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + # --- Step 4: Compressed indices --- + indexer_loss = None + + if self.force_unfused_dsa: + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where( + causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0 + ) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det + ) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + # compressed_kv is [n, b, hn]; expand to [n, b, np, hn] for loss + 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 * self.indexer.softmax_scale + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.indexer.pg_collection, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + ) + else: + _, topk_indices_compressed = self.indexer(x_det, qr_det, mask=causal_mask) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = (topk_indices_compressed >= 0) & ( + topk_indices_compressed < n_valid_per_pos + ) + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + # --- Step 5: Sparse attention --- + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + + else: + raise ValueError("Fused path is not supported for CompressedSparseAttention") + + # --- Step 6: Attach indexer loss --- + if indexer_loss is not None and self.training and torch.is_grad_enabled(): + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 761504e614e..ee5870ba3ee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -333,6 +333,22 @@ class TransformerConfig(ModelParallelConfig): dsa_indexer_k_norm_fp32: bool = False """Whether DSA indexer key LayerNorm should run on fp32 inputs.""" + #################### + # Compressed sparse attention + #################### + csa_window_size: int = 128 + """Sliding window size for compressed sparse attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in compressed sparse attention.""" + + csa_dense_mode: bool = False + """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be + disabled.""" + #################### # linear attention #################### diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index f7dc78ce9a2..681705a2d88 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -68,6 +68,10 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "csa_compress_ratios": None, + "csa_compress_rotary_base": 40000.0, + "csa_dense_mode": False, + "csa_window_size": 128, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_modules": [], 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 new file mode 100644 index 00000000000..19e96889293 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,1096 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _apply_rope, + get_compress_topk_idxs, + get_window_topk_idxs, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + + +def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + """Mock implementation of hadamard_transform for testing without the library installed.""" + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in both dsa and csa modules if not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# =========================================================================== +# Helper function tests +# =========================================================================== + + +class TestGetWindowTopkIdxs: + """Test get_window_topk_idxs helper.""" + + def test_basic_shape(self): + batch_size, seqlen, window_size = 2, 16, 4 + idxs = get_window_topk_idxs(window_size, batch_size, seqlen, torch.device("cpu")) + assert idxs.shape == (batch_size, seqlen, window_size) + + def test_causal_no_future(self): + """Indices should never exceed the query position.""" + seqlen, window_size = 32, 8 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + for i in range(seqlen): + valid = idxs[0, i][idxs[0, i] >= 0] + assert torch.all(valid <= i), f"Position {i} has future indices" + + def test_invalid_marked_minus_one(self): + """Early positions that cannot fill the window should use -1.""" + seqlen, window_size = 8, 4 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs[0, 0, 0] == -1 or idxs[0, 0, 0] == 0 + for pos in range(window_size, seqlen): + assert torch.all(idxs[0, pos] >= 0), f"Position {pos} has invalid -1" + + def test_window_larger_than_seqlen(self): + """Window larger than sequence length should still work.""" + seqlen, window_size = 4, 16 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs.shape == (1, seqlen, window_size) + + +class TestGetCompressTopkIdxs: + """Test get_compress_topk_idxs helper.""" + + def test_basic_shape(self): + ratio, batch_size, seqlen, offset = 4, 2, 32, 32 + idxs = get_compress_topk_idxs(ratio, batch_size, seqlen, offset, torch.device("cpu")) + n_compressed = seqlen // ratio + assert idxs.shape == (batch_size, seqlen, n_compressed) + + def test_offset_applied(self): + """Valid indices should be >= offset.""" + ratio, seqlen, offset = 4, 32, 100 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + valid = idxs[idxs >= 0] + if valid.numel() > 0: + assert torch.all(valid >= offset), "Valid indices should be offset" + + def test_causal_no_future(self): + """Compressed indices should respect causality.""" + ratio, seqlen, offset = 4, 32, 32 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + for i in range(seqlen): + n_valid = (i + 1) // ratio + valid = idxs[0, i][idxs[0, i] >= 0] + assert valid.numel() <= n_valid, f"Position {i} has too many valid compressed indices" + + def test_ratio_128(self): + """Test with large compression ratio.""" + ratio, seqlen, offset = 128, 256, 256 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + assert idxs.shape == (1, seqlen, seqlen // ratio) + + +# =========================================================================== +# unfused_compressed_sparse_attn tests +# =========================================================================== + + +class TestUnfusedCompressedSparseAttn: + """Test the unfused compressed sparse attention kernel.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_output_shape(self): + """Test output shape of unfused compressed sparse attention.""" + sq, b, np_, hn = 16, 2, 4, 64 + n_kv = sq + sq // 4 + topk = 8 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + + assert output.shape == (sq, b, np_ * hn) + assert output.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_indices_masked(self): + """Test that -1 indices are properly masked.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((b, sq, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, :, 0] = 0 + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + assert not torch.isnan(output).any(), "Output should not contain NaN" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_gradient_flow(self): + """Test that gradients flow through sparse attention.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + kv_full = torch.randn(n_kv, b, hn, dtype=torch.float32).cuda().requires_grad_(True) + attn_sink = torch.nn.Parameter(torch.zeros(np_, dtype=torch.float32).cuda()) + + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert kv_full.grad is not None + assert attn_sink.grad is not None + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def _make_mla_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + csa_compress_ratios=None, + csa_window_size=8, + csa_dense_mode=False, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + dsa_indexer_use_sparse_loss=False, +): + """Helper to create MLATransformerConfig for CSA tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0] * num_layers + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + csa_dense_mode=csa_dense_mode, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + + +def _make_compressor_submodules(): + """Create Compressor submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ) + + +def _make_csa_indexer_submodules(): + """Create CSAIndexer submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_weights_proj=ModuleSpec(module=TELinear), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + ) + + +def _make_csa_submodules(): + """Create CompressedSparseAttention submodules spec.""" + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressedSparseAttentionSubmodules( + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + indexer=ModuleSpec(module=CSAIndexer, submodules=_make_csa_indexer_submodules()), + ) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressor: + """Test Compressor module.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_output_shape(self, compress_ratio): + """Test that compressor produces correct output shape.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + + expected_len = seq_len // compress_ratio + assert output is not None + assert output.shape == (expected_len, batch_size, head_dim) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_too_short_input(self, compress_ratio): + """Test that compressor returns None when input is shorter than compress_ratio.""" + short_len = compress_ratio - 1 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + assert output is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_gradient_flow(self, compress_ratio): + """Test that gradients flow through the compressor.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + output = compressor(x) + loss = output.sum() + loss.backward() + + assert x.grad is not None + for name, param in compressor.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + +# =========================================================================== +# CSAIndexer tests +# =========================================================================== + + +@pytest.mark.parametrize("seqlen", [32, 128]) +class TestCSAIndexer: + """Test CSAIndexer module basic functionality.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_csa_indexer_constructor(self, seqlen): + """Test CSAIndexer initialization.""" + assert isinstance(self.indexer, CSAIndexer) + assert self.indexer.compress_ratio == self.compress_ratio + assert self.indexer.index_n_heads == self.config.dsa_indexer_n_heads + assert self.indexer.index_head_dim == self.config.dsa_indexer_head_dim + assert self.indexer.index_topk == self.config.dsa_indexer_topk + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward(self, seqlen): + """Test CSAIndexer forward pass.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + index_scores, topk_indices = self.indexer(x, qr) + n_compressed = seqlen // self.compress_ratio + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward_before_topk(self, seqlen): + """Test CSAIndexer forward_before_topk.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + q, k, weights = self.indexer.forward_before_topk(x, qr) + + assert q.shape == ( + seqlen, + batch_size, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim, + ) + n_compressed = seqlen // self.compress_ratio + assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) + assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_with_mask(self, seqlen): + """Test CSAIndexer with causal mask.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + n_compressed = seqlen // self.compress_ratio + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) + positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + + index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) + + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + + +# =========================================================================== +# CompressedSparseAttention tests +# =========================================================================== + + +class TestCompressedSparseAttentionRatio1: + """Test CompressedSparseAttention with compress_ratio=1 (window-only).""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config(csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.csa = CompressedSparseAttention( + config=cls.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=cls.pg_collection, + rotary_pos_emb=rotary_pos_emb, + compress_ratio=0, + ) + + yield + Utils.destroy_model_parallel() + + def test_ratio1_no_compressor(self): + """With ratio=1, compressor and indexer should not be built.""" + assert self.csa.compressor is None + assert self.csa.indexer is None + + def test_mtp_layer_number_is_offset(self): + """MTP attention layers are numbered after all decoder layers.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + compress_ratio=0, + is_mtp_layer=True, + ) + + assert csa.layer_number == self.config.num_layers + 1 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_forward(self): + """Test forward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_backward(self): + """Test backward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.train() + self.csa.cuda() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressedSparseAttentionCompressed: + """Test CompressedSparseAttention with compress_ratio > 1.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return a layer_number (1-indexed) whose compress_ratio matches.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_constructor(self, compress_ratio): + """Test that compressor/indexer are conditionally built.""" + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + assert csa.compressor is not None + if compress_ratio == 4: + assert csa.indexer is not None + elif compress_ratio == 128: + assert csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward(self, compress_ratio): + """Test forward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward(self, compress_ratio): + """Test backward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.train() + + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + for name, param in csa.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eval_mode(self, compress_ratio): + """Test forward pass in eval mode.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + csa.eval() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + with torch.no_grad(): + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +# =========================================================================== +# _apply_rope tests +# =========================================================================== + + +class TestApplyRope: + """Test ``_apply_rope`` — the layout-aware RoPE wrapper used by + Compressor / CSAIndexer / hybrid-attention callers. + + Behaviours covered: + + * 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs both work (3-D gets a temporary head-dim unsqueeze). + * Only the trailing ``pos_dim`` components are rotated; the leading + ``nope_dim`` slice is bit-exact unchanged. + * Both ``RotaryEmbedding`` (returns ``Tensor``) and + ``YarnRotaryEmbedding`` (returns ``(emb, mscale)`` tuple) — DSv4 + hybrid silently swaps the class based on ``compress_ratio``. + * Both unfused and fused (``config.apply_rope_fusion=True``) paths + produce the same output (within bf16 precision). + * For ``ratio > 1`` the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(0) + model_parallel_cuda_manual_seed(0) + cls = request.cls + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + # head_dim 32 = nope 24 + pos 8 + cls.config = _make_mla_config(v_head_dim=32, qk_pos_emb_head_dim=8) + yield + Utils.destroy_model_parallel() + + def _make_rotary(self, kind: str): + from megatron.core.models.common.embeddings import RotaryEmbedding, YarnRotaryEmbedding + + pos_dim = self.config.qk_pos_emb_head_dim + if kind == 'rope': + return RotaryEmbedding( + pos_dim, rotary_percent=1.0, rotary_base=10000, cp_group=self.pg_collection.cp + ) + if kind == 'yarn': + return YarnRotaryEmbedding( + pos_dim, + rotary_base=40000, + scaling_factor=40, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1.0, + mscale_all_dim=0.0, + cp_group=self.pg_collection.cp, + ) + raise ValueError(kind) + + def _config_with(self, *, apply_rope_fusion: bool): + # Reuse the class-level config; only flip the fusion flag. + cfg = self.config + cfg.apply_rope_fusion = apply_rope_fusion + return cfg + + _ROTARY_FUSION_COMBOS = [ + pytest.param('rope', False, id='rope-unfused'), + pytest.param('rope', True, id='rope-fused'), + pytest.param('yarn', False, id='yarn-unfused'), + pytest.param('yarn', True, id='yarn-fused'), + ] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize(("rotary_kind", "apply_rope_fusion"), _ROTARY_FUSION_COMBOS) + @pytest.mark.parametrize("input_ndim", [3, 4], ids=['3d', '4d']) + @pytest.mark.parametrize("ratio", [1, 4], ids=['ratio_1', 'ratio_4']) + def test_apply_rope(self, rotary_kind, apply_rope_fusion, input_ndim, ratio): + """Output shape == input shape; no NaN; nope-dim slice is + bit-exact unchanged. Sweeps the valid combinations of rotary + class × apply_rope_fusion × input rank × ratio. Yarn's + tuple-return is covered by the ``'yarn-*'`` combos. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads = 8, 2, 4 + cfg = self._config_with(apply_rope_fusion=apply_rope_fusion) + + shape = (seq, batch, head_dim) if input_ndim == 3 else (seq, batch, heads, head_dim) + x = torch.randn(*shape, dtype=torch.bfloat16, device='cuda') + # ``fused_mla_rope_inplace`` mutates the input — give it a copy so + # the nope-dim equality check below still has the original. + out = _apply_rope( + x.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + assert out.shape == x.shape + assert out.dtype == x.dtype + assert not torch.isnan(out).any() + # The leading nope_dim slice is the identity portion of RoPE. + assert torch.equal( + out[..., :nope], x[..., :nope] + ), "RoPE must not touch the first nope_dim components" + # Trailing pos_dim should rotate at non-zero positions. + pe_changed = (out[..., nope:] != x[..., nope:]).any(dim=-1).flatten() + assert pe_changed[ + 1: + ].any(), "RoPE should rotate the trailing pos_dim components for seq > 0" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_3d_input_matches_4d_with_single_head(self, rotary_kind): + """For a single-head input, the 3-D ``(s, b, d)`` and 4-D + ``(s, b, 1, d)`` invocations must produce numerically identical + output (3-D path just inserts a temporary head dim). + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch = 8, 2 + cfg = self._config_with(apply_rope_fusion=False) + + x_3d = torch.randn(seq, batch, head_dim, dtype=torch.bfloat16, device='cuda') + x_4d = x_3d.unsqueeze(-2) + + out_3d = _apply_rope( + x_3d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_4d = _apply_rope( + x_4d, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + + assert out_3d.shape == x_3d.shape + assert out_4d.shape == x_4d.shape + assert torch.equal(out_3d, out_4d.squeeze(-2)) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("rotary_kind", ['rope', 'yarn']) + def test_ratio_strides_rotary_table(self, rotary_kind): + """For ``ratio > 1``, the rotary table is built at + ``rotary_seq_len * ratio`` and strided by ``ratio``. The result + with ``ratio=k`` must equal an ``apply_rope`` call on the same + positions of a length-``rotary_seq_len * k`` table. + """ + rotary = self._make_rotary(rotary_kind).cuda() + nope = self.config.v_head_dim - self.config.qk_pos_emb_head_dim + pos = self.config.qk_pos_emb_head_dim + head_dim = nope + pos + seq, batch, heads, ratio = 4, 1, 2, 4 + cfg = self._config_with(apply_rope_fusion=False) + + x_comp = torch.randn(seq, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda') + out_comp = _apply_rope( + x_comp.clone(), + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + x_full = torch.zeros( + seq * ratio, batch, heads, head_dim, dtype=torch.bfloat16, device='cuda' + ) + x_full[::ratio][:seq] = x_comp + out_full = _apply_rope( + x_full, + nope, + pos, + rotary, + cfg, + rotary_seq_len=seq * ratio, + ratio=1, + cp_group=self.pg_collection.cp, + ) + out_ref = out_full[::ratio][:seq] + + assert torch.allclose(out_comp, out_ref, rtol=1e-3, atol=1e-3), ( + f"ratio={ratio} stride mismatch: " + f"max abs diff = {(out_comp - out_ref).abs().max().item():.3e}" + ) + + +# =========================================================================== +# csa_dense_mode tests +# =========================================================================== + + +class TestCompressedSparseAttentionDenseMode: + """Test that csa_dense_mode=True disables the indexer for ratio=4 layers.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], csa_window_size=8, csa_dense_mode=True + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_disables_indexer_for_ratio4(self): + """With csa_dense_mode=True, ratio=4 layers should NOT build an indexer.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + assert csa.compress_ratio == 4 + assert csa.compressor is not None, "Compressor should still be built" + assert csa.indexer is None, "Indexer should be disabled in dense mode" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_forward_ratio4(self): + """Forward pass should work for ratio=4 in dense mode (uses all compressed positions).""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any()