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..19bf0e78d9a --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,984 @@ +# 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.fp8_utils import get_fp8_disabled_context +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, mark_keep_in_fp32 +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 + +#: Bit-exact determinism status for the eager CSA operations introduced here. +#: The operations use CUDA reductions and indexed accumulation, but bit-exact +#: repeatability has not been certified, so the conservative status is unknown. +CSA_OPERATION_DETERMINISM: dict[str, str] = { + "unfused_sparse_attention": "unknown", + "non_compressed_lse": "unknown", + "compressor_pooling": "unknown", +} + +# --------------------------------------------------------------------------- +# 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) + + +@lru_cache(maxsize=8) +def _get_compress_causal_mask_cached( + ratio: int, seqlen: int, n_compressed: int, device_str: str +) -> torch.Tensor: + """Return the additive causal mask for compressed positions (cached).""" + compressed_positions = torch.arange(n_compressed, device=device_str).unsqueeze(0) + valid_counts = torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + return torch.where(compressed_positions >= valid_counts, float("-inf"), 0.0) + + +@lru_cache(maxsize=8) +def _get_compress_valid_counts_cached(ratio: int, seqlen: int, device_str: str) -> torch.Tensor: + """Return the number of causally valid compressed positions per query (cached).""" + return torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + + +# --------------------------------------------------------------------------- +# 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. + + Determinism: + Unknown. Bit-exact forward and backward repeatability has not been certified. + + 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() + if attn_sink.ndim != 1 or attn_sink.numel() != np_: + raise ValueError( + f"attn_sink must contain one value per query head ({np_}), " + f"got shape {tuple(attn_sink.shape)}." + ) + + # --- Gather KV at topk positions --- + # Flatten batch and KV position before gathering. Gathering from a logical + # [b, sq, n_kv, hn] expanded view makes gather backward allocate that entire + # dense shape before reducing the stride-0 query dimension. + n_kv = kv_full.size(0) + topk = topk_indices.size(-1) + kv_flat = kv_full.permute(1, 0, 2).reshape(b * n_kv, hn) + batch_offsets = (torch.arange(b, device=kv_full.device, dtype=torch.int64) * n_kv).view(b, 1, 1) + safe_indices = topk_indices.clamp(min=0).to(dtype=torch.int64) + batch_offsets + kv_gathered = kv_flat.index_select(0, safe_indices.reshape(-1)).view(b, sq, topk, hn) + + # --- 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 + + +@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, + softmax_scale: float, + chunk_size: int = 512, +) -> torch.Tensor: + """Return the detached sliding-window-plus-sink log mass for the CSA teacher. + + Determinism: + Unknown. Bit-exact CUDA reduction behavior has not been certified. + + Args: + query: Query tensor in ``[sq, batch, heads, head_dim]`` layout. + kv_full: Original (non-compressed) KV in ``[sk, batch, head_dim]`` layout. + attn_sink: Per-head sink logits in ``[heads]`` layout. + window_indices: Local per-batch window indices in ``[batch, sq, window]`` layout. + softmax_scale: Scale applied to query-key logits. + chunk_size: Maximum number of flattened query rows processed at once. + + Returns: + Detached FP32 log-sum-exp values in ``[batch, heads, sq]`` layout. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if query.ndim != 4: + raise ValueError(f"query must have shape [sq, batch, heads, dim], got {query.shape}") + if attn_sink.ndim != 1: + raise ValueError(f"attn_sink must be 1D, got shape {tuple(attn_sink.shape)}") + + 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( + "non-compressed KV must have 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( + "window_indices must have shape " + f"[{batch_size}, {seqlen_q}, window], got {tuple(window_indices.shape)}" + ) + if attn_sink.numel() != num_heads: + raise ValueError(f"attn_sink must contain {num_heads} values, got {attn_sink.numel()}") + if not (query.device == kv_full.device == attn_sink.device == window_indices.device): + raise ValueError("query, kv_full, attn_sink, and window_indices must share a device") + + n_kv = kv_full.shape[0] + 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) + + sink = attn_sink.detach().to(dtype=torch.float32).view(1, num_heads) + lse_chunks = [] + for start in range(0, q_flat.shape[0], chunk_size): + end = min(start + chunk_size, q_flat.shape[0]) + indices = global_indices[start:end] + gathered_kv = kv_flat.index_select(0, indices.clamp(min=0).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") + ) + lse_chunks.append(torch.logaddexp(torch.logsumexp(window_logits, dim=-1), sink)) + + 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) + return lse_flat.reshape(batch_size, seqlen_q, num_heads).permute(0, 2, 1).contiguous() + + +# --------------------------------------------------------------------------- +# 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 + + +def _pool_compressor_values( + kv: torch.Tensor, score: torch.Tensor, output_dtype: torch.dtype +) -> torch.Tensor: + """Pool compressor values with FP32 weights, products, and reduction.""" + weights = torch.softmax(score, dim=1, dtype=torch.float32) + return (kv.float() * weights).sum(dim=1).to(output_dtype) + + +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, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("Compressor requires an explicit ProcessGroupCollection") + 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 + + with get_fp8_disabled_context(config, is_init=True): + 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", + name=(name + ".linear_wkv") if name is not None else None, + ) + + 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", + name=(name + ".linear_wgate") if name is not None else None, + ) + + # keep to high precision (FP32 in the reference DeepSeek V4 checkpoint) + _ape = torch.empty( + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = mark_keep_in_fp32(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 backward_dw(self): + """Compute deferred weight gradients for the compressor projections.""" + self.linear_wkv.backward_dw() + self.linear_wgate.backward_dw() + + 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 _project(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Project compressor values and gates outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + kv, _ = self.linear_wkv(x) + score, _ = self.linear_wgate(x) + return kv, score + + def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states into shorter KV sequence. + + Determinism: + Unknown. Bit-exact CUDA pooling and gradient reductions have not been certified. + + 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, score = self._project(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 = _pool_compressor_values(kv, score, x.dtype) # [n_compressed, b, head_dim] + + kv = self.norm(kv) + + 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, + name: str | None = None, + ) -> None: + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError("CSAIndexer requires an explicit ProcessGroupCollection") + 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", + name=(name + ".linear_wq_b") if name is not None else None, + ) + + # The reference DeepSeek V4 checkpoint keeps this projection in BF16. + with get_fp8_disabled_context(config, is_init=True): + 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", + name=(name + ".linear_weights_proj") if name is not None else None, + ) + + # 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, + name=(name + ".compressor") if name is not None else None, + ) + + def backward_dw(self): + """Compute deferred weight gradients for the indexer projections.""" + self.linear_wq_b.backward_dw() + self.linear_weights_proj.backward_dw() + self.compressor.backward_dw() + + def _project_weights(self, x: torch.Tensor) -> torch.Tensor: + """Project indexer weights outside any enclosing FP8 context.""" + with get_fp8_disabled_context(self.config): + weights, _ = self.linear_weights_proj(x) + return weights + + 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._project_weights(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; which are built depends on the + ``compress_ratio`` passed by the caller: + + * ``ratio <= 1``: window-only (neither compressor nor indexer is built). + * ``ratio > 1``: window + compressed KV via ``Compressor``. + * ``ratio == 4`` and not ``config.csa_dense_mode``: additionally builds + ``CSAIndexer`` for learned top-k retrieval over compressed positions. Otherwise, + all causally valid compressed positions are attended. + """ + + 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, + name: str | None = None, + ): + """ + Args: + name (str | None): module instance name passed top-down from its parent module + """ + super().__init__(config=config) + + if pg_collection is None: + raise ValueError( + "CompressedSparseAttention requires an explicit ProcessGroupCollection" + ) + self.pg_collection = pg_collection + + tp_size = self.pg_collection.tp.size() + if tp_size != 1: + raise ValueError( + "CompressedSparseAttention supports only tensor-parallel size 1 in the " + f"native SBHD slice, got tp_size={tp_size}." + ) + + 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.num_attention_heads = config.num_attention_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim**-0.5 + self.softmax_scale = softmax_scale + + # Learnable attention sink per head, kept in reference-checkpoint FP32. + self.attn_sink = mark_keep_in_fp32( + nn.Parameter(torch.zeros(self.num_attention_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, + name=(name + ".compressor") if name is not None else None, + ) + 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, + name=(name + ".indexer") if name is not None else None, + ) + else: + self.indexer = None + + def backward_dw(self): + """Compute deferred gradients for the optional compressor and indexer projections.""" + if self.compressor is not None: + self.compressor.backward_dw() + if self.indexer is not None: + self.indexer.backward_dw() + + 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: Must be None; causal masking is applied internally. + 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] + """ + if attention_mask is not None: + raise ValueError( + "CompressedSparseAttention supports only an implicit causal mask in the " + "native SBHD slice; padding and document-boundary masks are not supported." + ) + if query.ndim != 4: + raise ValueError( + "CompressedSparseAttention query must have shape [seq, batch, heads, dim], " + f"got {tuple(query.shape)}." + ) + if query.size(2) != self.num_attention_heads: + raise ValueError( + "CompressedSparseAttention query head count must match the unsharded " + f"attention sink ({self.num_attention_heads}), got {query.size(2)}." + ) + 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.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 = ( + _get_compress_causal_mask_cached( + self.compress_ratio, sq, n_compressed, str(x.device) + ) + .unsqueeze(0) + .expand(b, -1, -1) + ) + + 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 = self.config.dsa_indexer_loss_coeff or 0.0 + 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, self.attn_sink, window_idxs, self.softmax_scale + ) + # The native reference intentionally recomputes sliding-window + # logits in the final attention call below. The teacher needs its + # detached denominator before indexer top-k is available; sharing + # it without materializing the full window gather requires a larger + # data-flow change. TODO(#6404): its fused training backend avoids + # this duplicate, but the native fallback should share the gathered + # window KV/logits instead of retaining this correctness-first helper. + 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, + self.config.dsa_indexer_use_sparse_loss, + 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( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_compressed = self.indexer(x_det, qr_det, mask=causal_mask) + + n_valid_per_pos = _get_compress_valid_counts_cached( + self.compress_ratio, sq, str(x.device) + ) + valid = (topk_indices_compressed >= 0) & (topk_indices_compressed < n_valid_per_pos) + compress_topk_idxs = torch.where(valid, topk_indices_compressed + offset, -1) + 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") + + # --- 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/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index bfead2a25c5..e00dec159ae 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -396,6 +396,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. @@ -421,6 +422,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). @@ -489,8 +494,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( @@ -504,7 +509,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, @@ -514,6 +519,80 @@ 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: + """Return selected-key teacher mass, optionally including omitted 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. When the absolute compressed + mass underflows FP32, all heads in a row receive the same log-domain shift; + the returned weights remain proportional and the caller L1-normalizes them. + """ + 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) + row_has_compressed_keys = expanded_valid_mask.any(dim=-1) + # Avoid the undefined ``-inf - -inf`` intermediate on fully masked rows. + # This is only a [batch, heads, seqlen] tensor, so it does not recreate the + # full-size temporary that the log-domain formulation is designed to avoid. + safe_compressed_lse = torch.where( + row_has_compressed_keys, compressed_lse, torch.zeros_like(compressed_lse) + ) + conditional_probabilities = torch.exp(masked_scores - safe_compressed_lse.unsqueeze(-1)) + del masked_scores + + full_lse = torch.logaddexp(non_compressed_lse.float(), compressed_lse) + log_compressed_mass = (compressed_lse - full_lse).masked_fill( + ~row_has_compressed_keys, float("-inf") + ) + + # The external window/sink mass can put every head's compressed mass below + # the FP32 normal range. A common per-row shift across heads + # preserves all relative teacher weights and cancels in the downstream L1 + # normalization. CSA currently requires TP1, so no cross-rank MAX is needed. + row_max = log_compressed_mass.amax(dim=1, keepdim=True) + needs_rescale = torch.isfinite(row_max) & (row_max < math.log(torch.finfo(torch.float32).tiny)) + common_shift = torch.where(needs_rescale, row_max, torch.zeros_like(row_max)) + compressed_mass = torch.exp(log_compressed_mass - common_shift) + return conditional_probabilities * compressed_mass.unsqueeze(-1) + + +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) + row_mass = target.sum(dim=-1, keepdim=True) + # External teacher mass can legitimately make the compressed mass smaller + # than INDEXER_LOSS_EPS (or even float32 tiny). Only an exactly zero row is + # degenerate; keep it zero rather than imposing a numerical floor. + safe_row_mass = torch.where(row_mass > 0, row_mass, torch.ones_like(row_mass)) + return target / safe_row_mass + + def _compute_index_scores( q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor, use_relu: bool = True ) -> torch.Tensor: @@ -627,6 +706,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( @@ -656,6 +736,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 @@ -680,6 +761,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") @@ -752,8 +834,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 @@ -776,7 +858,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 @@ -806,11 +890,12 @@ def bwd_fused_indexer_loss_naive( dtype=grad_kl_per_element.dtype ) - # For KL(target || softmax(logits)), the exact logit gradient is predict - target. - # Computing it through -target / (predict + eps) incorrectly suppresses gradients when - # valid predicted probabilities are smaller than eps. + # For KL(target || softmax(logits)), the exact logit gradient is + # predict * target.sum(-1) - target. Positive teacher rows are L1-normalized, + # while a fully masked zero-mass row must have zero gradient. + attention_target_mass = attention_scores_normalized.sum(dim=-1, keepdim=True) grad_index_scores_logits = ( - index_scores_softmax - attention_scores_normalized + index_scores_softmax * attention_target_mass - attention_scores_normalized ) * grad_kl_per_element del index_scores_softmax, attention_scores_normalized @@ -890,6 +975,7 @@ def bwd_fused_indexer_loss_naive( "query_valid_rows", "calculate_per_token_loss", "use_relu", + "non_compressed_lse", ) @@ -916,6 +1002,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. @@ -938,10 +1025,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 @@ -953,6 +1047,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 @@ -961,7 +1056,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, @@ -982,6 +1078,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 = { @@ -991,8 +1088,10 @@ 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) + 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/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 087c41284d1..875d0955e7a 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -350,6 +350,25 @@ 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.""" + + # TODO(#6402): consumed by DSv4 Hybrid attention orchestration, which selects the + # per-layer compression ratio and builds the compressed-KV rotary embedding. + # Neither field has a production reader in this primitive-only PR. + 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 b3ea98434c4..0f0a9d611bc 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..2f507e331df --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,1594 @@ +# 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 ( + CSA_OPERATION_DETERMINISM, + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + _apply_rope, + _compute_unfused_csa_non_compressed_lse, + _get_compress_causal_mask_cached, + _get_compress_valid_counts_cached, + _pool_compressor_values, + get_compress_topk_idxs, + get_window_topk_idxs, + 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 + +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 + + +class _DisabledContextTracker: + """Track whether a projection runs inside the FP8-disabled context.""" + + def __init__(self): + self.depth = 0 + self.entries = 0 + + def __call__(self, _config, is_init=False): + assert not is_init + return self + + def __enter__(self): + self.depth += 1 + self.entries += 1 + return self + + def __exit__(self, _exc_type, _exc_value, _traceback): + self.depth -= 1 + return False + + +@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 _SingleRankTP: + @staticmethod + def size(): + return 1 + + +class _SingleRankPG: + tp = _SingleRankTP() + + +class _TwoRankTP: + @staticmethod + def size(): + return 2 + + +class _TwoRankPG: + tp = _TwoRankTP() + + +def test_csa_rejects_explicit_attention_mask(): + """The native SBHD slice must not silently ignore padding or document boundaries.""" + with pytest.raises(ValueError, match="implicit causal mask"): + CompressedSparseAttention.forward( + None, query=None, key=None, value=None, attention_mask=torch.zeros(1, 1, 1, 1) + ) + + +def test_all_csa_operations_declare_determinism(): + """Every eager CSA operation declares a valid bit-exact determinism status.""" + valid_statuses = {"deterministic", "nondeterministic", "unknown"} + assert set(CSA_OPERATION_DETERMINISM) == { + "unfused_sparse_attention", + "non_compressed_lse", + "compressor_pooling", + } + assert set(CSA_OPERATION_DETERMINISM.values()) <= valid_statuses + + +def test_compressed_causal_metadata_is_cached_and_correct(): + ratio, seqlen, n_compressed = 4, 12, 3 + device_str = "cpu" + _get_compress_causal_mask_cached.cache_clear() + _get_compress_valid_counts_cached.cache_clear() + + mask = _get_compress_causal_mask_cached(ratio, seqlen, n_compressed, device_str) + valid_counts = _get_compress_valid_counts_cached(ratio, seqlen, device_str) + + assert mask is _get_compress_causal_mask_cached(ratio, seqlen, n_compressed, device_str) + assert valid_counts is _get_compress_valid_counts_cached(ratio, seqlen, device_str) + expected_counts = torch.arange(1, seqlen + 1).unsqueeze(1) // ratio + expected_mask = torch.where( + torch.arange(n_compressed).unsqueeze(0) >= expected_counts, float("-inf"), 0.0 + ) + torch.testing.assert_close(valid_counts, expected_counts) + torch.testing.assert_close(mask, expected_mask) + assert mask.unsqueeze(0).expand(2, -1, -1).stride(0) == 0 + + +def test_unfused_csa_non_compressed_lse_matches_window_and_sink_oracle(): + torch.manual_seed(17) + seqlen_q, batch_size, n_kv = 3, 2, 5 + num_heads, head_dim = 2, 4 + 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) + sink = torch.randn(num_heads, 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) + + actual = _compute_unfused_csa_non_compressed_lse( + query, kv_full, sink, window_indices, softmax_scale=1.0 + ) + + 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, +): + """Compute a small-loop CSA teacher oracle with the complete denominator.""" + 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 + ) + selected_position = selected.index(compressed_index) + head_mass = head_mass + torch.exp( + compressed_logits[selected_position] - 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_indexer_loss_uses_full_attention_denominator(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]]]) + non_compressed_lse = _compute_unfused_csa_non_compressed_lse( + query, window_kv, 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_reference = q.detach().clone().requires_grad_(True) + weights_reference = weights.detach().clone().requires_grad_(True) + k_reference = k.detach().clone().requires_grad_(True) + index_scores_reference, topk_reference = fused_qk_topk_naive( + q_reference, k_reference, weights_reference, index_topk + ) + loss_reference = compute_dsa_indexer_loss( + index_scores_reference, + topk_reference, + query.detach(), + key_for_loss.detach(), + 1.0, + loss_coeff, + sparse_loss, + _SingleRankPG(), + mask=compressed_mask, + non_compressed_lse=non_compressed_lse, + ) + loss_reference.backward() + + 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_reference.detach(), + topk_reference, + query.detach(), + compressed_kv.detach(), + window_kv.detach(), + window_indices, + sink.detach(), + sparse_loss=sparse_loss, + loss_coeff=loss_coeff, + ) + + torch.testing.assert_close(loss_actual, independent_loss) + torch.testing.assert_close(loss_actual, loss_reference) + torch.testing.assert_close(topk_actual, topk_reference) + torch.testing.assert_close(q.grad, q_reference.grad) + torch.testing.assert_close(weights.grad, weights_reference.grad) + torch.testing.assert_close(k.grad, k_reference.grad) + for teacher_tensor in (query, window_kv, compressed_kv, sink): + assert teacher_tensor.grad is None + + +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 +# =========================================================================== + + +def test_unfused_sparse_attention_rejects_sink_head_mismatch(): + with pytest.raises(ValueError, match="one value per query head"): + unfused_compressed_sparse_attn( + query=torch.zeros(2, 1, 2, 4), + kv_full=torch.zeros(2, 1, 4), + attn_sink=torch.zeros(1), + topk_indices=torch.zeros(1, 2, 1, dtype=torch.int32), + softmax_scale=1.0, + ) + + +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 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_batched_repeated_and_invalid_indices_match_loop_oracle(self): + """Batch-flattened gather preserves forward and backward sparse-attention semantics.""" + torch.manual_seed(41) + sq, b, np_, hn, n_kv = 3, 2, 2, 4, 5 + topk_indices = torch.tensor( + [[[0, 0, -1], [1, 3, 1], [4, -1, 2]], [[2, -1, 2], [4, 0, -1], [1, 1, 3]]], + dtype=torch.int32, + device="cuda", + ) + query = torch.randn(sq, b, np_, hn, device="cuda", requires_grad=True) + kv_full = torch.randn(n_kv, b, hn, device="cuda", requires_grad=True) + attn_sink = torch.randn(np_, device="cuda", requires_grad=True) + grad_output = torch.randn(sq, b, np_ * hn, device="cuda") + + with patch("torch.gather", side_effect=AssertionError("expanded gather must not be used")): + actual = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale=0.5 + ) + (actual * grad_output).sum().backward() + actual_grads = (query.grad.clone(), kv_full.grad.clone(), attn_sink.grad.clone()) + + query_ref = query.detach().clone().requires_grad_(True) + kv_ref = kv_full.detach().clone().requires_grad_(True) + sink_ref = attn_sink.detach().clone().requires_grad_(True) + rows = [] + for row in range(sq): + batches = [] + for batch in range(b): + heads = [] + for head in range(np_): + valid_indices = topk_indices[batch, row] + valid_indices = valid_indices[valid_indices >= 0].long() + logits = ( + torch.einsum( + "h,kh->k", query_ref[row, batch, head], kv_ref[valid_indices, batch] + ) + * 0.5 + ) + probabilities = torch.softmax( + torch.cat([logits, sink_ref[head : head + 1]]), dim=0 + ) + heads.append( + torch.einsum("k,kh->h", probabilities[:-1], kv_ref[valid_indices, batch]) + ) + batches.append(torch.cat(heads)) + rows.append(torch.stack(batches)) + expected = torch.stack(rows) + (expected * grad_output).sum().backward() + + torch.testing.assert_close(actual, expected) + for actual_grad, expected_grad in zip( + actual_grads, (query_ref.grad, kv_ref.grad, sink_ref.grad) + ): + torch.testing.assert_close(actual_grad, expected_grad) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def test_compressor_pooling_matches_fp32_forward_and_backward_oracle(): + torch.manual_seed(43) + kv = torch.randn(2, 4, 2, 6, dtype=torch.bfloat16, requires_grad=True) + score = torch.randn(2, 4, 2, 6, dtype=torch.bfloat16, requires_grad=True) + grad_output = torch.randn(2, 2, 6, dtype=torch.bfloat16) + + actual = _pool_compressor_values(kv, score, torch.bfloat16) + (actual * grad_output).sum().backward() + actual_grads = (kv.grad.clone(), score.grad.clone()) + + kv_ref = kv.detach().clone().requires_grad_(True) + score_ref = score.detach().clone().requires_grad_(True) + expected = ( + (kv_ref.float() * torch.softmax(score_ref, dim=1, dtype=torch.float32)) + .sum(dim=1) + .to(torch.bfloat16) + ) + (expected * grad_output).sum().backward() + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(actual_grads[0], kv_ref.grad, rtol=0, atol=0) + torch.testing.assert_close(actual_grads[1], score_ref.grad, rtol=0, atol=0) + + +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()), + ) + + +def test_compressed_sparse_attention_rejects_tensor_parallelism(): + config = _make_mla_config(num_attention_heads=2, csa_compress_ratios=[0] * 4) + with pytest.raises(ValueError, match="tensor-parallel size 1"): + CompressedSparseAttention( + config=config, + submodules=CompressedSparseAttentionSubmodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + pg_collection=_TwoRankPG(), + compress_ratio=0, + ) + + +def test_compressed_sparse_attention_rejects_query_head_mismatch(): + config = _make_mla_config(num_attention_heads=2, csa_compress_ratios=[0] * 4) + attention = CompressedSparseAttention( + config=config, + submodules=CompressedSparseAttentionSubmodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type="self", + pg_collection=_SingleRankPG(), + compress_ratio=0, + ) + + with pytest.raises(ValueError, match="query head count"): + attention( + query=torch.zeros(2, 1, 1, config.v_head_dim), key=None, value=None, attention_mask=None + ) + + +# =========================================================================== +# 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" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_projection_disables_fp8(self, compress_ratio, monkeypatch): + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + tracker = _DisabledContextTracker() + calls = [] + + for name, projection in ( + ('linear_wkv', compressor.linear_wkv), + ('linear_wgate', compressor.linear_wgate), + ): + original_forward = projection.forward + + def checked_forward(*args, _name=name, _forward=original_forward, **kwargs): + assert tracker.depth > 0, f"{_name} ran outside the FP8-disabled context" + calls.append(_name) + return _forward(*args, **kwargs) + + monkeypatch.setattr(projection, 'forward', checked_forward) + + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn( + compress_ratio * 2, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda' + ) + compressor(x) + + assert calls == ['linear_wkv', 'linear_wgate'] + assert tracker.entries == 1 + + +# =========================================================================== +# 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_weights_projection_disables_fp8(self, seqlen, monkeypatch): + tracker = _DisabledContextTracker() + self.indexer.cuda() + original_forward = self.indexer.linear_weights_proj.forward + + def checked_forward(*args, **kwargs): + assert tracker.depth > 0, "indexer weights projection ran under FP8" + return original_forward(*args, **kwargs) + + monkeypatch.setattr(self.indexer.linear_weights_proj, 'forward', checked_forward) + monkeypatch.setattr( + 'megatron.core.transformer.experimental_attention_variant.csa.get_fp8_disabled_context', + tracker, + ) + x = torch.randn(seqlen, 1, self.config.hidden_size, dtype=torch.bfloat16, device='cuda') + weights = self.indexer._project_weights(x) + + assert weights.shape == (seqlen, 1, self.config.dsa_indexer_n_heads) + assert tracker.entries == 1 + + @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() + + +class TestCSAHighPrecisionParams: + """Reference-checkpoint FP32 parameters survive BF16 model conversion.""" + + @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, 4, 4, 4]) + 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_ape_and_attn_sink_stay_fp32_after_bf16_conversion(self): + from megatron.core.transformer.module import Float16Module + + 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, + name="decoder.layers.0.self_attention.core_attention", + ) + + assert csa.attn_sink.dtype == torch.float32 + assert csa.compressor.ape.dtype == torch.float32 + assert csa.indexer.compressor.ape.dtype == torch.float32 + + bf16_module = Float16Module(config=self.config, module=csa) + + assert bf16_module.module.attn_sink.dtype == torch.float32 + assert bf16_module.module.compressor.ape.dtype == torch.float32 + assert bf16_module.module.indexer.compressor.ape.dtype == torch.float32 + assert bf16_module.module.compressor.linear_wkv.weight.dtype == torch.bfloat16 + assert bf16_module.module.compressor.linear_wgate.weight.dtype == torch.bfloat16 diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index 135c4802dd3..f7f190fbc62 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -28,8 +28,11 @@ DSAttention, DSAttentionSubmodules, FusedDSAIndexerLoss, + _compute_indexer_teacher_probabilities, + _normalize_indexer_teacher_target, _run_sparse_attention, _validate_nonpacked_cp_uniform_length, + bwd_fused_indexer_loss_naive, compute_dsa_indexer_loss, fused_qk_topk_naive, is_dsa_skip_topk_layer, @@ -50,6 +53,7 @@ build_fused_indexer_varlen_bounds, generate_varlen_mask_params_for_positions, masked_log_softmax, + masked_softmax, scatter_topk_into_index_mask, ) from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -1658,6 +1662,111 @@ def test_rotate_activation_dtype_check(self): rotate_activation(x) +def test_indexer_teacher_probability_matches_full_softmax_oracle(): + """External mass changes the denominator but not the compressed-key support.""" + attention_scores = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + valid_mask = torch.ones((1, 2, 2), dtype=torch.bool) + non_compressed_lse = torch.tensor([[[0.5, 1.5]]]) + + actual = _compute_indexer_teacher_probabilities( + attention_scores, valid_mask, non_compressed_lse + ) + expected_denominator = torch.logaddexp( + torch.logsumexp(attention_scores, dim=-1), non_compressed_lse + ) + expected = torch.exp(attention_scores - expected_denominator.unsqueeze(-1)) + torch.testing.assert_close(actual, expected) + + normalized = _normalize_indexer_teacher_target(actual.sum(dim=1), non_compressed_lse) + torch.testing.assert_close(normalized.sum(dim=-1), torch.ones((1, 2))) + + +def test_indexer_teacher_probability_rescales_underflowed_external_mass(): + """A huge external LSE preserves the relative compressed-key teacher target.""" + attention_scores = torch.tensor([[[[0.0, -1.0]], [[-2.0, 1.0]]]], dtype=torch.float32) + valid_mask = torch.ones((1, 1, 2), dtype=torch.bool) + non_compressed_lse = torch.tensor([[[1000.0], [1002.0]]], dtype=torch.float32) + + probabilities = _compute_indexer_teacher_probabilities( + attention_scores, valid_mask, non_compressed_lse + ) + normalized = _normalize_indexer_teacher_target(probabilities.sum(dim=1), non_compressed_lse) + + scores64 = attention_scores.double() + compressed_lse64 = torch.logsumexp(scores64, dim=-1) + full_lse64 = torch.logaddexp(non_compressed_lse.double(), compressed_lse64) + log_probabilities64 = scores64 - full_lse64.unsqueeze(-1) + expected = torch.softmax(torch.logsumexp(log_probabilities64, dim=1), dim=-1).float() + + assert torch.isfinite(probabilities).all() + torch.testing.assert_close(normalized.sum(dim=-1), torch.ones((1, 1))) + torch.testing.assert_close(normalized, expected) + + +def test_indexer_teacher_probability_keeps_fully_masked_rows_zero(): + """A fully masked compressed row remains finite and exactly zero.""" + attention_scores = torch.tensor([[[[1.0, 2.0]]]], requires_grad=True) + valid_mask = torch.zeros((1, 1, 2), dtype=torch.bool) + non_compressed_lse = torch.full((1, 1, 1), float("-inf")) + + target = _compute_indexer_teacher_probabilities( + attention_scores, valid_mask, non_compressed_lse + ) + assert torch.isfinite(target).all() + torch.testing.assert_close(target, torch.zeros_like(target)) + + normalized = _normalize_indexer_teacher_target(target.sum(dim=1), non_compressed_lse) + assert torch.isfinite(normalized).all() + torch.testing.assert_close(normalized, torch.zeros_like(normalized)) + target.sum().backward() + assert torch.isfinite(attention_scores.grad).all() + torch.testing.assert_close(attention_scores.grad, torch.zeros_like(attention_scores)) + + +def test_indexer_teacher_probability_preserves_legacy_path(): + """Without external mass, probability and target normalization stay unchanged.""" + attention_scores = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + valid_mask = torch.tensor([[[True, False], [False, False]]]) + expanded_valid_mask = valid_mask.unsqueeze(1) + + actual = _compute_indexer_teacher_probabilities(attention_scores, valid_mask) + expected = masked_softmax(attention_scores.float(), expanded_valid_mask, dim=-1) + torch.testing.assert_close(actual, expected) + + actual_normalized = _normalize_indexer_teacher_target(actual.sum(dim=1), None) + expected_normalized = dsa_indexer_loss.normalize_indexer_target(expected.sum(dim=1)) + torch.testing.assert_close(actual_normalized, expected_normalized) + + +def test_indexer_teacher_valid_zero_mass_has_zero_manual_gradient(): + """A valid row with zero teacher mass has zero manual KL gradient.""" + q = torch.ones((1, 1, 1, 1), dtype=torch.float32) + weights = torch.ones((1, 1, 1), dtype=torch.float32) + k = torch.ones((1, 1, 1), dtype=torch.float32) + query = torch.ones((1, 1, 1, 1), dtype=torch.float32) + key = torch.ones((1, 1, 1, 1), dtype=torch.float32) + + grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( + q=q, + weights=weights, + k=k, + query=query, + key=key, + topk_indices=torch.zeros((1, 1, 1), dtype=torch.long), + softmax_scale=1.0, + loss_coeff=1.0, + sparse_loss=False, + mask=torch.zeros((1, 1), dtype=torch.float32), + grad_loss=torch.tensor(1.0), + pg_collection=SimpleNamespace(tp=SimpleNamespace(size=lambda: 1)), + non_compressed_lse=torch.full((1, 1, 1), float("inf"), dtype=torch.float32), + ) + + torch.testing.assert_close(grad_q, torch.zeros_like(grad_q)) + torch.testing.assert_close(grad_weights, torch.zeros_like(grad_weights)) + torch.testing.assert_close(grad_k, torch.zeros_like(grad_k)) + + @pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) class TestComputeDSAIndexerLoss: """Test compute_dsa_indexer_loss function."""