diff --git a/gpt_builders.py b/gpt_builders.py index 13c914acc56..972ac738d50 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -123,6 +123,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, + args.experimental_attention_variant, moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, qk_l2_norm=args.qk_l2_norm, use_kitchen=config.use_kitchen, @@ -141,6 +142,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, + args.experimental_attention_variant, moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, normalization=args.normalization, use_kitchen=config.use_kitchen, diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py new file mode 100644 index 00000000000..b7f613e87d6 --- /dev/null +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -0,0 +1,114 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from typing import Optional + +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.models.backends import BackendSpecProvider +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexer, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, +) +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.multi_latent_attention import ( + MLASelfAttention, + MLASelfAttentionSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules + + +def get_dsa_module_spec_for_backend( + backend: BackendSpecProvider, + qk_layernorm: Optional[bool] = False, + qk_l2_norm: Optional[bool] = False, + multi_latent_attention: Optional[bool] = False, + num_experts: Optional[int] = None, + mlp: Optional[ModuleSpec] = None, +) -> ModuleSpec: + """Helper function to get module spec for Sparse Attention.""" + assert multi_latent_attention, "Currently only MLA supports sparse attention." + assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + linear_q_up_proj = ( + backend.column_parallel_layer_norm_linear() + if qk_layernorm + else backend.column_parallel_linear() + ) + linear_kv_up_proj = ( + backend.column_parallel_layer_norm_linear() + if qk_layernorm + else backend.column_parallel_linear() + ) + + # Because TransformerEngine does not support sparse attention yet, we use local + # implementation whether the backend is TransformerEngine or not. + core_attention = ModuleSpec( + module=DSAttention, + submodules=DSAttentionSubmodules( + indexer=ModuleSpec( + module=DSAIndexer, + submodules=DSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_wk=backend.linear(), + k_norm=backend.layer_norm(rms_norm=False, for_qk=True), + linear_weights_proj=backend.linear(), + ), + ) + ), + ) + + attention = ModuleSpec( + module=MLASelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=MLASelfAttentionSubmodules( + linear_q_proj=backend.column_parallel_linear(), + linear_q_down_proj=backend.linear(), + linear_q_up_proj=linear_q_up_proj, + linear_kv_down_proj=backend.linear(), + linear_kv_up_proj=linear_kv_up_proj, + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=IdentityOp, + kv_layernorm=IdentityOp, + ), + ) + + return ModuleSpec( + module=TransformerLayer, + submodules=TransformerLayerSubmodules( + input_layernorm=backend.layer_norm(), + self_attention=attention, + self_attn_bda=get_bias_dropout_add, + pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + ), + ) + + +def get_experimental_attention_variant_module_spec_for_backend( + backend: BackendSpecProvider, + experimental_attention_variant: Optional[str] = None, + qk_layernorm: Optional[bool] = False, + qk_l2_norm: Optional[bool] = False, + multi_latent_attention: Optional[bool] = False, + num_experts: Optional[int] = None, + mlp: Optional[ModuleSpec] = None, +) -> ModuleSpec: + """Helper function to get module spec for Attention""" + if experimental_attention_variant == "dsa": + return get_dsa_module_spec_for_backend( + backend=backend, + qk_layernorm=qk_layernorm, + qk_l2_norm=qk_l2_norm, + multi_latent_attention=multi_latent_attention, + num_experts=num_experts, + mlp=mlp, + ) + else: + raise ValueError( + f"Invalid experimental attention variant: {experimental_attention_variant}" + ) diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index 8615298a1e2..604b14c8ffe 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -9,6 +9,9 @@ InferenceSpecProvider, LocalSpecProvider, ) +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec_for_backend, +) from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules from megatron.core.transformer.enums import AttnMaskType, LayerType @@ -176,6 +179,7 @@ def get_gpt_layer_with_transformer_engine_spec( moe_grouped_gemm: Optional[bool] = False, qk_layernorm: Optional[bool] = False, multi_latent_attention: Optional[bool] = False, + experimental_attention_variant: Optional[str] = None, fp8: Optional[str] = None, # pylint: disable=unused-argument moe_use_legacy_grouped_gemm: Optional[bool] = False, qk_l2_norm: Optional[bool] = False, @@ -192,6 +196,9 @@ def get_gpt_layer_with_transformer_engine_spec( num_experts (int, optional): Number of experts. Defaults to None. moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False. qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False. + multi_latent_attention (bool, optional): To use MLA. Defaults to False. + experimental_attention_variant (str, optional): The type of experimental attention variant. + Defaults to None. fp8 (str, optional): Deprecated. For temporary Nemo compatibility. moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP. Defaults to False. @@ -232,6 +239,17 @@ def get_gpt_layer_with_transformer_engine_spec( use_te_activation_func=use_te_activation_func, ) + if experimental_attention_variant is not None: + return get_experimental_attention_variant_module_spec_for_backend( + backend=backend, + experimental_attention_variant=experimental_attention_variant, + qk_layernorm=qk_layernorm, + qk_l2_norm=qk_l2_norm, + multi_latent_attention=multi_latent_attention, + num_experts=num_experts, + mlp=mlp, + ) + if multi_latent_attention: assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." linear_q_up_proj = ( @@ -310,6 +328,7 @@ def get_gpt_layer_local_spec( moe_grouped_gemm: Optional[bool] = False, qk_layernorm: Optional[bool] = False, multi_latent_attention: Optional[bool] = False, + experimental_attention_variant: Optional[str] = None, fp8: Optional[str] = None, # pylint: disable=unused-argument moe_use_legacy_grouped_gemm: Optional[bool] = False, normalization: Optional[str] = None, @@ -325,6 +344,9 @@ def get_gpt_layer_local_spec( num_experts (int, optional): Number of experts. Defaults to None. moe_grouped_gemm (bool, optional): To use Grouped GEMM. Defaults to False. qk_layernorm (bool, optional): To use layernorm for queries/keys. Defaults to False. + multi_latent_attention (bool, optional): To use MLA. Defaults to False. + experimental_attention_variant (str, optional): The type of experimental attention variant. + Defaults to None. fp8 (str, optional): Deprecated. For temporary Nemo compatibility. moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP. Defaults to False. @@ -334,6 +356,11 @@ def get_gpt_layer_local_spec( ModuleSpec: Module specification with Megatron-Core modules """ + if experimental_attention_variant is not None: + raise NotImplementedError( + "Experimental attention variant is not supported with local spec yet." + ) + if use_kitchen: assert HAVE_KITCHEN backend = KitchenSpecProvider( diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py new file mode 100644 index 00000000000..88b4713dc60 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -0,0 +1,822 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +import copy +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch + +from megatron.core import parallel_state +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig + +try: + from fast_hadamard_transform import hadamard_transform +except ImportError: + hadamard_transform = None + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Apply Hadamard rotation activation. + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L424-L428 + + Args: + x: Input tensor (must be bfloat16). + + Returns: + Rotated tensor. + """ + assert ( + x.dtype == torch.bfloat16 + ), f"rotate_activation only support bf16 input, but got {x.dtype}" + assert hadamard_transform is not None, "fast_hadamard_transform is not installed." + hidden_size = x.size(-1) + return hadamard_transform(x, scale=hidden_size**-0.5) + + +class DSAIndexerLossLoggingHelper: + """Helper class for logging sparse attention indexer losses.""" + + tracker = {} + + @staticmethod + def save_loss_to_tracker( + loss: torch.Tensor, + layer_number: int, + num_layers: int, + reduce_group: torch.distributed.ProcessGroup = None, + avg_group: torch.distributed.ProcessGroup = None, + ): + """Save the indexer loss for logging. + + Args: + loss: The loss tensor. + layer_number: Layer index of the loss, 1-indexed. + num_layers: The number of total layers. + reduce_group: The group for reducing the loss. + avg_group: The group for averaging the loss. + """ + # Skip indexer loss logging if layer_number is None. + if layer_number is None: + return + + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" not in tracker: + tracker["values"] = torch.zeros(num_layers, device=torch.cuda.current_device()) + tracker["values"][layer_number - 1] += loss.detach() + tracker["reduce_group"] = reduce_group + tracker["avg_group"] = avg_group + + @staticmethod + def clean_loss_in_tracker(): + """Clear the indexer losses.""" + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" in tracker: + tracker["values"].zero_() + tracker["reduce_group"] = None + tracker["avg_group"] = None + + @staticmethod + def reduce_loss_in_tracker(): + """Collect and reduce the indexer losses across ranks.""" + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" not in tracker: + return + values = tracker["values"] + + torch.distributed.all_reduce( + values, group=parallel_state.get_pipeline_model_parallel_group() + ) + # Reduce indexer losses across ranks. + if tracker.get('reduce_group') is not None: + torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) + if tracker.get('avg_group') is not None: + torch.distributed.all_reduce( + values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG + ) + torch.distributed.all_reduce( + values, + group=parallel_state.get_data_parallel_group(with_context_parallel=False), + op=torch.distributed.ReduceOp.AVG, + ) + + @staticmethod + def track_indexer_metrics( + loss_scale: float, + iteration: int, + writer, + wandb_writer=None, + total_loss_dict=None, + per_layer_logging: bool = False, + ): + """Track the sparse attention indexer metrics for logging. + + Args: + loss_scale: Scale factor for the loss. + iteration: Current training iteration. + writer: TensorBoard writer. + wandb_writer: Weights & Biases writer. + total_loss_dict: Dictionary to accumulate total losses. + per_layer_logging: Whether to log per-layer losses. + """ + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" not in tracker: + return + + indexer_loss_values = tracker["values"] * loss_scale + num_layers = indexer_loss_values.shape[0] + + # Average across all layers (assuming all layers have sparse attention) + avg_indexer_loss = indexer_loss_values.sum() / num_layers + + # Log average loss + if total_loss_dict is not None: + if "indexer loss" in total_loss_dict: + total_loss_dict["indexer loss"] += avg_indexer_loss + else: + total_loss_dict["indexer loss"] = avg_indexer_loss + + if writer is not None: + writer.add_scalar("indexer loss", avg_indexer_loss, iteration) + + if wandb_writer is not None: + wandb_writer.log({"indexer loss": avg_indexer_loss}, iteration) + + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + + +def compute_dsa_indexer_loss( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + pg_collection: ProcessGroupCollection, +) -> torch.Tensor: + """ + Compute KL divergence loss between index_scores and true attention_scores. + + This loss trains the indexer to predict which tokens are important by matching the distribution + of true attention scores. + + Reference: Section 2.1 of + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf + + Args: + index_scores: Scores predicted by indexer [batch, seqlen_q, seqlen_k]. + topk_indices: Top-k indices [batch, seqlen_q, index_topk]. + query: Query tensor [seqlen_q, batch, heads, dim]. + key: Key tensor [seqlen_k, batch, heads, dim]. + softmax_scale: Scale coefficient after q @ k^T. + loss_coeff: Coefficient for the indexer KL divergence loss. + sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk + indices will be used to compute the loss. + pg_collection: Process group collection, must have TP process group. + + Returns: + index_loss: KL divergence loss (scalar). + """ + sq, b, np, hn = query.size() + sk = key.size(0) + + # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] + query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) + # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] + key = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) + # Compute attention scores [b * np, sq, sk] + attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale + # Reshape to [b, np, sq, sk] + attention_scores = attention_scores.reshape(b, np, sq, sk) + + # causal_mask [sq, sk] + causal_mask = torch.triu( + torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), + diagonal=1, + ) + # index_mask [b, sq, sk] + index_mask = torch.full( + (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device + ).scatter_(-1, topk_indices, 0) + + # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] + attention_scores += causal_mask.view(1, 1, sq, sk) + if sparse_loss: + # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] + attention_scores += index_mask.view(b, 1, sq, sk) + # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] + index_scores += index_mask + + # [b, np, sq, sk] -> [b, np, sq, sk] + attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) + # [b, sq, sk] -> [b, sq, sk] + index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + + # Sum attention scores across heads. + # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] + attention_scores = attention_scores.sum(dim=1) + if pg_collection.tp.size() > 1: + # attention scores are scattered to TP ranks in head dimension. + torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) + # L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are + # obtained from softmax so they are already non-negative. + attention_scores = attention_scores / attention_scores.sum(dim=-1, keepdim=True) + + # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) + # kl_per_element [b, sq, sk] + kl_per_element = attention_scores * ( + torch.log(attention_scores + 1e-10) - torch.log(index_scores + 1e-10) + ) + + # [b, sq, sk] -> [b, sq] -> [1] + # Each token has same weight in the loss. + kl_div = kl_per_element.sum(dim=-1).mean() + + # Scale by coefficient. + indexer_loss = kl_div * loss_coeff + + return indexer_loss + + +class DSAIndexerLossAutoScaler(torch.autograd.Function): + """An AutoScaler that triggers the backward pass and scales the grad for indexer loss. + + This custom autograd function attaches a KL divergence loss to the activation + to train the indexer to predict attention scores without affecting the forward pass. + """ + + main_loss_backward_scale: torch.Tensor = None + + @staticmethod + def forward(ctx, output: torch.Tensor, indexer_loss: torch.Tensor): + """Preserve the indexer_loss by storing it in the context to avoid garbage collection. + + Args: + output: The output tensor (activation). + indexer_loss: The indexer KL divergence loss tensor. + + Returns: + torch.Tensor: The output tensor unchanged. + """ + ctx.save_for_backward(indexer_loss) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + """Compute and scale the gradient for indexer loss. + + Args: + grad_output: The gradient of the output. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled indexer loss + gradient. + """ + (indexer_loss,) = ctx.saved_tensors + if DSAIndexerLossAutoScaler.main_loss_backward_scale is None: + DSAIndexerLossAutoScaler.main_loss_backward_scale = torch.tensor( + 1.0, device=indexer_loss.device + ) + indexer_loss_backward_scale = DSAIndexerLossAutoScaler.main_loss_backward_scale + scaled_indexer_loss_grad = torch.ones_like(indexer_loss) * indexer_loss_backward_scale + return grad_output, scaled_indexer_loss_grad + + @staticmethod + def set_loss_scale(scale: torch.Tensor): + """Set the scale of the indexer loss. + + Args: + scale: The scale value to set. + """ + if DSAIndexerLossAutoScaler.main_loss_backward_scale is None: + DSAIndexerLossAutoScaler.main_loss_backward_scale = scale + else: + DSAIndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) + + +@dataclass +class DSAIndexerSubmodules: + """ + Configuration class for specifying the submodules of an DSA Indexer. + + Args: + linear_wq_b: Linear projection for query bottleneck expansion. + linear_wk: Linear projection for key. + k_norm: Layer normalization for key. + linear_weights_proj: Linear projection for attention weights. + """ + + linear_wq_b: Union[ModuleSpec, type] = None + linear_wk: Union[ModuleSpec, type] = None + k_norm: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + + +@dataclass +class DSAttentionSubmodules: + """ + Configuration class for specifying the submodules of DSAttention. + + Args: + indexer: DSA Indexer module for computing sparse attention indices. + """ + + indexer: Union[ModuleSpec, type] = None + + +class DSAIndexer(MegatronModule): + """ + DSA Lightning Indexer for DeepSeek Sparse Attention. + + Computes index scores to identify the top-k most relevant key-value pairs for each query in + sparse attention. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480 + """ + + def __init__( + self, + config: TransformerConfig, + submodules: DSAIndexerSubmodules, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + """Initialize the indexer. + + Args: + config (TransformerConfig): The configuration for the transformer model. + submodules (DSAIndexerSubmodules): Indexer submodules specification. + pg_collection (ProcessGroupCollection, optional): Process groups for the indexer. + """ + super().__init__(config=config) + self.hidden_size = self.config.hidden_size + self.qk_pos_emb_head_dim = self.config.qk_pos_emb_head_dim + self.q_lora_rank = ( + self.config.q_lora_rank + if self.config.q_lora_rank is not None + else self.config.hidden_size + ) + + self.index_n_heads = self.config.dsa_indexer_n_heads + self.index_head_dim = self.config.dsa_indexer_head_dim + self.index_topk = self.config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim**-0.5 + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + self.pg_collection = pg_collection + + # Initialize Position Embedding. + if self.config.rope_type == 'rope': + self.rotary_pos_emb = RotaryEmbedding( + self.qk_pos_emb_head_dim, + rotary_percent=self.config.rotary_percent, + rotary_base=self.config.rotary_base, + cp_group=self.pg_collection.cp, + ) + elif self.config.rope_type == 'yarn': + self.rotary_pos_emb = YarnRotaryEmbedding( + self.qk_pos_emb_head_dim, + rotary_base=self.config.rotary_base, + scaling_factor=self.config.rotary_scaling_factor, + original_max_position_embeddings=self.config.original_max_position_embeddings, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + mscale=self.config.mscale, + mscale_all_dim=self.config.mscale_all_dim, + cp_group=self.pg_collection.cp, + ) + else: + raise ValueError( + f'Unsupported RoPE type: {self.config.rope_type}, supported types are "rope" and ' + f'"yarn"' + ) + + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + self.linear_wk = build_module( + submodules.linear_wk, + self.hidden_size, + self.index_head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + k_norm_config = copy.copy(self.config) + k_norm_config.normalization = "LayerNorm" + self.k_norm = build_module( + submodules.k_norm, + config=k_norm_config, + hidden_size=self.index_head_dim, + eps=self.config.layernorm_epsilon, + ) + + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + def _apply_rope(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor, mscale: float): + """Apply RoPE to the input tensor.""" + # x_nope [seqlen, batch, *, index_head_dim - qk_pos_emb_head_dim] + # x_pe [seqlen, batch, *, qk_pos_emb_head_dim] + x_nope, x_pe = torch.split( + x, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + ) + # [seqlen, batch, *, index_head_dim] + x = torch.cat([x_nope, x_pe], dim=-1) + return x + + def _compute_index_scores( + self, q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor + ) -> torch.Tensor: + """ + Perform index score using BF16 precision. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/kernel.py#L254-L274 + This is a BF16 implementation of the `fp8_index` logic: + 1. Compute attention scores: q @ k^T; + 2. Apply ReLU activation; + 3. Weight by attention weights; + 4. Sum across attention heads. + + Args: + q: BF16 [seqlen_q, batch, index_n_heads, index_head_dim], the query tensor. + weights: BF16 [seqlen_q, batch, index_n_heads], the attention weights. + k: BF16 [seqlen_k, batch, index_head_dim], the key tensor. + + Returns: + index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores. + """ + # Compute attention scores: q @ k^T + # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T + # -> [seqlen_q, batch, index_n_heads, seqlen_k] + index_scores = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) + + # Apply ReLU activation. + index_scores = torch.relu(index_scores) + + # Weight each head by attention weights. + # [seqlen_q, batch, index_n_heads, seqlen_k] * [seqlen_q, batch, index_n_heads, 1] + # -> [seqlen_q, batch, index_n_heads, seqlen_k] + index_scores = index_scores * weights.unsqueeze(-1) + + # Sum across attention heads. + # [seqlen_q, batch, index_n_heads, seqlen_k] -> [seqlen_q, batch, seqlen_k] + index_scores = index_scores.sum(dim=2) + + # Transpose to [batch, seqlen_q, seqlen_k]. + index_scores = index_scores.transpose(0, 1) + + return index_scores + + def forward_with_scores( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for DSA Indexer that returns both index scores and top-k indices. + + This is used when KL loss is enabled to compare indexer scores with true attention scores. + + Args: + x: hidden states [seqlen, batch, hidden_size]. + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + mask: Attention mask [batch, seqlen, seqlen]. + packed_seq_params: Packed sequence parameters for variable length sequences. + + Returns: + index_scores: Index scores [batch, seqlen, seqlen]. + topk_indices: Top-k indices [batch, seqlen, index_topk]. + """ + assert packed_seq_params is None, "Packed sequence is not supported for DSAttention" + + # ========================================= + # Prepare RoPE params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + None, None, x, self.config, packed_seq_params + ) + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + mscale = 1.0 + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # Gather inputs if sp is enabled + # ========================================= + if self.config.sequence_parallel and self.pg_collection.tp.size() > 1: + x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp) + qr = gather_from_sequence_parallel_region(qr, group=self.pg_collection.tp) + + # ========================================= + # Get sequence length and batch size + # ========================================= + seqlen, bsz, _ = x.size() + + # ========================================= + # q linear and apply rope to q + # ========================================= + # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim] + q, _ = self.linear_wq_b(qr) + # [seqlen, batch, index_n_heads * index_head_dim] + # -> [seqlen, batch, index_n_heads, index_head_dim] + q = q.reshape(seqlen, bsz, self.index_n_heads, self.index_head_dim) + q = self._apply_rope(q, rotary_pos_emb, mscale) + + # ========================================= + # k linear and apply rope to k + # ========================================= + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] + k, _ = self.linear_wk(x) + k = self.k_norm(k) + # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] + k = k.reshape(seqlen, bsz, 1, self.index_head_dim) + k = self._apply_rope(k, rotary_pos_emb, mscale) + # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] + k = k.reshape(seqlen, bsz, self.index_head_dim) + + # ========================================= + # Rotate activation + # ========================================= + q = rotate_activation(q) + k = rotate_activation(k) + + # ========================================= + # Compute index scores + # ========================================= + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] + weights, _ = self.linear_weights_proj(x) + weights = weights * (self.index_n_heads**-0.5) * self.softmax_scale + # [batch, seqlen, seqlen] + index_scores = self._compute_index_scores(q, weights, k) + if mask is not None: + assert mask.dtype == index_scores.dtype, "Mask dtype must match index scores dtype" + index_scores = index_scores + mask + + # ========================================= + # Select top-k indices + # ========================================= + topk_k = min(self.index_topk, seqlen) + # [batch, seqlen, index_topk] + topk_indices = index_scores.topk(topk_k, dim=-1)[1] + + return index_scores, topk_indices + + def forward( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ): + """ + Forward pass for DSA Indexer. + + Args: + x: hidden states [seqlen, batch, hidden_size]. + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + mask: Attention mask [batch, seqlen, seqlen]. + packed_seq_params: Packed sequence parameters for variable length sequences. + + Returns: + topk_indices: Top-k indices for sparse attention [batch, seqlen, index_topk]. + """ + _, topk_indices = self.forward_with_scores(x, qr, mask, packed_seq_params) + return topk_indices + + +def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): + """ + Unfused sparse attention implementation. + """ + sq, b, np, hn = query.size() + skv = key.size(0) + hnv = value.size(3) + + # =================================== + # Raw attention scores [b, np, sq, skv] + # =================================== + # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] + query = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) + # [skv, b, np, hn] -> [b, np, hn, skv] -> [b * np, hn, skv] + key = key.permute(1, 2, 3, 0).reshape(b * np, hn, skv) + # Compute attention scores [b * np, sq, skv] + attention_scores = torch.bmm(query.float(), key.float()) * softmax_scale + # Reshape to [b, np, sq, skv] + attention_scores = attention_scores.reshape(b, np, sq, skv) + + # =================================== + # Apply sparse mask from indexer + # =================================== + # index_mask [b, sq, skv] + index_mask = torch.full((b, sq, skv), float("-inf"), device=attention_scores.device) + index_mask.scatter_(-1, topk_indices, 0) + # causal_mask [sq, skv] + causal_mask = torch.triu( + torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=index_mask.device), + diagonal=1, + ) + # [b, sq, skv] + [1, sq, skv] -> [b, sq, skv] + index_mask += causal_mask.view(1, sq, skv) + # [b, np, sq, skv] + [b, 1, sq, skv] -> [b, np, sq, skv] + attention_scores += index_mask.unsqueeze(1) + attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) + + # =================================== + # Output + # =================================== + # [skv, b, np, hnv] -> [b, np, skv, hnv] -> [b * np, skv, hnv] + value = value.permute(1, 2, 0, 3).reshape(b * np, skv, hnv) + # Reshape attention_scores: [b, np, sq, skv] -> [b * np, sq, skv] + attention_scores = attention_scores.reshape(b * np, sq, skv) + # Compute output: [b * np, sq, hnv] + output = torch.bmm(attention_scores.to(value.dtype), value) + # Reshape output: [b * np, sq, hnv] -> [b, np, sq, hnv] -> [sq, b, np, hnv] + output = output.reshape(b, np, sq, hnv).permute(2, 0, 1, 3).contiguous() + # Flatten: [sq, b, np, hnv] -> [sq, b, np * hnv] + output = output.reshape(sq, b, np * hnv) + return output + + +class DSAttention(MegatronModule): + """ + This module implements sparse attention mechanism using an DSA Indexer to compute top-k + attention indices for reducing computational complexity. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597 + """ + + def __init__( + self, + config: TransformerConfig, + submodules: DSAttentionSubmodules, + 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: ProcessGroupCollection = None, + ): + super().__init__(config=config) + + self.layer_number = layer_number + + self.indexer = build_module( + submodules.indexer, config=self.config, pg_collection=pg_collection + ) + + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt( + k_channels if k_channels is not None else config.kv_channels + ) + self.softmax_scale = softmax_scale + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params: PackedSeqParams = None, + ): + """ + Forward pass for Sparse Attention. + + Args: + query: Query tensor [sq, b, np, hn]. + key: Key tensor [skv, b, np, hn]. + value: Value tensor [skv, b, np, hnv]. + x: Original hidden states [sq, b, hidden_size]. + qr: Low-rank query representation [sq, b, q_lora_rank]. + attention_mask: Attention mask tensor [b, 1, sq, sk]. + attn_mask_type: Type of attention mask. + attention_bias: Optional attention bias. + packed_seq_params: Packed sequence parameters. + + Returns: + output: Output tensor [sq, b, hidden_size] + """ + sq, b, np, hn = query.size() + skv = key.size(0) + hnv = value.size(3) + + # Detach x and qr to prevent gradients of indexer from flowing back to the main model. + x = x.detach() + qr = qr.detach() + + # Get a FP32 mask with -inf for masked positions. + if attn_mask_type is not None: + assert attn_mask_type == AttnMaskType.causal, 'Only causal mask is supported for now' + # Generate upper triangular mask with -inf above diagonal, 0 elsewhere + # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) + # float_mask [sq, skv] + float_mask = torch.triu( + torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device), + diagonal=1, + ) + else: + assert attention_mask.shape == (b, 1, sq, skv), 'attention_mask shape mismatch' + # [b, 1, sq, skv] -> [b, sq, skv] + mask = attention_mask.squeeze() + # float_mask [b, sq, skv] + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( + mask, float('-inf') + ) + + # =================================== + # Get index scores and top-k indices + # =================================== + index_scores, topk_indices = self.indexer.forward_with_scores( + x, qr, mask=float_mask, packed_seq_params=packed_seq_params + ) + + # =================================== + # Run sparse attention kernel + # =================================== + output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) + + # =================================== + # Attach indexer loss + # =================================== + if self.training and torch.is_grad_enabled(): + # Compute KL divergence loss between indexer scores and true attention scores + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + indexer_loss = compute_dsa_indexer_loss( + index_scores, + topk_indices, + query.detach(), + key.detach(), + self.softmax_scale, + indexer_loss_coeff, + getattr(self.config, "dsa_indexer_use_sparse_loss", False), + self.indexer.pg_collection, + ) + # Save indexer loss for logging + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + ) + # Attach loss to output + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + return output diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index d25f417fc50..e8296fb3019 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -244,7 +244,7 @@ def forward( # self or cross attn. # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128] with off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") as hidden_states: - query, key, value = self.get_query_key_value_tensors( + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, @@ -282,6 +282,12 @@ def forward( ) else: if inference_context is None or inference_context.is_static_batching(): + extra_kwargs = {} + if self.config.experimental_attention_variant == "dsa": + # For dsa we need to pass in the original hidden states and the compressed + # query representation. + extra_kwargs["x"] = hidden_states + extra_kwargs["qr"] = q_compressed with off_interface( self.offload_core_attention and self.training, query, "core_attn" ) as query: @@ -292,6 +298,7 @@ def forward( attention_mask, packed_seq_params=packed_seq_params, attn_mask_type=attn_mask_type, + **extra_kwargs, ) elif self.cache_mla_latents: # Dynamic batching attention kernel. @@ -617,6 +624,16 @@ def get_query_key_value_tensors( kv_compressed = kv_compressed.squeeze(1) k_pos_emb = k_pos_emb.squeeze(1) + # ========================================= + # Apply norm + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = self.q_layernorm(q_compressed) + + kv_compressed = self.kv_layernorm(kv_compressed) + # ========================================= # QKV up projection and RoPE apply # ========================================= @@ -627,7 +644,6 @@ def qkv_up_proj_and_rope_apply_for_cached_latent_kv( if self.config.q_lora_rank is not None: # q_compressed: [num_tokens, q_lora_rank] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] - q_compressed = self.q_layernorm(q_compressed) q, _ = self.linear_q_up_proj(q_compressed) else: # q_compressed: [num_tokens, hidden_size] @@ -637,8 +653,6 @@ def qkv_up_proj_and_rope_apply_for_cached_latent_kv( # q: [num_tokens, n, q_head_dim] q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) - kv_compressed = self.kv_layernorm(kv_compressed) - # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] k_pos_emb = torch.unsqueeze(k_pos_emb, -2) @@ -702,7 +716,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po if self.config.q_lora_rank is not None: # q_compressed: [num_tokens, q_lora_rank] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] - q_compressed = self.q_layernorm(q_compressed) q, _ = self.linear_q_up_proj(q_compressed) else: # q_compressed: [num_tokens, hidden_size] @@ -712,8 +725,6 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po # q: [num_tokens, n, q_head_dim] q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) - kv_compressed = self.kv_layernorm(kv_compressed) - # kv: [num_tokens, n * (qk_head_dim + v_head_dim)] kv, _ = self.linear_kv_up_proj(kv_compressed) @@ -838,7 +849,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb ) - return query, key, value + return query, key, value, q_compressed, kv_compressed def uncompress_kv_from_cache(self, kv_cached): """ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e94a7997595..4ec3e18530b 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -225,6 +225,28 @@ class TransformerConfig(ModelParallelConfig): A list of integers: Defines a custom pattern where 1 means skip RoPE and 0 means apply RoPE. For example, [0,1,1,0] means: apply RoPE, skip RoPE, skip RoPE, apply RoPE.""" + #################### + # attention variant + #################### + experimental_attention_variant: Optional[str] = None + """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + + dsa_indexer_n_heads: Optional[int] = None + """Number of DSA indexer heads.""" + + dsa_indexer_head_dim: Optional[int] = None + """Dimension per DSA indexer head.""" + + dsa_indexer_topk: Optional[int] = None + """Number of top-k tokens to select in DSA indexer.""" + + dsa_indexer_loss_coeff: Optional[float] = None + """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss.""" + + dsa_indexer_use_sparse_loss: bool = False + """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the + top-k indices.""" + #################### # initialization #################### @@ -1835,11 +1857,18 @@ def __post_init__(self): assert not self.add_qkv_bias assert not self.use_kitchen + if self.experimental_attention_variant == "dsa": + assert ( + self.context_parallel_size == 1 + ), "Currently context parallelism is not supported by DSAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" + if self.inference_fuse_tp_communication: assert self.transformer_impl == "inference_optimized", ( "inference_fuse_tp_communication is only supported " "for inference_optimized transformer implementation." ) + if self.batch_invariant_mode: assert ( self.attention_backend == AttnBackend.flash diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 49406ca30ff..979fbd4d5e1 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -70,6 +70,7 @@ def add_megatron_arguments(parser: argparse.ArgumentParser): parser = _add_vision_args(parser) parser = _add_moe_args(parser) parser = _add_mla_args(parser) + parser = _add_experimental_attention_variant_args(parser) parser = _add_heterogeneous_args(parser) parser = _add_logging_args(parser) parser = _add_straggler_detector_args(parser) @@ -3370,6 +3371,24 @@ def _add_mla_args(parser): return parser +def _add_experimental_attention_variant_args(parser): + group = parser.add_argument_group(title="experimental_attention_variant") + group.add_argument('--experimental-attention-variant', default=None, choices=['gated_delta_net', 'dsa'], type=str, + help='Type of attention variant to use. Currently support gated_delta_net and dsa.') + # DSA + group.add_argument('--dsa-indexer-n-heads', default=None, type=int, + help='Number of indexer heads for sparse attention. If not set, defaults to num-attention-heads.') + group.add_argument('--dsa-indexer-head-dim', default=None, type=int, + help='Dimension per indexer head for sparse attention. If not set, defaults to kv-channels.') + group.add_argument('--dsa-indexer-topk', default=None, type=int, + help='Number of top-k tokens to select in sparse attention indexer.') + group.add_argument('--dsa-indexer-loss-coeff', default=None, type=float, + help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') + group.add_argument('--dsa-indexer-use-sparse-loss', action='store_true', + help='Use sparse indexer loss. If set, the indexer loss will be computed using the top-k indices.') + + return parser + def _add_heterogeneous_args(parser): """ Heterogeneous models refer to transformer architectures where individual layers can differ diff --git a/megatron/training/training.py b/megatron/training/training.py index 5ac4d947c4a..9551c89bf60 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -107,6 +107,7 @@ from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler from megatron.core.transformer.moe import upcycling_utils from megatron.core.transformer.moe.moe_utils import track_moe_metrics +from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper from megatron.core.parallel_state import ( destroy_global_memory_buffer, @@ -1771,6 +1772,16 @@ def training_log( MTPLossLoggingHelper.track_mtp_metrics( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) + # Track sparse attention indexer loss + if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: + indexer_loss_scale = 1 / get_num_microbatches() + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=indexer_loss_scale, + iteration=iteration, + writer=writer, + wandb_writer=wandb_writer, + total_loss_dict=total_loss_dict, + ) if iteration % args.log_interval == 0: if args.record_memory_history and (is_last_rank() or torch.distributed.get_backend() == 'fake'): snapshot = torch.cuda.memory._snapshot() diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py index c4cfbc2e56c..d3a58710907 100644 --- a/tests/unit_tests/models/test_mamba_moe_model.py +++ b/tests/unit_tests/models/test_mamba_moe_model.py @@ -79,11 +79,17 @@ "disable_bf16_reduced_precision_matmul": False, "disable_parameter_transpose_cache": False, "distribute_saved_activations": False, + "dsa_indexer_head_dim": None, + "dsa_indexer_loss_coeff": None, + "dsa_indexer_n_heads": None, + "dsa_indexer_topk": None, + "dsa_indexer_use_sparse_loss": False, "embedding_init_method": {}, "embedding_init_method_std": 0.014, "enable_autocast": False, "enable_cuda_graph": False, "ep_overlap_early_attn_memory_release": False, + "experimental_attention_variant": None, "expert_model_parallel_size": 4, "expert_tensor_parallel_size": 1, "external_cuda_graph": False, diff --git a/tests/unit_tests/transformer/test_attention_variant_dsa.py b/tests/unit_tests/transformer/test_attention_variant_dsa.py new file mode 100644 index 00000000000..bd106aa6f0e --- /dev/null +++ b/tests/unit_tests/transformer/test_attention_variant_dsa.py @@ -0,0 +1,1271 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +import megatron.core.parallel_state as parallel_state +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexer, + DSAIndexerLossAutoScaler, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, + compute_dsa_indexer_loss, + rotate_activation, +) +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. + + This is a simple identity-like transformation that preserves shape and applies scaling. + """ + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in dsa module if not installed.""" + if not HAVE_HADAMARD: + with patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + mock_hadamard_transform, + ): + yield + else: + yield + + +class TestRotateActivation: + """Test rotate_activation function.""" + + @pytest.fixture(scope='function', 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() + + def test_rotate_activation_shape(self): + """Test that rotate_activation preserves shape.""" + batch_size = 2 + seq_len = 16 + hidden_size = 128 + + x = torch.randn(seq_len, batch_size, hidden_size, dtype=torch.bfloat16).cuda() + output = rotate_activation(x) + + assert output.shape == x.shape + assert output.dtype == torch.bfloat16 + + def test_rotate_activation_dtype_check(self): + """Test that rotate_activation only accepts bfloat16.""" + x = torch.randn(16, 2, 128, dtype=torch.float32).cuda() + + with pytest.raises(AssertionError, match="only support bf16"): + rotate_activation(x) + + +@pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) +class TestComputeDSAIndexerLoss: + """Test compute_dsa_indexer_loss function.""" + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp']) + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_loss_shape(self, seqlen_and_topk): + """Test that indexer loss returns a scalar.""" + batch_size = 2 + seqlen = seqlen_and_topk[0] + num_heads = 4 + head_dim = 128 + index_topk = seqlen_and_topk[1] + + # Create dummy index scores + index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() + + # Apply causal mask to index_scores before computing topk + causal_mask = torch.triu( + torch.full( + (seqlen, seqlen), float('-inf'), dtype=torch.float32, device=index_scores.device + ), + diagonal=1, + ) + # [batch_size, seqlen, seqlen] + [seqlen, seqlen] -> [batch_size, seqlen, seqlen] + masked_index_scores = index_scores + causal_mask + + # Get topk indices from masked index_scores + topk_k = min(index_topk, seqlen) + topk_indices = masked_index_scores.topk(topk_k, dim=-1)[1] + + query = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + softmax_scale = head_dim**-0.5 + + loss = compute_dsa_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + query=query, + key=key, + softmax_scale=softmax_scale, + loss_coeff=1.0, + sparse_loss=False, + pg_collection=self.pg_collection, + ) + + assert loss.shape == torch.Size([]) + assert loss.dtype == torch.float32 + assert loss >= 0 # KL divergence should be non-negative + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_loss_sparse(self, seqlen_and_topk): + """Test sparse indexer loss computation.""" + batch_size = 2 + seqlen = seqlen_and_topk[0] + num_heads = 4 + head_dim = 128 + index_topk = seqlen_and_topk[1] + + # Create dummy index scores + index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() + + # Apply causal mask to index_scores before computing topk + causal_mask = torch.triu( + torch.full( + (seqlen, seqlen), float('-inf'), dtype=torch.float32, device=index_scores.device + ), + diagonal=1, + ) + # [batch_size, seqlen, seqlen] + [seqlen, seqlen] -> [batch_size, seqlen, seqlen] + masked_index_scores = index_scores + causal_mask + + # Get topk indices from masked index_scores + topk_k = min(index_topk, seqlen) + topk_indices = masked_index_scores.topk(topk_k, dim=-1)[1] + + query = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + softmax_scale = head_dim**-0.5 + + loss_sparse = compute_dsa_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + query=query, + key=key, + softmax_scale=softmax_scale, + loss_coeff=1.0, + sparse_loss=True, + pg_collection=self.pg_collection, + ) + + loss_dense = compute_dsa_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + query=query, + key=key, + softmax_scale=softmax_scale, + loss_coeff=1.0, + sparse_loss=False, + pg_collection=self.pg_collection, + ) + + # Sparse loss should be different from dense loss + if seqlen > index_topk: + assert loss_sparse != loss_dense + else: + assert loss_sparse == loss_dense + assert loss_sparse >= 0 + assert loss_dense >= 0 + + +class TestDSAIndexerLossAutoScaler: + """Test DSAIndexerLossAutoScaler autograd function.""" + + @pytest.fixture(scope='function', 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_forward_pass(self): + """Test that forward pass preserves output.""" + output = torch.randn(16, 2, 128).cuda() + output.requires_grad_(True) + indexer_loss = torch.tensor(0.5).cuda() + indexer_loss.requires_grad_(True) + + result = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + assert torch.allclose(result, output, atol=0, rtol=0) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward_pass(self): + """Test that backward pass triggers indexer loss backward and scales gradient correctly.""" + output = torch.randn(16, 2, 128).cuda() + output.requires_grad_(True) + + # Create indexer_loss with computation graph + # This simulates compute_dsa_indexer_loss which computes KL divergence + dummy_input = torch.randn(10).cuda() + dummy_input.requires_grad_(True) + indexer_loss = dummy_input.mean() + + # Set loss scale + scale = torch.tensor(2.0).cuda() + DSAIndexerLossAutoScaler.set_loss_scale(scale) + + # Apply the autograd function + result = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + # Trigger backward + main_loss = result.sum() + main_loss.backward() + + # Check that gradients flow back to output + assert output.grad is not None, "Gradient should flow back to parameters" + + # Check that indexer_loss backward was triggered + assert dummy_input.grad is not None, "Indexer loss backward should be triggered" + + # Verify the gradient is scaled correctly + expected_grad_per_element = scale.item() / len(dummy_input) + assert torch.allclose( + dummy_input.grad, + torch.full_like(dummy_input, expected_grad_per_element), + rtol=0, + atol=0, + ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}" + + +@pytest.mark.parametrize("seqlen", [16, 64]) +class TestDSAIndexer: + """Test DSA Indexer module basic functionality with TP=1.""" + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + # Create MLA config with sparse attention parameters + self.index_topk = 32 + self.config = MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + # MLA specific configs + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=64, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + # Sparse attention specific configs + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=self.index_topk, + ) + + # Create indexer submodules spec + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + indexer_submodules = DSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + self.indexer = DSAIndexer(self.config, indexer_submodules, self.pg_collection) + + yield + Utils.destroy_model_parallel() + + def test_dsa_indexer_constructor(self, seqlen): + """Test indexer initialization.""" + assert isinstance(self.indexer, DSAIndexer) + assert self.indexer.hidden_size == 256 + assert self.indexer.index_n_heads == 8 + assert self.indexer.index_head_dim == 64 + assert self.indexer.index_topk == 32 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_forward(self, seqlen): + """Test indexer forward pass.""" + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + + # Forward pass + topk_indices = self.indexer(x, qr) + + # Check output shape + assert topk_indices.shape == (batch_size, seqlen, min(self.config.dsa_indexer_topk, seqlen)) + assert topk_indices.dtype == torch.long + assert torch.all((topk_indices >= 0) & (topk_indices < seqlen)) + # Make sure no duplicate indices are selected + assert torch.all( + torch.sort(topk_indices, dim=-1).values[:, :, 1:] + != torch.sort(topk_indices, dim=-1).values[:, :, :-1] + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_forward_with_scores(self, seqlen): + """Test indexer forward pass with scores.""" + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + + # Forward pass with scores + index_scores, topk_indices = self.indexer.forward_with_scores(x, qr) + + # Check output shapes + assert index_scores.shape == (batch_size, seqlen, seqlen) + assert topk_indices.shape == (batch_size, seqlen, min(self.config.dsa_indexer_topk, seqlen)) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + assert torch.all((topk_indices >= 0) & (topk_indices < seqlen)) + # Make sure no duplicate indices are selected + assert torch.all( + torch.sort(topk_indices, dim=-1).values[:, :, 1:] + != torch.sort(topk_indices, dim=-1).values[:, :, :-1] + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_with_mask(self, seqlen): + """Test indexer with attention mask.""" + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + mask = torch.triu( + torch.full((batch_size, seqlen, seqlen), float('-inf'), dtype=torch.float32).cuda(), + diagonal=1, + ) + + # Forward pass with mask + index_scores, topk_indices = self.indexer.forward_with_scores(x, qr, mask=mask) + + # Check that masked positions are not selected + # For causal mask, topk_indices[b, i, :] should all be <= i (except for the case that + # i < index_topk). + for b in range(batch_size): + for i in range(seqlen): + assert torch.all(topk_indices[b, i] <= max(self.index_topk, i)) + + +class TestDSAttention: + """Test DSAttention module basic functionality with TP=1.""" + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + # Create MLA config with sparse attention parameters + self.config = MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + # MLA specific configs + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=64, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + # Sparse attention specific configs + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=False, + ) + + # Create sparse attention submodules spec + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + indexer_submodules = DSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + indexer_spec = ModuleSpec(module=DSAIndexer, submodules=indexer_submodules) + sparse_attention_submodules = DSAttentionSubmodules(indexer=indexer_spec) + + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + self.sparse_attention = DSAttention( + config=self.config, + submodules=sparse_attention_submodules, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_dsa_constructor(self): + """Test sparse attention initialization.""" + assert isinstance(self.sparse_attention, DSAttention) + assert hasattr(self.sparse_attention, 'indexer') + assert isinstance(self.sparse_attention.indexer, DSAIndexer) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_forward(self): + """Test sparse attention forward pass.""" + seq_len = 16 + batch_size = 2 + num_heads = self.config.num_attention_heads + head_dim = self.config.hidden_size // num_heads + + self.sparse_attention.cuda() + + # Create input tensors [seq_len, batch, num_heads, head_dim] + query = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + value = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + # Original hidden states and low-rank query + 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() + + # Create causal attention mask + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + # Forward pass + output = self.sparse_attention( + query=query, + key=key, + value=value, + x=x, + qr=qr, + attention_mask=attention_mask, + attn_mask_type=AttnMaskType.causal, + ) + + # Check output shape + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_backward(self): + """Test sparse attention backward pass with indexer loss.""" + seq_len = 16 + batch_size = 2 + num_heads = self.config.num_attention_heads + head_dim = self.config.hidden_size // num_heads + + self.sparse_attention.train() + self.sparse_attention.cuda() + + # Create input tensors + query = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + value = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + # Original hidden states and low-rank query + 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() + + # Create causal attention mask + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + # Forward pass + output = self.sparse_attention( + query=query, + key=key, + value=value, + x=x, + qr=qr, + attention_mask=attention_mask, + attn_mask_type=AttnMaskType.causal, + ) + + # Backward pass + loss = output.sum() + loss.backward() + + # Check that gradients are computed for inputs + assert query.grad is not None + assert key.grad is not None + assert value.grad is not None + + # Check that indexer parameters have gradients + for name, param in self.sparse_attention.indexer.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Indexer parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_topk_selection(self): + """Test that sparse attention correctly selects top-k indices.""" + seq_len = 16 + batch_size = 2 + num_heads = self.config.num_attention_heads + head_dim = self.config.hidden_size // num_heads + + self.sparse_attention.eval() + self.sparse_attention.cuda() + + # Create input tensors + query = torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + value = torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + + # Original hidden states and low-rank query + 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() + + # Create causal attention mask + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + with torch.no_grad(): + # Get topk indices from indexer + _, topk_indices = self.sparse_attention.indexer.forward_with_scores(x, qr) + + # Forward pass + output = self.sparse_attention( + query=query, + key=key, + value=value, + x=x, + qr=qr, + attention_mask=attention_mask, + attn_mask_type=AttnMaskType.causal, + ) + + # Check that topk_indices are valid + assert torch.all(topk_indices >= 0) + assert torch.all(topk_indices < seq_len) + assert topk_indices.shape[2] == min(self.config.dsa_indexer_topk, seq_len) + + +# ====================================================================================== +# Tensor Parallel Consistency Tests +# ====================================================================================== + + +@pytest.mark.parametrize("tensor_model_parallel_size", [2, 4, 8]) +@pytest.mark.parametrize("sequence_parallel", [False, True]) +class TestIndexerTensorParallel: + """Test DSA Indexer with different TP sizes and SP settings, compare with TP=1 baseline.""" + + def _create_config(self, sequence_parallel=False): + """Helper to create MLA config.""" + # Get TP size from parallel_state + tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size() + + return MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + # MLA specific configs + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=64, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + # Sparse attention specific configs + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + ) + + def _create_indexer(self, config, pg_collection): + """Helper to create indexer.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + indexer_submodules = DSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + + return DSAIndexer(config, indexer_submodules, pg_collection) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_weight_consistency(self, tensor_model_parallel_size, sequence_parallel): + """Test that indexer weights are identical across ALL GPUs.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config(sequence_parallel=sequence_parallel) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + indexer = self._create_indexer(config, pg_collection).cuda() + + # Check that all weights are identical across ALL ranks (not just TP group) + world_size = torch.distributed.get_world_size() + world_rank = torch.distributed.get_rank() + + if world_size > 1: + for name, param in indexer.named_parameters(): + # Gather weights from ALL ranks in WORLD group + param_list = [torch.zeros_like(param.data) for _ in range(world_size)] + torch.distributed.all_gather(param_list, param.data) + + # All weights should be identical across all GPUs + for i in range(1, world_size): + assert torch.allclose( + param_list[0], param_list[i], rtol=0, atol=0 + ), f"Parameter {name} differs between rank 0 and rank {i} (world)" + + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_forward_consistency(self, tensor_model_parallel_size, sequence_parallel): + """Test that indexer gives consistent results across different TP sizes and SP settings.""" + # First run with TP=1 to get baseline + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config_tp1 = self._create_config(sequence_parallel=False) # TP=1 doesn't use SP + pg_collection_tp1 = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + indexer_tp1 = self._create_indexer(config_tp1, pg_collection_tp1).cuda() + + seq_len = 64 + batch_size = 2 + + # Create one common input (all ranks create same input with same seed) + x_input = torch.randn( + seq_len, batch_size, config_tp1.hidden_size, dtype=torch.bfloat16 + ).cuda() + qr_input = torch.randn( + seq_len, batch_size, config_tp1.q_lora_rank, dtype=torch.bfloat16 + ).cuda() + + # Forward pass with gradients enabled + index_scores_tp1, topk_indices_tp1 = indexer_tp1.forward_with_scores(x_input, qr_input) + + # Backward pass + loss_tp1 = index_scores_tp1.sum() + loss_tp1.backward() + + # Save gradients from TP=1 + indexer_tp1_grads = { + name: param.grad.clone().cpu() + for name, param in indexer_tp1.named_parameters() + if param.grad is not None + } + + Utils.destroy_model_parallel() + + # Now run with target TP size + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config_tpn = self._create_config(sequence_parallel=sequence_parallel) + pg_collection_tpn = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + indexer_tpn = self._create_indexer(config_tpn, pg_collection_tpn).cuda() + + # Prepare input: split along seqlen if SP is enabled + if sequence_parallel: + tp_rank = parallel_state.get_tensor_model_parallel_rank() + seq_per_rank = seq_len // tensor_model_parallel_size + start_idx = tp_rank * seq_per_rank + end_idx = (tp_rank + 1) * seq_per_rank + x_tpn = x_input[start_idx:end_idx] + qr_tpn = qr_input[start_idx:end_idx] + else: + # No SP: all TP ranks see full input + x_tpn = x_input + qr_tpn = qr_input + + # Forward pass with gradients enabled + index_scores_tpn, topk_indices_tpn = indexer_tpn.forward_with_scores(x_tpn, qr_tpn) + + # Backward pass + loss_tpn = index_scores_tpn.sum() + loss_tpn.backward() + + # Compare forward outputs + assert index_scores_tpn.shape == index_scores_tp1.shape + assert topk_indices_tpn.shape == topk_indices_tp1.shape + + # Check that index scores are close (allow for floating point accumulation errors) + assert torch.allclose( + index_scores_tpn, index_scores_tp1, rtol=0, atol=0 + ), f"Index scores mismatch between TP=1 and TP={tensor_model_parallel_size}, SP={sequence_parallel}" + + # Check that topk indices are exactly the same + assert torch.equal( + topk_indices_tpn, topk_indices_tp1 + ), f"Top-k indices mismatch between TP=1 and TP={tensor_model_parallel_size}, SP={sequence_parallel}" + + # Compare gradients - indexer grads should be identical (duplicated weights) + for name, param in indexer_tpn.named_parameters(): + if param.grad is not None and name in indexer_tp1_grads: + assert torch.allclose( + param.grad.cpu(), indexer_tp1_grads[name], rtol=0, atol=0 + ), f"Indexer gradient {name} mismatch between TP=1 and TP={tensor_model_parallel_size}, SP={sequence_parallel}" + + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_gradient_sync(self, tensor_model_parallel_size, sequence_parallel): + """Test that gradients are properly synchronized within TP group.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config(sequence_parallel=sequence_parallel) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + indexer = self._create_indexer(config, pg_collection).cuda() + + seq_len = 64 + batch_size = 2 + + # Create one common input (all ranks create same input with same seed) + x_input = torch.randn(seq_len, batch_size, config.hidden_size, dtype=torch.bfloat16).cuda() + qr_input = torch.randn(seq_len, batch_size, config.q_lora_rank, dtype=torch.bfloat16).cuda() + + # Prepare input: split along seqlen if SP is enabled + if sequence_parallel: + tp_rank = parallel_state.get_tensor_model_parallel_rank() + tp_size = parallel_state.get_tensor_model_parallel_world_size() + seq_per_rank = seq_len // tp_size + start_idx = tp_rank * seq_per_rank + end_idx = (tp_rank + 1) * seq_per_rank + x = x_input[start_idx:end_idx] + qr = qr_input[start_idx:end_idx] + else: + # No SP: all TP ranks see full input + x = x_input + qr = qr_input + + # Forward and backward + index_scores, topk_indices = indexer.forward_with_scores(x, qr) + loss = index_scores.sum() + loss.backward() + + # Check that all parameters have gradients + for name, param in indexer.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + # After TP sync, check that gradients are identical within TP group + # Note: We only check TP group because DDP sync happens separately + tp_size = parallel_state.get_tensor_model_parallel_world_size() + if tp_size > 1: + for name, param in indexer.named_parameters(): + if param.requires_grad and param.grad is not None: + # Gather gradients from all ranks in TP group only + grad_list = [torch.zeros_like(param.grad) for _ in range(tp_size)] + torch.distributed.all_gather(grad_list, param.grad, group=pg_collection.tp) + + # All gradients should be identical within TP group after sync + for i in range(1, tp_size): + assert torch.allclose( + grad_list[0], grad_list[i], rtol=0, atol=0 + ), f"Gradient for {name} differs between TP rank 0 and rank {i} after TP sync" + + Utils.destroy_model_parallel() + + +@pytest.mark.parametrize("tensor_model_parallel_size", [2, 4]) +@pytest.mark.parametrize("sequence_parallel", [False, True]) +@pytest.mark.parametrize("use_sparse_indexer_loss", [False, True]) +class TestDSAttentionTensorParallel: + """Test DSAttention with different TP sizes, SP settings, and sparse indexer loss.""" + + def _create_config(self, sequence_parallel=False, use_sparse_indexer_loss=False): + """Helper to create MLA config.""" + # Get TP size from parallel_state + tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size() + + return MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + # MLA specific configs + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=64, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + # Sparse attention specific configs + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=32, + dsa_indexer_loss_coeff=1.0, + dsa_indexer_use_sparse_loss=use_sparse_indexer_loss, + ) + + def _create_sparse_attention(self, config, pg_collection): + """Helper to create sparse attention.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + indexer_submodules = DSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + indexer_spec = ModuleSpec(module=DSAIndexer, submodules=indexer_submodules) + sparse_attention_submodules = DSAttentionSubmodules(indexer=indexer_spec) + + return DSAttention( + config=config, + submodules=sparse_attention_submodules, + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=pg_collection, + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_weight_consistency( + self, tensor_model_parallel_size, sequence_parallel, use_sparse_indexer_loss + ): + """Test that sparse attention indexer weights are identical across ALL GPUs.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config( + sequence_parallel=sequence_parallel, use_sparse_indexer_loss=use_sparse_indexer_loss + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() + + # Check that all indexer weights are identical across ALL ranks + world_size = torch.distributed.get_world_size() + world_rank = torch.distributed.get_rank() + + if world_size > 1: + for name, param in sparse_attention.indexer.named_parameters(): + # Gather weights from ALL ranks in WORLD group + param_list = [torch.zeros_like(param.data) for _ in range(world_size)] + torch.distributed.all_gather(param_list, param.data) + + # All weights should be identical across all GPUs + for i in range(1, world_size): + torch.testing.assert_close(param_list[0], param_list[i], rtol=0, atol=0) + + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_forward_consistency( + self, tensor_model_parallel_size, sequence_parallel, use_sparse_indexer_loss + ): + """Test that sparse attention gives consistent results across different TP, SP, and sparse loss settings.""" + # First run with TP=1 to get baseline + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config_tp1 = self._create_config( + sequence_parallel=False, use_sparse_indexer_loss=use_sparse_indexer_loss + ) # TP=1 doesn't use SP + pg_collection_tp1 = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + sparse_attention_tp1 = self._create_sparse_attention(config_tp1, pg_collection_tp1).cuda() + + seq_len = 64 + batch_size = 2 + num_heads = config_tp1.num_attention_heads + head_dim = config_tp1.hidden_size // num_heads + + # Create one common input (all ranks create same input with same seed) + query_input = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + key_input = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + value_input = ( + torch.randn(seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + x_input = torch.randn( + seq_len, batch_size, config_tp1.hidden_size, dtype=torch.bfloat16 + ).cuda() + qr_input = torch.randn( + seq_len, batch_size, config_tp1.q_lora_rank, dtype=torch.bfloat16 + ).cuda() + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + # Forward pass with gradients enabled + sparse_attention_tp1.train() + output_tp1 = sparse_attention_tp1( + query=query_input, + key=key_input, + value=value_input, + x=x_input, + qr=qr_input, + attention_mask=attention_mask, + attn_mask_type=AttnMaskType.causal, + ) + + # Backward pass + loss_tp1 = output_tp1.sum() + loss_tp1.backward() + + # Save gradients from TP=1 + indexer_tp1_grads = { + name: param.grad.clone() + for name, param in sparse_attention_tp1.indexer.named_parameters() + if param.grad is not None + } + query_tp1_grad = query_input.grad.clone().cpu() + key_tp1_grad = key_input.grad.clone().cpu() + value_tp1_grad = value_input.grad.clone().cpu() + + Utils.destroy_model_parallel() + + # Now run with target TP size + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config_tpn = self._create_config( + sequence_parallel=sequence_parallel, use_sparse_indexer_loss=use_sparse_indexer_loss + ) + pg_collection_tpn = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + sparse_attention_tpn = self._create_sparse_attention(config_tpn, pg_collection_tpn).cuda() + + # Create one common input (all ranks create same input with same seed) + query_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + key_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + value_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + x_input = torch.randn( + seq_len, batch_size, config_tp1.hidden_size, dtype=torch.bfloat16 + ).cuda() + qr_input = torch.randn( + seq_len, batch_size, config_tp1.q_lora_rank, dtype=torch.bfloat16 + ).cuda() + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + # Prepare input: split along seqlen if SP is enabled + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if sequence_parallel: + seq_per_rank = seq_len // tensor_model_parallel_size + start_idx = tp_rank * seq_per_rank + end_idx = (tp_rank + 1) * seq_per_rank + x_tpn = x_input[start_idx:end_idx] + qr_tpn = qr_input[start_idx:end_idx] + else: + x_tpn = x_input + qr_tpn = qr_input + + query_input = query_input.detach() + key_input = key_input.detach() + value_input = value_input.detach() + head_per_rank = num_heads // tensor_model_parallel_size + start_head = tp_rank * head_per_rank + end_head = (tp_rank + 1) * head_per_rank + query_tpn = query_input[:, :, start_head:end_head, :].clone().requires_grad_(True) + key_tpn = key_input[:, :, start_head:end_head, :].clone().requires_grad_(True) + value_tpn = value_input[:, :, start_head:end_head, :].clone().requires_grad_(True) + attention_mask_tpn = attention_mask + + # Forward pass with gradients enabled + sparse_attention_tpn.train() + output_tpn = sparse_attention_tpn( + query=query_tpn, + key=key_tpn, + value=value_tpn, + x=x_tpn, + qr=qr_tpn, + attention_mask=attention_mask_tpn, + attn_mask_type=AttnMaskType.causal, + ) + + # Backward pass + loss_tpn = output_tpn.sum() + loss_tpn.backward() + + from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region + + output_tpn_gathered = gather_from_tensor_model_parallel_region( + output_tpn, group=pg_collection_tpn.tp + ) + assert output_tpn_gathered.shape == output_tp1.shape + assert torch.allclose( + output_tpn_gathered.detach(), output_tp1.detach(), rtol=0, atol=0 + ), f"Sparse attention outputs mismatch between TP=1 and TP={tensor_model_parallel_size}, SP={sequence_parallel}, sparse_loss={use_sparse_indexer_loss}" + + # 1. Check indexer gradients. + for name, param in sparse_attention_tpn.indexer.named_parameters(): + if param.grad is not None and name in indexer_tp1_grads: + torch.testing.assert_close( + param.grad, indexer_tp1_grads[name], rtol=1e-5, atol=1e-5 + ) + + # 2. Query/Key/Value gradients need to be gathered along num_heads dim (dim 2) if SP is enabled + # Flatten last two dims: [seq_len, batch, num_heads, head_dim] -> [seq_len, batch, num_heads * head_dim] + sq, b, nh, hd = query_tpn.grad.shape + query_grad_flat = query_tpn.grad.reshape(sq, b, nh * hd) + key_grad_flat = key_tpn.grad.reshape(sq, b, nh * hd) + value_grad_flat = value_tpn.grad.reshape(sq, b, nh * hd) + + # Gather along last dim + query_grad_gathered_flat = gather_from_tensor_model_parallel_region( + query_grad_flat, group=pg_collection_tpn.tp + ) + key_grad_gathered_flat = gather_from_tensor_model_parallel_region( + key_grad_flat, group=pg_collection_tpn.tp + ) + value_grad_gathered_flat = gather_from_tensor_model_parallel_region( + value_grad_flat, group=pg_collection_tpn.tp + ) + + # Reshape back: [seq_len, batch, num_heads * head_dim] -> [seq_len, batch, num_heads, head_dim] + query_tpn_grad_gathered = query_grad_gathered_flat.reshape(sq, b, num_heads, hd) + key_tpn_grad_gathered = key_grad_gathered_flat.reshape(sq, b, num_heads, hd) + value_tpn_grad_gathered = value_grad_gathered_flat.reshape(sq, b, num_heads, hd) + + assert torch.allclose( + query_tpn_grad_gathered.cpu(), query_tp1_grad, rtol=0, atol=0 + ), f"Query gradient mismatch between TP=1 and TP={tensor_model_parallel_size}" + assert torch.allclose( + key_tpn_grad_gathered.cpu(), key_tp1_grad, rtol=0, atol=0 + ), f"Key gradient mismatch between TP=1 and TP={tensor_model_parallel_size}" + assert torch.allclose( + value_tpn_grad_gathered.cpu(), value_tp1_grad, rtol=0, atol=0 + ), f"Value gradient mismatch between TP=1 and TP={tensor_model_parallel_size}" + + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_gradient_sync( + self, tensor_model_parallel_size, sequence_parallel, use_sparse_indexer_loss + ): + """Test that indexer gradients are properly synchronized within TP group.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tensor_model_parallel_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config( + sequence_parallel=sequence_parallel, use_sparse_indexer_loss=use_sparse_indexer_loss + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() + sparse_attention.train() + + seq_len = 64 + batch_size = 2 + num_heads = config.num_attention_heads + head_dim = config.hidden_size // num_heads + + # Create one common input (all ranks create same input with same seed) + query_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + key_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + value_input = torch.randn( + seq_len, batch_size, num_heads, head_dim, dtype=torch.bfloat16 + ).cuda() + x_input = torch.randn(seq_len, batch_size, config.hidden_size, dtype=torch.bfloat16).cuda() + qr_input = torch.randn(seq_len, batch_size, config.q_lora_rank, dtype=torch.bfloat16).cuda() + + # Prepare input: split along seqlen if SP is enabled + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if sequence_parallel: + tp_size = parallel_state.get_tensor_model_parallel_world_size() + seq_per_rank = seq_len // tp_size + start_idx = tp_rank * seq_per_rank + end_idx = (tp_rank + 1) * seq_per_rank + x = x_input[start_idx:end_idx] + qr = qr_input[start_idx:end_idx] + else: + x = x_input + qr = qr_input + + # query, key, value should be split along num_heads dim + head_per_rank = num_heads // tensor_model_parallel_size + start_head = tp_rank * head_per_rank + end_head = (tp_rank + 1) * head_per_rank + query = query_input[:, :, start_head:end_head, :] + key = key_input[:, :, start_head:end_head, :] + value = value_input[:, :, start_head:end_head, :] + + attention_mask = torch.ones(batch_size, 1, seq_len, seq_len, dtype=torch.bool).cuda() + attention_mask = torch.tril(attention_mask) + + query.requires_grad_(True) + key.requires_grad_(True) + value.requires_grad_(True) + + # Forward and backward + output = sparse_attention( + query=query, + key=key, + value=value, + x=x, + qr=qr, + attention_mask=attention_mask, + attn_mask_type=AttnMaskType.causal, + ) + + loss = output.sum() + loss.backward() + + # Check that gradients exist before sync + assert query.grad is not None + assert key.grad is not None + assert value.grad is not None + + # Check that indexer parameters have gradients + for name, param in sparse_attention.indexer.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Indexer parameter {name} has no gradient" + + # Check that indexer gradients are identical within TP group + tp_size = parallel_state.get_tensor_model_parallel_world_size() + if tp_size > 1: + for name, param in sparse_attention.indexer.named_parameters(): + if param.requires_grad and param.grad is not None: + # Gather gradients from all ranks in TP group only + grad_list = [torch.zeros_like(param.grad) for _ in range(tp_size)] + torch.distributed.all_gather(grad_list, param.grad, group=pg_collection.tp) + + # All gradients should be identical within TP group after sync + for i in range(1, tp_size): + assert torch.allclose( + grad_list[0], grad_list[i], rtol=0, atol=0 + ), f"Indexer gradient for {name} differs between TP rank 0 and rank {i} after TP sync" + + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/transformer/test_multi_latent_attention.py b/tests/unit_tests/transformer/test_multi_latent_attention.py index fb259671287..bc8514ee561 100644 --- a/tests/unit_tests/transformer/test_multi_latent_attention.py +++ b/tests/unit_tests/transformer/test_multi_latent_attention.py @@ -293,13 +293,17 @@ def test_gpu_forward_thd_padded(self): assert bias.shape[0] == config.hidden_size # Test that the get_query_key_value_tensors function properly handles padded cu_seqlens - query, key, value = self.parallel_attention.get_query_key_value_tensors( - hidden_states, None, None, packed_seq_params, None + query, key, value, q_compressed, kv_compressed = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states, None, None, packed_seq_params, None + ) ) assert query is not None assert key is not None assert value is not None + assert q_compressed is not None + assert kv_compressed is not None assert query.is_contiguous() assert key.is_contiguous() assert value.is_contiguous() @@ -370,7 +374,9 @@ def test_up_proj_recomputed_gpu_forward(self): ) hidden_states = hidden_states.cuda() - q, k, v = checkpointed_parallel_attention.get_query_key_value_tensors(hidden_states) + q, k, v, q_compressed, kv_compressed = ( + checkpointed_parallel_attention.get_query_key_value_tensors(hidden_states) + ) assert q.is_contiguous() assert k.is_contiguous() assert v.is_contiguous() @@ -675,18 +681,30 @@ def test_gpu_forward_thd_precision(self): packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) # fine-grained check - query_sbhd, key_sbhd, value_sbhd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_sbhd, None, None, None, None + query_sbhd, key_sbhd, value_sbhd, q_compressed_sbhd, kv_compressed_sbhd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_sbhd, None, None, None, None + ) ) - query_thd, key_thd, value_thd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_thd, None, None, packed_seq_params, None + query_thd, key_thd, value_thd, q_compressed_thd, kv_compressed_thd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_thd, None, None, packed_seq_params, None + ) ) _query_sbhd = query_sbhd.transpose(0, 1).contiguous().view(*query_thd.shape) _key_sbhd = key_sbhd.transpose(0, 1).contiguous().view(*key_thd.shape) _value_sbhd = value_sbhd.transpose(0, 1).contiguous().view(*value_thd.shape) + _q_compressed_sbhd = ( + q_compressed_sbhd.transpose(0, 1).contiguous().view(*q_compressed_thd.shape) + ) + _kv_compressed_sbhd = ( + kv_compressed_sbhd.transpose(0, 1).contiguous().view(*kv_compressed_thd.shape) + ) assert torch.equal(_query_sbhd, query_thd) assert torch.equal(_key_sbhd, key_thd) assert torch.equal(_value_sbhd, value_thd) + assert torch.equal(_q_compressed_sbhd, q_compressed_thd) + assert torch.equal(_kv_compressed_sbhd, kv_compressed_thd) core_attn_out_sbhd = self.parallel_attention.core_attention( query_sbhd, @@ -828,18 +846,30 @@ def test_gpu_forward_thd_precision(self): packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) # fine-grained check - query_sbhd, key_sbhd, value_sbhd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_sbhd, None, None, None, None + query_sbhd, key_sbhd, value_sbhd, q_compressed_sbhd, kv_compressed_sbhd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_sbhd, None, None, None, None + ) ) - query_thd, key_thd, value_thd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_thd, None, None, packed_seq_params, None + query_thd, key_thd, value_thd, q_compressed_thd, kv_compressed_thd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_thd, None, None, packed_seq_params, None + ) ) _query_sbhd = query_sbhd.transpose(0, 1).contiguous().view(*query_thd.shape) _key_sbhd = key_sbhd.transpose(0, 1).contiguous().view(*key_thd.shape) _value_sbhd = value_sbhd.transpose(0, 1).contiguous().view(*value_thd.shape) + _q_compressed_sbhd = ( + q_compressed_sbhd.transpose(0, 1).contiguous().view(*q_compressed_thd.shape) + ) + _kv_compressed_sbhd = ( + kv_compressed_sbhd.transpose(0, 1).contiguous().view(*kv_compressed_thd.shape) + ) torch.testing.assert_close(_query_sbhd, query_thd, atol=1e-6, rtol=1e-6) torch.testing.assert_close(_key_sbhd, key_thd, atol=1e-6, rtol=1e-6) torch.testing.assert_close(_value_sbhd, value_thd, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(_q_compressed_sbhd, q_compressed_thd, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(_kv_compressed_sbhd, kv_compressed_thd, atol=1e-6, rtol=1e-6) core_attn_out_sbhd = self.parallel_attention.core_attention( query_sbhd, @@ -967,18 +997,30 @@ def test_gpu_forward_thd_precision(self): packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) # fine-grained check - query_sbhd, key_sbhd, value_sbhd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_sbhd, None, None, None, None + query_sbhd, key_sbhd, value_sbhd, q_compressed_sbhd, kv_compressed_sbhd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_sbhd, None, None, None, None + ) ) - query_thd, key_thd, value_thd = self.parallel_attention.get_query_key_value_tensors( - hidden_states_thd, None, None, packed_seq_params, None + query_thd, key_thd, value_thd, q_compressed_thd, kv_compressed_thd = ( + self.parallel_attention.get_query_key_value_tensors( + hidden_states_thd, None, None, packed_seq_params, None + ) ) _query_sbhd = query_sbhd.transpose(0, 1).contiguous().view(*query_thd.shape) _key_sbhd = key_sbhd.transpose(0, 1).contiguous().view(*key_thd.shape) _value_sbhd = value_sbhd.transpose(0, 1).contiguous().view(*value_thd.shape) + _q_compressed_sbhd = ( + q_compressed_sbhd.transpose(0, 1).contiguous().view(*q_compressed_thd.shape) + ) + _kv_compressed_sbhd = ( + kv_compressed_sbhd.transpose(0, 1).contiguous().view(*kv_compressed_thd.shape) + ) assert torch.equal(_query_sbhd, query_thd) assert torch.equal(_key_sbhd, key_thd) assert torch.equal(_value_sbhd, value_thd) + assert torch.equal(_q_compressed_sbhd, q_compressed_thd) + assert torch.equal(_kv_compressed_sbhd, kv_compressed_thd) core_attn_out_sbhd = self.parallel_attention.core_attention( query_sbhd,