From 2c7f33eb70270290ffd92f9da0102e93bf8ab209 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 6 Nov 2025 10:57:38 +0800 Subject: [PATCH 01/28] Initial draft version Signed-off-by: kunlunl --- gpt_builders.py | 2 + megatron/core/models/gpt/gpt_layer_specs.py | 30 +- .../gpt/sparse_attention_module_specs.py | 62 ++ megatron/core/transformer/attention.py | 2 + .../transformer/multi_latent_attention.py | 1 + megatron/core/transformer/sparse_attention.py | 662 ++++++++++++++++++ .../core/transformer/transformer_config.py | 18 + megatron/training/arguments.py | 15 + 8 files changed, 790 insertions(+), 2 deletions(-) create mode 100644 megatron/core/models/gpt/sparse_attention_module_specs.py create mode 100644 megatron/core/transformer/sparse_attention.py diff --git a/gpt_builders.py b/gpt_builders.py index 591f74bb20c..6af8f64e116 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -112,6 +112,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, + args.sparse_attention_type, args.linear_attention_type, moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, qk_l2_norm=args.qk_l2_norm, @@ -123,6 +124,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, + args.sparse_attention_type, args.linear_attention_type, moe_use_legacy_grouped_gemm=args.moe_use_legacy_grouped_gemm, normalization=args.normalization, diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index e3ef7f20141..dfa7d148a09 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -9,6 +9,9 @@ get_linear_attention_module_spec_for_backend, ) from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend +from megatron.core.models.gpt.sparse_attention_module_specs import ( + get_sparse_attention_module_spec_for_backend, +) from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules from megatron.core.transformer.enums import AttnMaskType, LayerType from megatron.core.transformer.identity_op import IdentityOp @@ -77,6 +80,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, + sparse_attention_type: Optional[str] = None, linear_attention_type: Optional[str] = None, fp8: Optional[str] = None, # pylint: disable=unused-argument moe_use_legacy_grouped_gemm: Optional[bool] = False, @@ -94,6 +98,7 @@ def get_gpt_layer_with_transformer_engine_spec( 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 multi-latent attention. Defaults to False. + sparse_attention_type (str, optional): The type of sparse attention. Defaults to None. linear_attention_type (str, optional): The type of linear attention. Defaults to None. fp8 (str, optional): Deprecated. For temporary Nemo compatibility. moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP. @@ -129,6 +134,7 @@ def get_gpt_layer_with_transformer_engine_spec( attention = get_attention_module_spec_for_backend( backend=backend, sharded_state_dict_keys_map=sharded_state_dict_keys_map, + sparse_attention_type=sparse_attention_type, linear_attention_type=linear_attention_type, qk_layernorm=qk_layernorm, qk_l2_norm=qk_l2_norm, @@ -161,6 +167,7 @@ def get_gpt_layer_local_spec( moe_grouped_gemm: Optional[bool] = False, qk_layernorm: Optional[bool] = False, multi_latent_attention: Optional[bool] = False, + sparse_attention_type: Optional[str] = None, linear_attention_type: Optional[str] = None, fp8: Optional[str] = None, # pylint: disable=unused-argument moe_use_legacy_grouped_gemm: Optional[bool] = False, @@ -176,6 +183,7 @@ def get_gpt_layer_local_spec( 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 multi-latent attention. Defaults to False. + sparse_attention_type (str, optional): The type of sparse attention. Defaults to None. linear_attention_type (str, optional): The type of linear attention. Defaults to None. fp8 (str, optional): Deprecated. For temporary Nemo compatibility. moe_use_legacy_grouped_gemm (bool, optional): Force use the legacy GroupedMLP. @@ -200,6 +208,9 @@ def get_gpt_layer_local_spec( " and will be removed soon. Please update your code accordingly." ) + if sparse_attention_type is not None: + raise NotImplementedError("Sparse attention is not supported with local spec yet.") + if linear_attention_type is not None: raise NotImplementedError("Linear attention is not supported with local spec yet.") @@ -208,6 +219,7 @@ def get_gpt_layer_local_spec( attention = get_attention_module_spec_for_backend( backend=backend, sharded_state_dict_keys_map=sharded_state_dict_keys_map, + sparse_attention_type=sparse_attention_type, linear_attention_type=linear_attention_type, qk_layernorm=qk_layernorm, qk_l2_norm=qk_l2_norm, @@ -272,6 +284,7 @@ def get_transformer_layer_spec_for_backend( def get_attention_module_spec_for_backend( backend: BackendSpecProvider, sharded_state_dict_keys_map: dict, + sparse_attention_type: Optional[str] = None, linear_attention_type: Optional[str] = None, qk_layernorm: Optional[bool] = False, qk_l2_norm: Optional[bool] = False, @@ -292,6 +305,16 @@ def get_attention_module_spec_for_backend( rms_norm = normalization == "RMSNorm" qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) + if sparse_attention_type is not None: + assert multi_latent_attention, "Currently only MLA supports sparse attention." + core_attention = get_sparse_attention_module_spec_for_backend( + backend=backend, + sparse_attention_type=sparse_attention_type, + normalization=normalization, + ) + else: + core_attention = backend.core_attention() + if multi_latent_attention: assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." linear_q_down_proj = ( @@ -328,7 +351,7 @@ def get_attention_module_spec_for_backend( linear_q_up_proj=linear_q_up_proj, linear_kv_down_proj=linear_kv_down_proj, linear_kv_up_proj=linear_kv_up_proj, - core_attention=backend.core_attention(), + core_attention=core_attention, linear_proj=backend.row_parallel_linear(), q_layernorm=qk_norm, kv_layernorm=qk_norm, @@ -352,7 +375,7 @@ def get_attention_module_spec_for_backend( params={"attn_mask_type": AttnMaskType.causal}, submodules=SelfAttentionSubmodules( linear_qkv=linear_qkv, - core_attention=backend.core_attention(), + core_attention=core_attention, linear_proj=backend.row_parallel_linear(), q_layernorm=qk_norm, k_layernorm=qk_norm, @@ -522,15 +545,18 @@ def get_gpt_decoder_block_spec( continue linear_attention_type = config.linear_attention_type multi_latent_attention = None + sparse_attention_type = None else: linear_attention_type = None multi_latent_attention = config.multi_latent_attention + sparse_attention_type = config.sparse_attention_type layer_spec_key = f"{mlp_type}_{attention_type}" layer_spec_dict[layer_spec_key] = get_layer_spec_fn( num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, multi_latent_attention=multi_latent_attention, + sparse_attention_type=sparse_attention_type, linear_attention_type=linear_attention_type, **get_layer_spec_kwargs, ) diff --git a/megatron/core/models/gpt/sparse_attention_module_specs.py b/megatron/core/models/gpt/sparse_attention_module_specs.py new file mode 100644 index 00000000000..b6ba992b73e --- /dev/null +++ b/megatron/core/models/gpt/sparse_attention_module_specs.py @@ -0,0 +1,62 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from typing import Optional + +from megatron.core.models.backends import BackendSpecProvider +from megatron.core.transformer.sparse_attention import ( + Indexer, + IndexerSubmodules, + SparseAttention, + SparseAttentionSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec + + +def get_indexer_spec_for_backend( + backend: BackendSpecProvider, normalization: Optional[str] = None +) -> ModuleSpec: + """Helper function to get Indexer module spec for a given backend. + + Args: + backend: Backend specification provider (TE or Local). + normalization: Normalization type ("RMSNorm" or None for LayerNorm). + + Returns: + ModuleSpec for Indexer with appropriate submodules. + """ + rms_norm = normalization == "RMSNorm" + return ModuleSpec( + module=Indexer, + submodules=IndexerSubmodules( + linear_wq_b=backend.linear(), + linear_wk=backend.linear(), + k_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=True), + linear_weights_proj=backend.linear(), + ), + ) + + +def get_sparse_attention_module_spec_for_backend( + backend: BackendSpecProvider, sparse_attention_type: str, normalization: Optional[str] = None +) -> ModuleSpec: + """Helper function to get module spec for Sparse Attention. + + Args: + backend: Backend specification provider (TE or Local). + sparse_attention_type: Type of sparse attention. + normalization: Normalization type ("RMSNorm" or None for LayerNorm). + + Returns: + ModuleSpec for the sparse attention implementation with appropriate submodules. + """ + if sparse_attention_type == "dsa": + # Because TransformerEngine does not support sparse attention yet, we use local + # implementation whether the backend is TransformerEngine or not. + return ModuleSpec( + module=SparseAttention, + submodules=SparseAttentionSubmodules( + indexer=get_indexer_spec_for_backend(backend, normalization=normalization), + ), + ) + else: + raise ValueError(f"Invalid sparse attention type: {sparse_attention_type}") diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index d4e990041ca..ce916c64052 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -172,6 +172,7 @@ def __init__( self.key_hidden_size = self.hidden_size_per_attention_head self.val_hidden_size = self.hidden_size_per_attention_head + # TODO: This is built twice when using MLA, should be refactored. self.core_attention = build_module( submodules.core_attention, config=self.config, @@ -188,6 +189,7 @@ def __init__( and "core_attn" in self.config.recompute_modules ) + # TODO: This is built twice when using MLA, should be refactored. # Output. self.linear_proj = build_module( submodules.linear_proj, diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index a8893ebec36..103eaf91a72 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -148,6 +148,7 @@ def __init__( "'rope' and 'yarn'" ) + # TODO(kunlunl): Support sparse attention. self.core_attention = build_module( submodules.core_attention, config=self.config, diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py new file mode 100644 index 00000000000..8e0a448e8c8 --- /dev/null +++ b/megatron/core/transformer/sparse_attention.py @@ -0,0 +1,662 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import torch + +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + _yarn_get_mscale, + apply_rotary_pos_emb, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +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 +from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + +# TODO(kunlunl): Add third-party fused kernels. + + +class IndexerLossAutoScaler(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 IndexerLossAutoScaler.main_loss_backward_scale is None: + IndexerLossAutoScaler.main_loss_backward_scale = torch.tensor( + 1.0, device=indexer_loss.device + ) + indexer_loss_backward_scale = IndexerLossAutoScaler.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. Please ensure that the scale passed in + matches the scale of the main_loss. + """ + if IndexerLossAutoScaler.main_loss_backward_scale is None: + IndexerLossAutoScaler.main_loss_backward_scale = scale + else: + IndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) + + +def compute_indexer_loss( + index_scores: torch.Tensor, + attention_scores: torch.Tensor, + indexer_loss_coeff: float, +) -> torch.Tensor: + """ + Compute KL divergence loss between indexer scores and true attention scores. + + This loss trains the indexer to predict which tokens are important + by matching the distribution of true attention scores. + + Args: + index_scores: Scores predicted by indexer [batch, seq, seq] + attention_scores: True attention scores from q@k [batch, heads, seq, seq] + indexer_loss_coeff: Coefficient for the indexer KL divergence loss + + Returns: + indexer_loss: KL divergence loss (scalar) + """ + # Average attention scores across heads to get target distribution + # [batch, heads, seq, seq] -> [batch, seq, seq] + target_scores = attention_scores.mean(dim=1) + + # Convert to probabilities with softmax + # Apply softmax over the last dimension (keys) + index_probs = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + target_probs = torch.nn.functional.softmax(target_scores, dim=-1, dtype=torch.float32) + + # Compute KL divergence: KL(target || index) + # KL(P || Q) = Σ P(x) * log(P(x) / Q(x)) + kl_div = torch.nn.functional.kl_div( + index_probs.log(), + target_probs, + reduction='batchmean', + log_target=False, + ) + + # Scale by coefficient + indexer_loss = kl_div * indexer_loss_coeff + + return indexer_loss + + +def rotate_activation(x: torch.Tensor) -> torch.Tensor: + """Apply Hadamard rotation activation. + + Args: + x: Input tensor (must be bfloat16) + + Returns: + Rotated tensor + """ + assert x.dtype == torch.bfloat16 + from fast_hadamard_transform import hadamard_transform + hidden_size = x.size(-1) + return hadamard_transform(x, scale=hidden_size ** -0.5) + + +def compute_index_score(q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor) -> torch.Tensor: + """ + Compute index scores for sparse attention (BF16 version). + + This is a BF16 implementation of the FP8 index kernel logic: + 1. Compute attention scores: q @ k^T + 2. Apply ReLU activation + 3. Weight by attention weights + 4. Sum across attention heads + + Args: + q: Query tensor [batch, seq_len, n_heads, head_dim] (bf16) + weights: Attention weights [batch, seq_len, n_heads, 1] (bf16) + k: Key tensor [batch, seq_len, head_dim] (bf16) + + Returns: + index_score: Index scores [batch, seq_len, seq_len] (bf16) + + Note: Original FP8 kernel signature was: + fp8_index(q, q_s, k, k_s) where q_s and k_s are scaling factors + Here we use BF16 directly without separate scaling factors. + """ + # q: [bsz, seqlen, n_heads, head_dim] + # k: [bsz, seqlen, head_dim] + # weights: [bsz, seqlen, n_heads, 1] + + # Compute attention scores: q @ k^T + # [bsz, seqlen, n_heads, head_dim] @ [bsz, seqlen, head_dim]^T + # -> [bsz, seqlen, n_heads, seqlen] + index_score = torch.einsum('bshd,btd->bsht', q, k) + + # Apply ReLU activation (for throughput efficiency) + index_score = torch.relu(index_score) + + # Weight each head by attention weights + # [bsz, seqlen, n_heads, seqlen] * [bsz, seqlen, n_heads, 1] + index_score = index_score * weights + + # Sum across attention heads + # [bsz, seqlen, n_heads, seqlen] -> [bsz, seqlen, seqlen] + index_score = index_score.sum(dim=2) + + return index_score + + + +@dataclass +class IndexerSubmodules: + """ + Configuration class for specifying the submodules of an 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 SparseAttentionSubmodules: + """ + Configuration class for specifying the submodules of SparseAttention. + + Args: + indexer: Indexer module for computing sparse attention indices. + """ + indexer: Union[ModuleSpec, type] = None + + +class Indexer(MegatronModule): + """ + Lightning Indexer for DeepSeek Sparse Attention. + + Computes index scores to identify the top-k most relevant key-value pairs + for each query position in sparse attention. + + Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py + + Args: + config: Transformer configuration. + submodules: Indexer submodules specification. + dim: Model hidden dimension. + n_heads: Number of attention heads. + head_dim: Dimension per attention head. + rope_head_dim: Dimension for rotary position embeddings. + index_topk: Number of top-k indices to select. + q_lora_rank: Rank for low-rank query projection. + pg_collection: Process group collection for tensor parallelism. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: IndexerSubmodules, + dim: int, + n_heads: int, + head_dim: int, + rope_head_dim: int, + index_topk: int, + q_lora_rank: int, + pg_collection: ProcessGroupCollection = None, + ): + super().__init__(config=config) + + self.config = config + self.dim = dim + self.n_heads = n_heads + self.head_dim = head_dim + self.rope_head_dim = rope_head_dim + self.index_topk = index_topk + self.q_lora_rank = q_lora_rank + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + self.pg_collection = pg_collection + world_size = pg_collection.tp.size() + self.n_local_heads = n_heads // world_size + + # Initialize Rotary Position Embedding. + # Use rope_type from config if available, default to "rope". + if self.config.rope_type == "rope": + self.rotary_pos_emb = RotaryEmbedding( + self.rope_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.rope_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 'yarn'" + ) + + # Build linear layers using build_module + self.wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.n_heads * self.head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + ) + + self.wk = build_module( + submodules.linear_wk, + self.dim, + self.head_dim, + config=self.config, + init_method=self.config.init_method, + bias=False, + ) + + self.k_norm = build_module( + submodules.k_norm, + config=self.config, + hidden_size=self.head_dim, + eps=self.config.layernorm_epsilon, + ) + + # TODO(kunlunl): The dtype of this module should be torch.get_default_dtype(). + self.weights_proj = build_module( + submodules.linear_weights_proj, + self.dim, + self.n_heads, + config=self.config, + init_method=self.config.init_method, + bias=False, + ) + + self.softmax_scale: float = self.head_dim ** -0.5 + + def forward( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional = None, + ): + """ + Forward pass for Indexer. + + Args: + x: Input tensor [batch, seq_len, hidden_dim] or [seq, batch, hidden_dim] + qr: Query representation tensor [batch, seq_len, q_lora_rank] + mask: Attention mask + packed_seq_params: Packed sequence parameters for variable length sequences + + Returns: + topk_indices: Top-k indices for sparse attention [batch, seq_len, index_topk] + """ + # Call forward_with_scores and only return indices + _, topk_indices = self.forward_with_scores(x, qr, mask, packed_seq_params) + return topk_indices + + def forward_with_scores( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for 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: Input tensor [batch, seq_len, hidden_dim] + qr: Query representation tensor [batch, seq_len, q_lora_rank] + mask: Attention mask + packed_seq_params: Packed sequence parameters + + Returns: + index_scores: Index scores [batch, seq_len, seq_len] + topk_indices: Top-k indices [batch, seq_len, index_topk] + """ + bsz, seqlen, _ = x.size() + + # Compute rotary position embeddings internally + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(seqlen, packed_seq=packed_seq) + mscale = 1.0 + else: # yarn + rotary_pos_emb, mscale = self.rotary_pos_emb(seqlen, packed_seq=packed_seq) + + q = self.wq_b(qr) + q = q.reshape(bsz, seqlen, -1, self.head_dim) + q_pe, q_nope = torch.split( + q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + # Apply RoPE to query position embedding part + q_pe = apply_rotary_pos_emb( + q_pe, + rotary_pos_emb, + config=self.config, + cp_group=self.pg_collection.cp, + mscale=mscale, + ) + q = torch.cat([q_pe, q_nope], dim=-1) + + k = self.wk(x) + k = self.k_norm(k) + k_pe, k_nope = torch.split( + k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + ) + + # Apply RoPE to key position embedding part + k_pe = k_pe.unsqueeze(2) # [batch, seq_len, 1, rope_head_dim] + k_pe = apply_rotary_pos_emb( + k_pe, + rotary_pos_emb, + config=self.config, + cp_group=self.pg_collection.cp, + mscale=mscale, + ) + k_pe = k_pe.squeeze(2) # [batch, seq_len, rope_head_dim] + k = torch.cat([k_pe, k_nope], dim=-1) + + q = rotate_activation(q) + k = rotate_activation(k) + + weights = self.weights_proj(x) * self.n_heads ** -0.5 + weights = weights.unsqueeze(-1) * self.softmax_scale + + # Compute index scores (BF16 version of the FP8 kernel) + index_scores = compute_index_score(q.contiguous(), weights, k.contiguous()) + + if mask is not None: + index_scores = index_scores + mask + + # Select top-k indices + topk_k = min(self.index_topk, seqlen) + topk_indices = index_scores.topk(topk_k, dim=-1)[1] + + return index_scores, topk_indices + + +class SparseAttention(MegatronModule): + """ + This module implements sparse attention mechanism using an 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 + + Args: + config: Transformer configuration. + submodules: Sparse attention submodules specification. + layer_number: Layer number in the model. + attn_mask_type: Type of attention mask. + attention_type: Type of attention. + attention_dropout: Dropout probability for attention weights. + softmax_scale: Scale factor for softmax. + k_channels: Number of channels in key tensor. + v_channels: Number of channels in value tensor. + cp_comm_type: Context parallel communication type. + pg_collection: Process group collection for distributed training. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: SparseAttentionSubmodules, + 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.config: TransformerConfig = config + + assert ( + self.config.context_parallel_size == 1 + ), "Currently context parallelism is not supported by SparseAttention!" + + self.layer_number = max(1, layer_number) + self.attn_mask_type = attn_mask_type + self.attention_type = attention_type + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp']) + else: + assert hasattr( + pg_collection, 'tp' + ), "SparseAttention pg_collection must have tp process group" + self.pg_collection = pg_collection + + world_size = pg_collection.tp.size() + + 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 + + # Build indexer - required for sparse attention + assert submodules.indexer is not None, "Indexer is required for SparseAttention" + self.indexer = build_module( + submodules.indexer, + config=self.config, + dim=self.config.hidden_size, + n_heads=self.config.index_n_heads, + head_dim=self.config.index_head_dim, + rope_head_dim=self.config.qk_pos_emb_head_dim, + index_topk=self.config.index_topk, + q_lora_rank=self.config.q_lora_rank, + pg_collection=self.pg_collection, + ) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params: PackedSeqParams = None, + x: Optional[torch.Tensor] = None, + qr: Optional[torch.Tensor] = None, + ): + """ + Forward pass for Sparse Attention. + + Args: + query: Query tensor [sq, b, np, hn] + key: Key tensor [sk, b, np, hn] + value: Value tensor [sk, b, np, hn] + attention_mask: Attention mask tensor + attn_mask_type: Type of attention mask + attention_bias: Optional attention bias + packed_seq_params: Packed sequence parameters + x: Original input hidden states [b, s, h] (needed for indexer, will be inferred if not provided) + qr: Low-rank query representation [b, s, q_lora_rank] (for MLA, will be inferred if not provided) + + Returns: + context: Output tensor [sq, b, hp] + """ + # Input shape: [sq, b, np, hn] + sq, b, np, hn = query.size() + sk = key.size(0) + + # Prepare inputs for indexer (expects batch-first [b, s, h] format) + if x is None: + # Convert query from [sq, b, np, hn] to [b, sq, np*hn] + x = query.transpose(0, 1).reshape(b, sq, np * hn) + if qr is None: + # For non-MLA, use x as qr + qr = x + + # =================================== + # Raw attention scores [b, np, sq, sk] + # =================================== + output_size = (b, np, sq, sk) + + # Reshape for batch matrix multiplication + # [sq, b, np, hn] -> [b * np, sq, hn] + query_reshaped = query.transpose(0, 1).reshape(b * np, sq, hn) + # [sk, b, np, hn] -> [b * np, sk, hn] + key_reshaped = key.transpose(0, 1).reshape(b * np, sk, hn) + + # Compute attention scores: [b * np, sq, sk] + attention_scores = torch.bmm( + query_reshaped, + key_reshaped.transpose(1, 2) + ) * self.softmax_scale + + # Reshape to [b, np, sq, sk] + attention_scores = attention_scores.view(*output_size) + + # =================================== + # Use Indexer for sparse selection + # =================================== + # Get index scores and top-k indices + # Note: We need to get index_scores before topk for KL loss computation + index_scores, topk_indices = self.indexer.forward_with_scores( + x, qr, mask=None, packed_seq_params=packed_seq_params + ) + + # =================================== + # Compute and attach indexer loss + # =================================== + if self.training and torch.is_grad_enabled(): + # Get indexer loss coefficient from config + indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) + + if indexer_loss_coeff > 0: + # Compute KL divergence loss between indexer scores and true attention scores + indexer_loss = compute_indexer_loss( + index_scores, + attention_scores.detach(), # Don't backprop through attention scores + indexer_loss_coeff + ) + + # Attach loss to query activation (will be backpropagated) + # This doesn't change the forward pass but triggers gradient flow in backward + query = IndexerLossAutoScaler.apply(query, indexer_loss) + + # =================================== + # Apply sparse mask from indexer + # =================================== + # topk_indices: [b, sq, topk] + # Create sparse mask + index_mask = torch.full( + (b, sq, sk), float("-inf"), device=query.device, dtype=attention_scores.dtype + ) + # Fill top-k positions with 0 (allow attention) + index_mask.scatter_(-1, topk_indices, 0) + + # Expand index_mask to [b, np, sq, sk] + index_mask = index_mask.unsqueeze(1).expand(-1, np, -1, -1) + + # Combine with regular attention mask + if attention_mask is not None: + index_mask = index_mask + attention_mask + + attention_scores = attention_scores + index_mask + + # =================================== + # Attention probabilities [b, np, sq, sk] + # =================================== + attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) + attention_probs = attention_probs.to(query.dtype) + + # =================================== + # Context layer [sq, b, hp] + # =================================== + # Reshape value: [sk, b, np, hn] -> [b * np, sk, hn] + value_reshaped = value.transpose(0, 1).reshape(b * np, sk, hn) + + # Reshape attention_probs: [b, np, sq, sk] -> [b * np, sq, sk] + attention_probs_reshaped = attention_probs.view(b * np, sq, sk) + + # Compute context: [b * np, sq, hn] + context = torch.bmm(attention_probs_reshaped, value_reshaped) + + # Reshape context: [b * np, sq, hn] -> [b, np, sq, hn] -> [sq, b, np, hn] + context = context.view(b, np, sq, hn).permute(2, 0, 1, 3).contiguous() + + # Flatten: [sq, b, np, hn] -> [sq, b, np*hn] + context = context.view(sq, b, np * hn) + + return context + + def sharded_state_dict( + self, + prefix: str = '', + sharded_offsets: Tuple[Tuple[int, int, int]] = (), + metadata: Optional[dict] = None, + ) -> ShardedStateDict: + """Sharded state dict for the learnable softmax offset parameter""" + if self.config.softmax_type == "learnable": + state_dict = self.state_dict(prefix="", keep_vars=True) + else: + state_dict = {} + return make_sharded_tensors_for_checkpoint( + state_dict, prefix, {'softmax_offset': 0}, sharded_offsets + ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index b39b7706feb..9a050e71416 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -219,6 +219,24 @@ class TransformerConfig(ModelParallelConfig): """Number of SMs to use for HybridEP. In pure NVL scenarios, 16 SMs can generally achieve good bandwidth.""" + #################### + # sparse attention + #################### + sparse_attention_type: Optional[str] = None + """Type of sparse attention to use. Currently only supports dsa (DeepSeek Sparse Attention).""" + + index_n_heads: Optional[int] = None + """Number of indexer heads for sparse attention. If None, defaults to num_attention_heads.""" + + index_head_dim: Optional[int] = None + """Dimension per indexer head. If None, defaults to kv_channels.""" + + index_topk: int = 256 + """Number of top-k tokens to select in sparse attention indexer.""" + + indexer_loss_coeff: float = 0.0 + """Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.""" + #################### # linear attention #################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index bdf915a8ae1..c7e914f1d5f 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -69,6 +69,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_sparse_attention_args(parser) parser = _add_linear_attention_args(parser) parser = _add_heterogeneous_args(parser) parser = _add_logging_args(parser) @@ -3266,6 +3267,20 @@ def _add_mla_args(parser): return parser +def _add_sparse_attention_args(parser): + group = parser.add_argument_group(title="sparse_attention") + group.add_argument('--sparse-attention-type', default=None, choices=['dsa'], type=str, + help="Type of sparse attention to use. Currently support dsa (DeepSeek Sparse Attention).") + group.add_argument('--index-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('--index-head-dim', default=None, type=int, + help='Dimension per indexer head for sparse attention. If not set, defaults to kv-channels.') + group.add_argument('--index-topk', default=256, type=int, + help='Number of top-k tokens to select in sparse attention indexer.') + group.add_argument('--indexer-loss-coeff', default=0.0, type=float, + help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') + return parser + def _add_linear_attention_args(parser): group = parser.add_argument_group(title="la") group.add_argument('--linear-attention-type', default=None, choices=['gated_delta_net', 'mamba'], type=str, From 501ad656249c3724440709a64c14bf10b4242bb5 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 10 Nov 2025 16:36:08 +0800 Subject: [PATCH 02/28] Run through single GPU version Signed-off-by: kunlunl --- .../transformer/multi_latent_attention.py | 77 +++++++--- megatron/core/transformer/sparse_attention.py | 134 ++++++++++-------- 2 files changed, 127 insertions(+), 84 deletions(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 103eaf91a72..2ca88d83815 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -234,13 +234,22 @@ def forward( # Get the query, key and value tensors based on the type of attention - # self or cross attn. # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128] - query, key, value = self.get_query_key_value_tensors( - hidden_states, - key_value_states, - position_ids, - packed_seq_params, - inference_context=inference_context, - ) + if self.config.sparse_attention_type is None: + query, key, value = self.get_query_key_value_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + inference_context=inference_context, + ) + else: + query, key, value, q_compressed, _ = self.get_query_key_value_and_compressed_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + inference_context=inference_context, + ) # =================================================== # Adjust key, value for inference @@ -268,14 +277,28 @@ def forward( ) else: if inference_context is None or inference_context.is_static_batching(): - core_attn_out = self.core_attention( - query, - key, - value, - attention_mask, - packed_seq_params=packed_seq_params, - attn_mask_type=attn_mask_type, - ) + if self.config.sparse_attention_type is None: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=packed_seq_params, + attn_mask_type=attn_mask_type, + ) + else: + # For sparse attention, use a specialized forward. + core_attn_out = self.core_attention( + query, + key, + value, + hidden_states, + q_compressed, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=None, + packed_seq_params=packed_seq_params, + ) elif self.cache_mla_latents: # Dynamic batching attention kernel. q, k, v = (query, key, value) @@ -461,7 +484,7 @@ def __init__( eps=self.config.layernorm_epsilon, ) - def get_query_key_value_tensors( + def get_query_key_value_and_compressed_tensors( self, hidden_states, key_value_states=None, @@ -579,6 +602,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 # ========================================= @@ -589,7 +622,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] @@ -599,8 +631,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) @@ -664,7 +694,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] @@ -674,8 +703,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) @@ -799,7 +826,11 @@ 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 get_query_key_value_tensors(self, *args, **kwargs): + # Only return q, k, v + return get_query_key_value_and_compressed_tensors(self, *args, **kwargs)[:3] def uncompress_kv_from_cache(self, kv_cached): """ diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 8e0a448e8c8..b4c47cd3278 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -296,6 +296,9 @@ def __init__( config=self.config, init_method=self.config.init_method, bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", ) self.wk = build_module( @@ -305,6 +308,9 @@ def __init__( config=self.config, init_method=self.config.init_method, bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", ) self.k_norm = build_module( @@ -322,6 +328,9 @@ def __init__( config=self.config, init_method=self.config.init_method, bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", ) self.softmax_scale: float = self.head_dim ** -0.5 @@ -337,8 +346,8 @@ def forward( Forward pass for Indexer. Args: - x: Input tensor [batch, seq_len, hidden_dim] or [seq, batch, hidden_dim] - qr: Query representation tensor [batch, seq_len, q_lora_rank] + x: Input tensor [seq, batch, hidden_dim] + qr: Query representation tensor [seq, batch, q_lora_rank] mask: Attention mask packed_seq_params: Packed sequence parameters for variable length sequences @@ -362,8 +371,8 @@ def forward_with_scores( This is used when KL loss is enabled to compare indexer scores with true attention scores. Args: - x: Input tensor [batch, seq_len, hidden_dim] - qr: Query representation tensor [batch, seq_len, q_lora_rank] + x: Input tensor [seq, batch, hidden_dim] + qr: Query representation tensor [seq, batch, q_lora_rank] mask: Attention mask packed_seq_params: Packed sequence parameters @@ -371,7 +380,8 @@ def forward_with_scores( index_scores: Index scores [batch, seq_len, seq_len] topk_indices: Top-k indices [batch, seq_len, index_topk] """ - bsz, seqlen, _ = x.size() + # Input format: [seq, batch, hidden] + seqlen, bsz, _ = x.size() # Compute rotary position embeddings internally packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' @@ -381,13 +391,12 @@ def forward_with_scores( else: # yarn rotary_pos_emb, mscale = self.rotary_pos_emb(seqlen, packed_seq=packed_seq) - q = self.wq_b(qr) - q = q.reshape(bsz, seqlen, -1, self.head_dim) + q, _ = self.wq_b(qr) # [seq, batch, heads*head_dim] + q = q.reshape(seqlen, bsz, -1, self.head_dim) # [seq, batch, heads, head_dim] q_pe, q_nope = torch.split( q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - # Apply RoPE to query position embedding part q_pe = apply_rotary_pos_emb( q_pe, rotary_pos_emb, @@ -395,16 +404,15 @@ def forward_with_scores( cp_group=self.pg_collection.cp, mscale=mscale, ) - q = torch.cat([q_pe, q_nope], dim=-1) + q = torch.cat([q_pe, q_nope], dim=-1) # [seq, batch, heads, head_dim] - k = self.wk(x) - k = self.k_norm(k) + k, _ = self.wk(x) # [seq, batch, head_dim] + k = self.k_norm(k) # [seq, batch, head_dim] k_pe, k_nope = torch.split( k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 ) - # Apply RoPE to key position embedding part - k_pe = k_pe.unsqueeze(2) # [batch, seq_len, 1, rope_head_dim] + k_pe = k_pe.unsqueeze(2) # [seq, batch, 1, rope_head_dim] k_pe = apply_rotary_pos_emb( k_pe, rotary_pos_emb, @@ -412,16 +420,25 @@ def forward_with_scores( cp_group=self.pg_collection.cp, mscale=mscale, ) - k_pe = k_pe.squeeze(2) # [batch, seq_len, rope_head_dim] - k = torch.cat([k_pe, k_nope], dim=-1) - + k_pe = k_pe.squeeze(2) # [seq, batch, rope_head_dim] + k = torch.cat([k_pe, k_nope], dim=-1) # [seq, batch, head_dim] + + # For compute_index_score, we need batch-first format + # Transpose to [batch, seq, heads, head_dim] + q = q.transpose(0, 1) # [seq, batch, heads, head_dim] -> [batch, seq, heads, head_dim] + k = k.transpose(0, 1) # [seq, batch, head_dim] -> [batch, seq, head_dim] + q = rotate_activation(q) k = rotate_activation(k) - weights = self.weights_proj(x) * self.n_heads ** -0.5 + # weights_proj expects seq-first, so use original x + weights, _ = self.weights_proj(x) # [seq, batch, n_heads] + weights = weights.transpose(0, 1) # -> [batch, seq, n_heads] + weights = weights * self.n_heads ** -0.5 weights = weights.unsqueeze(-1) * self.softmax_scale # Compute index scores (BF16 version of the FP8 kernel) + # All inputs are now batch-first index_scores = compute_index_score(q.contiguous(), weights, k.contiguous()) if mask is not None: @@ -516,12 +533,12 @@ def forward( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, attention_mask: torch.Tensor, attn_mask_type: AttnMaskType = None, attention_bias: torch.Tensor = None, packed_seq_params: PackedSeqParams = None, - x: Optional[torch.Tensor] = None, - qr: Optional[torch.Tensor] = None, ): """ Forward pass for Sparse Attention. @@ -543,14 +560,19 @@ def forward( # Input shape: [sq, b, np, hn] sq, b, np, hn = query.size() sk = key.size(0) + v_hn = value.size(3) # Value head dimension may differ from query/key - # Prepare inputs for indexer (expects batch-first [b, s, h] format) - if x is None: - # Convert query from [sq, b, np, hn] to [b, sq, np*hn] - x = query.transpose(0, 1).reshape(b, sq, np * hn) - if qr is None: - # For non-MLA, use x as qr - qr = x + # =================================== + # Use Indexer for sparse selection (do this first to compute loss before using query) + # =================================== + # Get index scores and top-k indices + # Note: We need to get index_scores before topk for KL loss computation + # Detach x and qr to prevent gradients from flowing back to the main model + # Indexer is trained solely through the indexer_loss (KL divergence) + # TODO(kunlunl): Should x and qr be detached? + index_scores, topk_indices = self.indexer.forward_with_scores( + x.detach(), qr.detach(), mask=None, packed_seq_params=packed_seq_params + ) # =================================== # Raw attention scores [b, np, sq, sk] @@ -572,34 +594,6 @@ def forward( # Reshape to [b, np, sq, sk] attention_scores = attention_scores.view(*output_size) - # =================================== - # Use Indexer for sparse selection - # =================================== - # Get index scores and top-k indices - # Note: We need to get index_scores before topk for KL loss computation - index_scores, topk_indices = self.indexer.forward_with_scores( - x, qr, mask=None, packed_seq_params=packed_seq_params - ) - - # =================================== - # Compute and attach indexer loss - # =================================== - if self.training and torch.is_grad_enabled(): - # Get indexer loss coefficient from config - indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) - - if indexer_loss_coeff > 0: - # Compute KL divergence loss between indexer scores and true attention scores - indexer_loss = compute_indexer_loss( - index_scores, - attention_scores.detach(), # Don't backprop through attention scores - indexer_loss_coeff - ) - - # Attach loss to query activation (will be backpropagated) - # This doesn't change the forward pass but triggers gradient flow in backward - query = IndexerLossAutoScaler.apply(query, indexer_loss) - # =================================== # Apply sparse mask from indexer # =================================== @@ -629,20 +623,38 @@ def forward( # =================================== # Context layer [sq, b, hp] # =================================== - # Reshape value: [sk, b, np, hn] -> [b * np, sk, hn] - value_reshaped = value.transpose(0, 1).reshape(b * np, sk, hn) + # Reshape value: [sk, b, np, v_hn] -> [b * np, sk, v_hn] + value_reshaped = value.transpose(0, 1).reshape(b * np, sk, v_hn) # Reshape attention_probs: [b, np, sq, sk] -> [b * np, sq, sk] attention_probs_reshaped = attention_probs.view(b * np, sq, sk) - # Compute context: [b * np, sq, hn] + # Compute context: [b * np, sq, v_hn] context = torch.bmm(attention_probs_reshaped, value_reshaped) - # Reshape context: [b * np, sq, hn] -> [b, np, sq, hn] -> [sq, b, np, hn] - context = context.view(b, np, sq, hn).permute(2, 0, 1, 3).contiguous() + # Reshape context: [b * np, sq, v_hn] -> [b, np, sq, v_hn] -> [sq, b, np, v_hn] + context = context.view(b, np, sq, v_hn).permute(2, 0, 1, 3).contiguous() + + # Flatten: [sq, b, np, v_hn] -> [sq, b, np*v_hn] + context = context.view(sq, b, np * v_hn) + + # =================================== + # Attach indexer loss (training only) + # =================================== + if self.training and torch.is_grad_enabled(): + # Get indexer loss coefficient from config + indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) + + if indexer_loss_coeff > 0: + # Compute KL divergence loss between indexer scores and true attention scores + indexer_loss = compute_indexer_loss( + index_scores, + attention_scores.detach(), # Don't backprop through attention scores + indexer_loss_coeff + ) - # Flatten: [sq, b, np, hn] -> [sq, b, np*hn] - context = context.view(sq, b, np * hn) + # Attach loss to context output (will trigger backward through indexer) + context = IndexerLossAutoScaler.apply(context, indexer_loss) return context From c3754a78bd0182050013da5d6ce6f63a3d6a9253 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 10 Nov 2025 20:14:50 +0800 Subject: [PATCH 03/28] Fix the usage of attention_mask Signed-off-by: kunlunl --- .../transformer/multi_latent_attention.py | 9 ++-- megatron/core/transformer/sparse_attention.py | 46 +++++++++++-------- megatron/training/arguments.py | 2 +- 3 files changed, 32 insertions(+), 25 deletions(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 2ca88d83815..462d6935aa7 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -148,7 +148,6 @@ def __init__( "'rope' and 'yarn'" ) - # TODO(kunlunl): Support sparse attention. self.core_attention = build_module( submodules.core_attention, config=self.config, @@ -292,9 +291,9 @@ def forward( query, key, value, - hidden_states, - q_compressed, - attention_mask, + x=hidden_states, + qr=q_compressed, + attention_mask=attention_mask, attn_mask_type=attn_mask_type, attention_bias=None, packed_seq_params=packed_seq_params, @@ -829,7 +828,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po return query, key, value, q_compressed, kv_compressed def get_query_key_value_tensors(self, *args, **kwargs): - # Only return q, k, v + # Only return query, key and value. return get_query_key_value_and_compressed_tensors(self, *args, **kwargs)[:3] def uncompress_kv_from_cache(self, kv_cached): diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index b4c47cd3278..ee3b2d88eed 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -598,21 +598,30 @@ def forward( # Apply sparse mask from indexer # =================================== # topk_indices: [b, sq, topk] - # Create sparse mask - index_mask = torch.full( - (b, sq, sk), float("-inf"), device=query.device, dtype=attention_scores.dtype + + # Step 1: Create sparse selection mask (only allow top-k) + # Start with all positions masked (True = masked) + sparse_mask = torch.ones( + (b, sq, sk), dtype=torch.bool, device=query.device ) - # Fill top-k positions with 0 (allow attention) - index_mask.scatter_(-1, topk_indices, 0) + # Allow top-k positions (False = not masked) + sparse_mask.scatter_(-1, topk_indices, False) - # Expand index_mask to [b, np, sq, sk] - index_mask = index_mask.unsqueeze(1).expand(-1, np, -1, -1) + # Expand to [b, np, sq, sk] + sparse_mask = sparse_mask.unsqueeze(1).expand(-1, np, -1, -1) - # Combine with regular attention mask + # Step 2: Combine with regular attention mask if provided if attention_mask is not None: - index_mask = index_mask + attention_mask + # attention_mask is boolean: True = masked, False = allowed + # Combine with OR: position is masked if EITHER mask says so + if attention_mask.dtype == torch.bool: + sparse_mask = sparse_mask | attention_mask + else: + # If attention_mask is float (-inf for masked), convert to boolean + sparse_mask = sparse_mask | (attention_mask < -1000.0) - attention_scores = attention_scores + index_mask + # Step 3: Apply combined mask to attention scores + attention_scores = attention_scores.masked_fill(sparse_mask, float("-inf")) # =================================== # Attention probabilities [b, np, sq, sk] @@ -645,16 +654,15 @@ def forward( # Get indexer loss coefficient from config indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) - if indexer_loss_coeff > 0: - # Compute KL divergence loss between indexer scores and true attention scores - indexer_loss = compute_indexer_loss( - index_scores, - attention_scores.detach(), # Don't backprop through attention scores - indexer_loss_coeff - ) + # Compute KL divergence loss between indexer scores and true attention scores + indexer_loss = compute_indexer_loss( + index_scores, + attention_scores.detach(), # Don't backprop through attention scores + indexer_loss_coeff + ) - # Attach loss to context output (will trigger backward through indexer) - context = IndexerLossAutoScaler.apply(context, indexer_loss) + # Attach loss to context output (will trigger backward through indexer) + context = IndexerLossAutoScaler.apply(context, indexer_loss) return context diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c7e914f1d5f..95d883ad714 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3275,7 +3275,7 @@ def _add_sparse_attention_args(parser): help='Number of indexer heads for sparse attention. If not set, defaults to num-attention-heads.') group.add_argument('--index-head-dim', default=None, type=int, help='Dimension per indexer head for sparse attention. If not set, defaults to kv-channels.') - group.add_argument('--index-topk', default=256, type=int, + group.add_argument('--index-topk', default=None, type=int, help='Number of top-k tokens to select in sparse attention indexer.') group.add_argument('--indexer-loss-coeff', default=0.0, type=float, help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') From 451a36bc9a25c2e1b310d551bc8a307d15971fe7 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 11 Nov 2025 03:51:03 +0800 Subject: [PATCH 04/28] Update mask & indexer loss Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 502 +++++++++--------- 1 file changed, 253 insertions(+), 249 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index ee3b2d88eed..c0eb849185e 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -10,7 +10,6 @@ from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, - _yarn_get_mscale, apply_rotary_pos_emb, ) from megatron.core.packed_seq_params import PackedSeqParams @@ -21,7 +20,122 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import ( + fused_apply_mla_rope_for_kv, + fused_apply_mla_rope_for_q, + ) +except: + fused_apply_mla_rope_for_kv = None + fused_apply_mla_rope_for_q = None + # TODO(kunlunl): Add third-party fused kernels. +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 + 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) + + +def compute_index_score(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, bsz, n_heads, head_dim], the query tensor. + weights : FP32 [seqlen_q, bsz, n_heads], the attention weights. + k : BF16 [seqlen_k, bsz, head_dim], the key tensor. + + Returns: + index_score: FP32 [bsz, seqlen_q, seqlen_k], the index scores. + """ + # Compute attention scores: q @ k^T + # [seqlen_q, bsz, n_heads, head_dim] @ [seqlen_k, bsz, head_dim]^T -> [seqlen_q, bsz, n_heads, seqlen_k] + index_score = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) + + # Apply ReLU activation. + index_score = torch.relu(index_score) + + # Weight each head by attention weights. + # [seqlen_q, bsz, n_heads, seqlen_k] * [seqlen_q, bsz, n_heads, 1] -> [seqlen_q, bsz, n_heads, seqlen_k] + index_score = index_score * weights.unsqueeze(-1) + + # Sum across attention heads. + # [seqlen_q, bsz, n_heads, seqlen_k] -> [seqlen_q, bsz, seqlen_k] + index_score = index_score.sum(dim=2) + + # Transpose to [bsz, seqlen_q, seqlen_k]. + index_score = index_score.transpose(0, 1) + + return index_score + + +def compute_indexer_loss( + index_scores: torch.Tensor, + attention_scores: torch.Tensor, + indexer_loss_coeff: float, +) -> 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 [bsz, seqlen_q, seqlen_k]. + attention_scores: True attention scores from q @ k^T [bsz, heads, seqlen_q, seqlen_k]. + indexer_loss_coeff: Coefficient for the indexer KL divergence loss. + + Returns: + index_loss: KL divergence loss (scalar). + """ + # Sum attention scores across heads. + # [bsz, heads, seqlen_q, seqlen_k] -> [bsz, seqlen_q, seqlen_k] + target_scores = attention_scores.sum(dim=1) + + # L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are + # obtained from softmax so they are already non-negative. + target_probs = target_scores / target_scores.sum(dim=-1, keepdim=True) + + # Convert index_scores to probabilities with softmax. + index_probs = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + + # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) + kl_per_element = ( + target_probs * (torch.log(target_probs + 1e-10) - torch.log(index_probs + 1e-10)) + ) + kl_div = kl_per_element.sum(dim=-1).mean() + + # Scale by coefficient. + indexer_loss = kl_div * indexer_loss_coeff + + return indexer_loss class IndexerLossAutoScaler(torch.autograd.Function): @@ -71,8 +185,7 @@ def set_loss_scale(scale: torch.Tensor): """Set the scale of the indexer loss. Args: - scale: The scale value to set. Please ensure that the scale passed in - matches the scale of the main_loss. + scale: The scale value to set. """ if IndexerLossAutoScaler.main_loss_backward_scale is None: IndexerLossAutoScaler.main_loss_backward_scale = scale @@ -80,110 +193,6 @@ def set_loss_scale(scale: torch.Tensor): IndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) -def compute_indexer_loss( - index_scores: torch.Tensor, - attention_scores: torch.Tensor, - indexer_loss_coeff: float, -) -> torch.Tensor: - """ - Compute KL divergence loss between indexer scores and true attention scores. - - This loss trains the indexer to predict which tokens are important - by matching the distribution of true attention scores. - - Args: - index_scores: Scores predicted by indexer [batch, seq, seq] - attention_scores: True attention scores from q@k [batch, heads, seq, seq] - indexer_loss_coeff: Coefficient for the indexer KL divergence loss - - Returns: - indexer_loss: KL divergence loss (scalar) - """ - # Average attention scores across heads to get target distribution - # [batch, heads, seq, seq] -> [batch, seq, seq] - target_scores = attention_scores.mean(dim=1) - - # Convert to probabilities with softmax - # Apply softmax over the last dimension (keys) - index_probs = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) - target_probs = torch.nn.functional.softmax(target_scores, dim=-1, dtype=torch.float32) - - # Compute KL divergence: KL(target || index) - # KL(P || Q) = Σ P(x) * log(P(x) / Q(x)) - kl_div = torch.nn.functional.kl_div( - index_probs.log(), - target_probs, - reduction='batchmean', - log_target=False, - ) - - # Scale by coefficient - indexer_loss = kl_div * indexer_loss_coeff - - return indexer_loss - - -def rotate_activation(x: torch.Tensor) -> torch.Tensor: - """Apply Hadamard rotation activation. - - Args: - x: Input tensor (must be bfloat16) - - Returns: - Rotated tensor - """ - assert x.dtype == torch.bfloat16 - from fast_hadamard_transform import hadamard_transform - hidden_size = x.size(-1) - return hadamard_transform(x, scale=hidden_size ** -0.5) - - -def compute_index_score(q: torch.Tensor, weights: torch.Tensor, k: torch.Tensor) -> torch.Tensor: - """ - Compute index scores for sparse attention (BF16 version). - - This is a BF16 implementation of the FP8 index kernel logic: - 1. Compute attention scores: q @ k^T - 2. Apply ReLU activation - 3. Weight by attention weights - 4. Sum across attention heads - - Args: - q: Query tensor [batch, seq_len, n_heads, head_dim] (bf16) - weights: Attention weights [batch, seq_len, n_heads, 1] (bf16) - k: Key tensor [batch, seq_len, head_dim] (bf16) - - Returns: - index_score: Index scores [batch, seq_len, seq_len] (bf16) - - Note: Original FP8 kernel signature was: - fp8_index(q, q_s, k, k_s) where q_s and k_s are scaling factors - Here we use BF16 directly without separate scaling factors. - """ - # q: [bsz, seqlen, n_heads, head_dim] - # k: [bsz, seqlen, head_dim] - # weights: [bsz, seqlen, n_heads, 1] - - # Compute attention scores: q @ k^T - # [bsz, seqlen, n_heads, head_dim] @ [bsz, seqlen, head_dim]^T - # -> [bsz, seqlen, n_heads, seqlen] - index_score = torch.einsum('bshd,btd->bsht', q, k) - - # Apply ReLU activation (for throughput efficiency) - index_score = torch.relu(index_score) - - # Weight each head by attention weights - # [bsz, seqlen, n_heads, seqlen] * [bsz, seqlen, n_heads, 1] - index_score = index_score * weights - - # Sum across attention heads - # [bsz, seqlen, n_heads, seqlen] -> [bsz, seqlen, seqlen] - index_score = index_score.sum(dim=2) - - return index_score - - - @dataclass class IndexerSubmodules: """ @@ -219,7 +228,7 @@ class Indexer(MegatronModule): Computes index scores to identify the top-k most relevant key-value pairs for each query position in sparse attention. - Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py + Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480 Args: config: Transformer configuration. @@ -254,16 +263,15 @@ def __init__( self.rope_head_dim = rope_head_dim self.index_topk = index_topk self.q_lora_rank = q_lora_rank + self.softmax_scale: float = self.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 world_size = pg_collection.tp.size() self.n_local_heads = n_heads // world_size - # Initialize Rotary Position Embedding. - # Use rope_type from config if available, default to "rope". + # Initialize Position Embedding. if self.config.rope_type == "rope": self.rotary_pos_emb = RotaryEmbedding( self.rope_head_dim, @@ -285,10 +293,10 @@ def __init__( ) else: raise ValueError( - f"Unsupported RoPE type: {self.config.rope_type}, supported types are 'rope' and 'yarn'" + f"Unsupported RoPE type: {self.config.rope_type}, supported types are " + "'rope' and 'yarn'" ) - # Build linear layers using build_module self.wq_b = build_module( submodules.linear_wq_b, self.q_lora_rank, @@ -333,28 +341,25 @@ def __init__( parallel_mode="duplicated", ) - self.softmax_scale: float = self.head_dim ** -0.5 - def forward( self, x: torch.Tensor, qr: torch.Tensor, mask: Optional[torch.Tensor] = None, - packed_seq_params: Optional = None, + packed_seq_params: Optional[PackedSeqParams] = None, ): """ Forward pass for Indexer. Args: - x: Input tensor [seq, batch, hidden_dim] - qr: Query representation tensor [seq, batch, q_lora_rank] - mask: Attention mask - packed_seq_params: Packed sequence parameters for variable length sequences + x: hidden states [seqlen, batch, hidden_dim]. + 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, seq_len, index_topk] + topk_indices: Top-k indices for sparse attention [batch, seqlen, index_topk]. """ - # Call forward_with_scores and only return indices _, topk_indices = self.forward_with_scores(x, qr, mask, packed_seq_params) return topk_indices @@ -363,7 +368,7 @@ def forward_with_scores( x: torch.Tensor, qr: torch.Tensor, mask: Optional[torch.Tensor] = None, - packed_seq_params: Optional = None, + packed_seq_params: Optional[PackedSeqParams] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Forward pass for Indexer that returns both index scores and top-k indices. @@ -371,81 +376,100 @@ def forward_with_scores( This is used when KL loss is enabled to compare indexer scores with true attention scores. Args: - x: Input tensor [seq, batch, hidden_dim] - qr: Query representation tensor [seq, batch, q_lora_rank] - mask: Attention mask - packed_seq_params: Packed sequence parameters + x: hidden states [seqlen, batch, hidden_dim]. + 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, seq_len, seq_len] - topk_indices: Top-k indices [batch, seq_len, index_topk] + index_scores: Index scores [batch, seqlen, seqlen] + topk_indices: Top-k indices [batch, seqlen, index_topk] """ - # Input format: [seq, batch, hidden] + assert packed_seq_params is None, "Packed sequence is not supported for SparseAttention" + assert not self.config.apply_rope_fusion, "RoPE fusion is not supported for SparseAttention" + seqlen, bsz, _ = x.size() - # Compute rotary position embeddings internally - packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + # ========================================= + # 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(seqlen, packed_seq=packed_seq) + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) mscale = 1.0 - else: # yarn - rotary_pos_emb, mscale = self.rotary_pos_emb(seqlen, packed_seq=packed_seq) - - q, _ = self.wq_b(qr) # [seq, batch, heads*head_dim] - q = q.reshape(seqlen, bsz, -1, self.head_dim) # [seq, batch, heads, head_dim] - q_pe, q_nope = torch.split( - q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # Apply RoPE to q + # ========================================= + # [seqlen, batch, q_lora_rank] -> [seqlen, batch, n_heads * head_dim] + q, _ = self.wq_b(qr) + # [seqlen, batch, n_heads * head_dim] -> [seqlen, batch, n_heads, head_dim] + q = q.reshape(seqlen, bsz, self.n_heads, self.head_dim) + q_nope, q_pe = torch.split( + q, [self.head_dim - self.rope_head_dim, self.rope_head_dim], dim=-1 ) - q_pe = apply_rotary_pos_emb( q_pe, rotary_pos_emb, config=self.config, - cp_group=self.pg_collection.cp, + cu_seqlens=None, mscale=mscale, + cp_group=self.pg_collection.cp, ) - q = torch.cat([q_pe, q_nope], dim=-1) # [seq, batch, heads, head_dim] - - k, _ = self.wk(x) # [seq, batch, head_dim] - k = self.k_norm(k) # [seq, batch, head_dim] - k_pe, k_nope = torch.split( - k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1 + # [seqlen, batch, n_heads, head_dim] + q = torch.cat([q_nope, q_pe], dim=-1) + + # ========================================= + # Apply RoPE to k + # ========================================= + # [seqlen, batch, hidden_dim] -> [seqlen, batch, head_dim] + k, _ = self.wk(x) + k = self.k_norm(k) + # [seqlen, batch, head_dim] -> [seqlen, batch, 1, head_dim] + k = k.reshape(seqlen, bsz, 1, self.head_dim) + k_nope, k_pe = torch.split( + k, [self.head_dim - self.rope_head_dim, self.rope_head_dim], dim=-1 ) - - k_pe = k_pe.unsqueeze(2) # [seq, batch, 1, rope_head_dim] k_pe = apply_rotary_pos_emb( k_pe, rotary_pos_emb, config=self.config, - cp_group=self.pg_collection.cp, + cu_seqlens=None, mscale=mscale, + cp_group=self.pg_collection.cp, ) - k_pe = k_pe.squeeze(2) # [seq, batch, rope_head_dim] - k = torch.cat([k_pe, k_nope], dim=-1) # [seq, batch, head_dim] - - # For compute_index_score, we need batch-first format - # Transpose to [batch, seq, heads, head_dim] - q = q.transpose(0, 1) # [seq, batch, heads, head_dim] -> [batch, seq, heads, head_dim] - k = k.transpose(0, 1) # [seq, batch, head_dim] -> [batch, seq, head_dim] - + # [seqlen, batch, 1, head_dim] + k = torch.cat([k_nope, k_pe], dim=-1) + # [seqlen, batch, head_dim] + k = k.reshape(seqlen, bsz, self.head_dim) + + # ========================================= + # Rotate activation + # ========================================= q = rotate_activation(q) k = rotate_activation(k) - # weights_proj expects seq-first, so use original x - weights, _ = self.weights_proj(x) # [seq, batch, n_heads] - weights = weights.transpose(0, 1) # -> [batch, seq, n_heads] - weights = weights * self.n_heads ** -0.5 - weights = weights.unsqueeze(-1) * self.softmax_scale - - # Compute index scores (BF16 version of the FP8 kernel) - # All inputs are now batch-first - index_scores = compute_index_score(q.contiguous(), weights, k.contiguous()) - + # ========================================= + # Compute index scores + # ========================================= + # [seqlen, batch, hidden_dim] -> [seqlen, batch, n_heads] + weights, _ = self.weights_proj(x) + weights = weights * (self.n_heads ** -0.5) * self.softmax_scale + # [batcch, seqlen, seqlen] + index_scores = compute_index_score(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 @@ -456,7 +480,7 @@ class SparseAttention(MegatronModule): This module implements sparse attention mechanism using an 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 + Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597 Args: config: Transformer configuration. @@ -514,7 +538,6 @@ def __init__( ) self.softmax_scale = softmax_scale - # Build indexer - required for sparse attention assert submodules.indexer is not None, "Indexer is required for SparseAttention" self.indexer = build_module( submodules.indexer, @@ -544,127 +567,107 @@ def forward( Forward pass for Sparse Attention. Args: - query: Query tensor [sq, b, np, hn] - key: Key tensor [sk, b, np, hn] - value: Value tensor [sk, b, np, hn] - attention_mask: Attention mask tensor - attn_mask_type: Type of attention mask - attention_bias: Optional attention bias - packed_seq_params: Packed sequence parameters - x: Original input hidden states [b, s, h] (needed for indexer, will be inferred if not provided) - qr: Low-rank query representation [b, s, q_lora_rank] (for MLA, will be inferred if not provided) + query: Query tensor [seqlen_q, bsz, n_heads, head_dim]. + key: Key tensor [seqlen_k, bsz, n_heads, head_dim]. + value: Value tensor [seqlen_k, bsz, n_heads, head_dim_v]. + x: Original hidden states [seqlen_q, bsz, hidden_dim]. + qr: Low-rank query representation [seqlen_q, bsz, q_lora_rank]. + attention_mask: Attention mask tensor. + attn_mask_type: Type of attention mask. + attention_bias: Optional attention bias. + packed_seq_params: Packed sequence parameters. Returns: context: Output tensor [sq, b, hp] """ - # Input shape: [sq, b, np, hn] sq, b, np, hn = query.size() sk = key.size(0) - v_hn = value.size(3) # Value head dimension may differ from query/key + # Value head dimension may differ from query/key. + v_hn = value.size(3) + + # Detach x and qr to prevent gradients of indexer from flowing back to the main model. + # TODO(kunlunl): Should x and qr be detached? + x = x.detach() + qr = qr.detach() + + # Get a FP32 mask with -inf for masked positions. + if attention_mask is not None: + mask = attention_mask.squeeze() + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( + mask, float('-inf') + ) + else: + float_mask = None - # =================================== - # Use Indexer for sparse selection (do this first to compute loss before using query) # =================================== # Get index scores and top-k indices - # Note: We need to get index_scores before topk for KL loss computation - # Detach x and qr to prevent gradients from flowing back to the main model - # Indexer is trained solely through the indexer_loss (KL divergence) - # TODO(kunlunl): Should x and qr be detached? + # =================================== index_scores, topk_indices = self.indexer.forward_with_scores( - x.detach(), qr.detach(), mask=None, packed_seq_params=packed_seq_params + x, qr, mask=float_mask, packed_seq_params=packed_seq_params ) # =================================== # Raw attention scores [b, np, sq, sk] # =================================== - output_size = (b, np, sq, sk) - - # Reshape for batch matrix multiplication - # [sq, b, np, hn] -> [b * np, sq, hn] - query_reshaped = query.transpose(0, 1).reshape(b * np, sq, hn) - # [sk, b, np, hn] -> [b * np, sk, hn] - key_reshaped = key.transpose(0, 1).reshape(b * np, sk, hn) - + # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] + query_reshaped = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) + # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] + key_reshaped = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) # Compute attention scores: [b * np, sq, sk] attention_scores = torch.bmm( - query_reshaped, - key_reshaped.transpose(1, 2) + query_reshaped.float(), key_reshaped.float() ) * self.softmax_scale - # Reshape to [b, np, sq, sk] - attention_scores = attention_scores.view(*output_size) + attention_scores = attention_scores.view(b, np, sq, sk) # =================================== # Apply sparse mask from indexer # =================================== - # topk_indices: [b, sq, topk] - - # Step 1: Create sparse selection mask (only allow top-k) - # Start with all positions masked (True = masked) - sparse_mask = torch.ones( - (b, sq, sk), dtype=torch.bool, device=query.device - ) - # Allow top-k positions (False = not masked) - sparse_mask.scatter_(-1, topk_indices, False) - - # Expand to [b, np, sq, sk] - sparse_mask = sparse_mask.unsqueeze(1).expand(-1, np, -1, -1) - - # Step 2: Combine with regular attention mask if provided - if attention_mask is not None: - # attention_mask is boolean: True = masked, False = allowed - # Combine with OR: position is masked if EITHER mask says so - if attention_mask.dtype == torch.bool: - sparse_mask = sparse_mask | attention_mask - else: - # If attention_mask is float (-inf for masked), convert to boolean - sparse_mask = sparse_mask | (attention_mask < -1000.0) - - # Step 3: Apply combined mask to attention scores - attention_scores = attention_scores.masked_fill(sparse_mask, float("-inf")) + # index_mask [b, sq, sk] + index_mask = torch.full((b, sq, sk), float("-inf"), device=x.device) + index_mask.scatter_(-1, topk_indices, 0) + if float_mask is not None: + index_mask += float_mask + attention_scores += index_mask.unsqueeze(1) # =================================== # Attention probabilities [b, np, sq, sk] # =================================== - attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) - attention_probs = attention_probs.to(query.dtype) + attention_probs_fp32 = torch.nn.functional.softmax( + attention_scores, dim=-1, dtype=torch.float32 + ) + attention_probs = attention_probs_fp32.to(query.dtype) # =================================== - # Context layer [sq, b, hp] + # Output # =================================== - # Reshape value: [sk, b, np, v_hn] -> [b * np, sk, v_hn] - value_reshaped = value.transpose(0, 1).reshape(b * np, sk, v_hn) - + # [sk, b, np, v_hn] -> [b, np, sk, v_hn] -> [b * np, sk, v_hn] + value_reshaped = value.permute(1, 2, 0, 3).reshape(b * np, sk, v_hn) # Reshape attention_probs: [b, np, sq, sk] -> [b * np, sq, sk] attention_probs_reshaped = attention_probs.view(b * np, sq, sk) - - # Compute context: [b * np, sq, v_hn] - context = torch.bmm(attention_probs_reshaped, value_reshaped) - - # Reshape context: [b * np, sq, v_hn] -> [b, np, sq, v_hn] -> [sq, b, np, v_hn] - context = context.view(b, np, sq, v_hn).permute(2, 0, 1, 3).contiguous() - - # Flatten: [sq, b, np, v_hn] -> [sq, b, np*v_hn] - context = context.view(sq, b, np * v_hn) + # Compute output: [b * np, sq, v_hn] + output = torch.bmm(attention_probs_reshaped, value_reshaped) + # Reshape output: [b * np, sq, v_hn] -> [b, np, sq, v_hn] -> [sq, b, np, v_hn] + output = output.view(b, np, sq, v_hn).permute(2, 0, 1, 3).contiguous() + # Flatten: [sq, b, np, v_hn] -> [sq, b, np * v_hn] + output = output.view(sq, b, np * v_hn) # =================================== - # Attach indexer loss (training only) + # Attach indexer loss # =================================== if self.training and torch.is_grad_enabled(): # Get indexer loss coefficient from config indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) - # Compute KL divergence loss between indexer scores and true attention scores indexer_loss = compute_indexer_loss( index_scores, - attention_scores.detach(), # Don't backprop through attention scores - indexer_loss_coeff + attention_probs_fp32.detach(), + indexer_loss_coeff, ) + # Attach loss to output output (will trigger backward through indexer) + output = IndexerLossAutoScaler.apply(output, indexer_loss) - # Attach loss to context output (will trigger backward through indexer) - context = IndexerLossAutoScaler.apply(context, indexer_loss) - - return context + return output def sharded_state_dict( self, @@ -673,6 +676,7 @@ def sharded_state_dict( metadata: Optional[dict] = None, ) -> ShardedStateDict: """Sharded state dict for the learnable softmax offset parameter""" + # TODO(kunlunl): Add checkpointing for indexer. if self.config.softmax_type == "learnable": state_dict = self.state_dict(prefix="", keep_vars=True) else: From 186dcdc673fccafed4f121bee95825998790ee13 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 11 Nov 2025 19:54:56 +0800 Subject: [PATCH 05/28] Minor changes about code style Signed-off-by: kunlunl --- megatron/core/transformer/multi_latent_attention.py | 7 ++++++- megatron/core/transformer/sparse_attention.py | 9 --------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 462d6935aa7..762369da0fd 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -242,6 +242,7 @@ def forward( inference_context=inference_context, ) else: + # TODO(kunlunl): Is this a universal usage of sparse attention? query, key, value, q_compressed, _ = self.get_query_key_value_and_compressed_tensors( hidden_states, key_value_states, @@ -287,6 +288,7 @@ def forward( ) else: # For sparse attention, use a specialized forward. + # TODO(kunlunl): Is there a unified interface for sparse attention? core_attn_out = self.core_attention( query, key, @@ -828,8 +830,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po return query, key, value, q_compressed, kv_compressed def get_query_key_value_tensors(self, *args, **kwargs): + query, key, value, q_compressed, kv_compressed = ( + self.get_query_key_value_and_compressed_tensors(self, *args, **kwargs) + ) # Only return query, key and value. - return get_query_key_value_and_compressed_tensors(self, *args, **kwargs)[:3] + return query, key, value def uncompress_kv_from_cache(self, kv_cached): """ diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index c0eb849185e..581239c46d1 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -20,15 +20,6 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint -try: - from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_apply_mla_rope_for_kv, - fused_apply_mla_rope_for_q, - ) -except: - fused_apply_mla_rope_for_kv = None - fused_apply_mla_rope_for_q = None - # TODO(kunlunl): Add third-party fused kernels. try: from fast_hadamard_transform import hadamard_transform From 527140623ec0f5d1a351b40aac873ae3d3665fc0 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 17 Nov 2025 21:33:01 +0800 Subject: [PATCH 06/28] Fix TP Signed-off-by: kunlunl --- .../gpt/sparse_attention_module_specs.py | 2 +- megatron/core/transformer/attention.py | 4 +- .../transformer/multi_latent_attention.py | 5 +- megatron/core/transformer/sparse_attention.py | 497 ++++++++---------- .../core/transformer/transformer_config.py | 8 +- megatron/training/arguments.py | 2 + 6 files changed, 241 insertions(+), 277 deletions(-) diff --git a/megatron/core/models/gpt/sparse_attention_module_specs.py b/megatron/core/models/gpt/sparse_attention_module_specs.py index b6ba992b73e..a81da994518 100644 --- a/megatron/core/models/gpt/sparse_attention_module_specs.py +++ b/megatron/core/models/gpt/sparse_attention_module_specs.py @@ -55,7 +55,7 @@ def get_sparse_attention_module_spec_for_backend( return ModuleSpec( module=SparseAttention, submodules=SparseAttentionSubmodules( - indexer=get_indexer_spec_for_backend(backend, normalization=normalization), + indexer=get_indexer_spec_for_backend(backend, normalization=normalization) ), ) else: diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index ce916c64052..1daca372fe1 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -48,7 +48,9 @@ rearrange = None try: - from flashattn_hopper.flash_attn_interface import _flash_attn_forward + from flashattn_hopper.flash_attn_interface import ( + _flash_attn_forward, + ) from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 762369da0fd..d064338dff0 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -830,8 +830,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po return query, key, value, q_compressed, kv_compressed def get_query_key_value_tensors(self, *args, **kwargs): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ query, key, value, q_compressed, kv_compressed = ( - self.get_query_key_value_and_compressed_tensors(self, *args, **kwargs) + self.get_query_key_value_and_compressed_tensors(*args, **kwargs) ) # Only return query, key and value. return query, key, value diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 581239c46d1..aed7ff1bb29 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -6,7 +6,6 @@ import torch -from megatron.core.dist_checkpointing.mapping import ShardedStateDict from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -14,11 +13,11 @@ ) 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 -from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint # TODO(kunlunl): Add third-party fused kernels. try: @@ -29,66 +28,28 @@ 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 + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L424-L428 Args: - x: Input tensor (must be bfloat16) + x: Input tensor (must be bfloat16). Returns: - Rotated tensor + Rotated tensor. """ assert x.dtype == torch.bfloat16 - assert hadamard_transform is not None, ( - "fast_hadamard_transform is not installed." - ) + 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) - - -def compute_index_score(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, bsz, n_heads, head_dim], the query tensor. - weights : FP32 [seqlen_q, bsz, n_heads], the attention weights. - k : BF16 [seqlen_k, bsz, head_dim], the key tensor. - - Returns: - index_score: FP32 [bsz, seqlen_q, seqlen_k], the index scores. - """ - # Compute attention scores: q @ k^T - # [seqlen_q, bsz, n_heads, head_dim] @ [seqlen_k, bsz, head_dim]^T -> [seqlen_q, bsz, n_heads, seqlen_k] - index_score = torch.einsum('sbhd,tbd->sbht', q.float(), k.float()) - - # Apply ReLU activation. - index_score = torch.relu(index_score) - - # Weight each head by attention weights. - # [seqlen_q, bsz, n_heads, seqlen_k] * [seqlen_q, bsz, n_heads, 1] -> [seqlen_q, bsz, n_heads, seqlen_k] - index_score = index_score * weights.unsqueeze(-1) - - # Sum across attention heads. - # [seqlen_q, bsz, n_heads, seqlen_k] -> [seqlen_q, bsz, seqlen_k] - index_score = index_score.sum(dim=2) - - # Transpose to [bsz, seqlen_q, seqlen_k]. - index_score = index_score.transpose(0, 1) - - return index_score + return hadamard_transform(x, scale=hidden_size**-0.5) def compute_indexer_loss( index_scores: torch.Tensor, + topk_indices: torch.Tensor, attention_scores: torch.Tensor, indexer_loss_coeff: float, + use_sparse_indexer_loss: bool, + pg_collection: ProcessGroupCollection, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -96,19 +57,27 @@ def compute_indexer_loss( 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 + 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 [bsz, seqlen_q, seqlen_k]. - attention_scores: True attention scores from q @ k^T [bsz, heads, seqlen_q, seqlen_k]. + index_scores: Scores predicted by indexer [batch, seqlen_q, seqlen_k]. + topk_indices: Top-k indices [batch, seqlen_q, index_topk]. + attention_scores: True attention scores from q @ k^T [batch, heads, seqlen_q, seqlen_k]. indexer_loss_coeff: Coefficient for the indexer KL divergence loss. + use_sparse_indexer_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). """ # Sum attention scores across heads. - # [bsz, heads, seqlen_q, seqlen_k] -> [bsz, seqlen_q, seqlen_k] + # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] target_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(target_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. @@ -118,9 +87,14 @@ def compute_indexer_loss( index_probs = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) - kl_per_element = ( - target_probs * (torch.log(target_probs + 1e-10) - torch.log(index_probs + 1e-10)) + kl_per_element = target_probs * ( + torch.log(target_probs + 1e-10) - torch.log(index_probs + 1e-10) ) + + if use_sparse_indexer_loss: + sparse_mask = torch.zeros_like(kl_per_element).scatter_(-1, topk_indices, 1) + kl_per_element = kl_per_element * sparse_mask + kl_div = kl_per_element.sum(dim=-1).mean() # Scale by coefficient. @@ -160,7 +134,8 @@ def backward(ctx, grad_output: torch.Tensor): grad_output: The gradient of the output. Returns: - Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled indexer loss gradient. + Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled indexer loss + gradient. """ (indexer_loss,) = ctx.saved_tensors if IndexerLossAutoScaler.main_loss_backward_scale is None: @@ -195,6 +170,7 @@ class IndexerSubmodules: 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 @@ -209,6 +185,7 @@ class SparseAttentionSubmodules: Args: indexer: Indexer module for computing sparse attention indices. """ + indexer: Union[ModuleSpec, type] = None @@ -216,63 +193,50 @@ class Indexer(MegatronModule): """ Lightning Indexer for DeepSeek Sparse Attention. - Computes index scores to identify the top-k most relevant key-value pairs - for each query position in sparse attention. - - Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480 + Computes index scores to identify the top-k most relevant key-value pairs for each query in + sparse attention. - Args: - config: Transformer configuration. - submodules: Indexer submodules specification. - dim: Model hidden dimension. - n_heads: Number of attention heads. - head_dim: Dimension per attention head. - rope_head_dim: Dimension for rotary position embeddings. - index_topk: Number of top-k indices to select. - q_lora_rank: Rank for low-rank query projection. - pg_collection: Process group collection for tensor parallelism. + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480 """ def __init__( self, config: TransformerConfig, submodules: IndexerSubmodules, - dim: int, - n_heads: int, - head_dim: int, - rope_head_dim: int, - index_topk: int, - q_lora_rank: int, - pg_collection: ProcessGroupCollection = None, - ): - super().__init__(config=config) + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + """Initialize the indexer. - self.config = config - self.dim = dim - self.n_heads = n_heads - self.head_dim = head_dim - self.rope_head_dim = rope_head_dim - self.index_topk = index_topk - self.q_lora_rank = q_lora_rank - self.softmax_scale: float = self.head_dim ** -0.5 + Args: + config (TransformerConfig): The configuration for the transformer model. + submodules (IndexerSubmodules): 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 + self.index_n_heads = self.config.index_n_heads + self.index_head_dim = self.config.index_head_dim + self.index_topk = self.config.index_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 - world_size = pg_collection.tp.size() - self.n_local_heads = n_heads // world_size # Initialize Position Embedding. - if self.config.rope_type == "rope": + if self.config.rope_type == 'rope': self.rotary_pos_emb = RotaryEmbedding( - self.rope_head_dim, + 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": + elif self.config.rope_type == 'yarn': self.rotary_pos_emb = YarnRotaryEmbedding( - self.rope_head_dim, + 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, @@ -284,14 +248,14 @@ def __init__( ) else: raise ValueError( - f"Unsupported RoPE type: {self.config.rope_type}, supported types are " - "'rope' and 'yarn'" + f'Unsupported RoPE type: {self.config.rope_type}, supported types are "rope" and ' + f'"yarn"' ) self.wq_b = build_module( submodules.linear_wq_b, self.q_lora_rank, - self.n_heads * self.head_dim, + self.index_n_heads * self.index_head_dim, config=self.config, init_method=self.config.init_method, bias=False, @@ -302,8 +266,8 @@ def __init__( self.wk = build_module( submodules.linear_wk, - self.dim, - self.head_dim, + self.hidden_size, + self.index_head_dim, config=self.config, init_method=self.config.init_method, bias=False, @@ -315,15 +279,15 @@ def __init__( self.k_norm = build_module( submodules.k_norm, config=self.config, - hidden_size=self.head_dim, + hidden_size=self.index_head_dim, eps=self.config.layernorm_epsilon, ) # TODO(kunlunl): The dtype of this module should be torch.get_default_dtype(). self.weights_proj = build_module( submodules.linear_weights_proj, - self.dim, - self.n_heads, + self.hidden_size, + self.index_n_heads, config=self.config, init_method=self.config.init_method, bias=False, @@ -332,27 +296,71 @@ def __init__( parallel_mode="duplicated", ) - def forward( - self, - x: torch.Tensor, - qr: torch.Tensor, - mask: Optional[torch.Tensor] = None, - packed_seq_params: Optional[PackedSeqParams] = None, - ): + for param in self.parameters(): + setattr(param, 'sequence_parallel', self.config.sequence_parallel) + + 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: """ - Forward pass for Indexer. + 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: - x: hidden states [seqlen, batch, hidden_dim]. - 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. + 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: - topk_indices: Top-k indices for sparse attention [batch, seqlen, index_topk]. + index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores. """ - _, topk_indices = self.forward_with_scores(x, qr, mask, packed_seq_params) - return topk_indices + # 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, @@ -367,20 +375,18 @@ def forward_with_scores( This is used when KL loss is enabled to compare indexer scores with true attention scores. Args: - x: hidden states [seqlen, batch, hidden_dim]. + 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] + 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 SparseAttention" assert not self.config.apply_rope_fusion, "RoPE fusion is not supported for SparseAttention" - seqlen, bsz, _ = x.size() - # ========================================= # Prepare RoPE params # ========================================= @@ -394,49 +400,38 @@ def forward_with_scores( rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) # ========================================= - # Apply RoPE to q + # 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, batch, q_lora_rank] -> [seqlen, batch, n_heads * head_dim] + 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.wq_b(qr) - # [seqlen, batch, n_heads * head_dim] -> [seqlen, batch, n_heads, head_dim] - q = q.reshape(seqlen, bsz, self.n_heads, self.head_dim) - q_nope, q_pe = torch.split( - q, [self.head_dim - self.rope_head_dim, self.rope_head_dim], dim=-1 - ) - q_pe = apply_rotary_pos_emb( - q_pe, - rotary_pos_emb, - config=self.config, - cu_seqlens=None, - mscale=mscale, - cp_group=self.pg_collection.cp, - ) - # [seqlen, batch, n_heads, head_dim] - q = torch.cat([q_nope, q_pe], dim=-1) + # [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) # ========================================= - # Apply RoPE to k + # k linear and apply rope to k # ========================================= - # [seqlen, batch, hidden_dim] -> [seqlen, batch, head_dim] + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] k, _ = self.wk(x) k = self.k_norm(k) - # [seqlen, batch, head_dim] -> [seqlen, batch, 1, head_dim] - k = k.reshape(seqlen, bsz, 1, self.head_dim) - k_nope, k_pe = torch.split( - k, [self.head_dim - self.rope_head_dim, self.rope_head_dim], dim=-1 - ) - k_pe = apply_rotary_pos_emb( - k_pe, - rotary_pos_emb, - config=self.config, - cu_seqlens=None, - mscale=mscale, - cp_group=self.pg_collection.cp, - ) - # [seqlen, batch, 1, head_dim] - k = torch.cat([k_nope, k_pe], dim=-1) - # [seqlen, batch, head_dim] - k = k.reshape(seqlen, bsz, self.head_dim) + # [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 @@ -447,11 +442,11 @@ def forward_with_scores( # ========================================= # Compute index scores # ========================================= - # [seqlen, batch, hidden_dim] -> [seqlen, batch, n_heads] + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] weights, _ = self.weights_proj(x) - weights = weights * (self.n_heads ** -0.5) * self.softmax_scale - # [batcch, seqlen, seqlen] - index_scores = compute_index_score(q, weights, k) + 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 @@ -465,26 +460,36 @@ def forward_with_scores( 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 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 + class SparseAttention(MegatronModule): """ This module implements sparse attention mechanism using an 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 - - Args: - config: Transformer configuration. - submodules: Sparse attention submodules specification. - layer_number: Layer number in the model. - attn_mask_type: Type of attention mask. - attention_type: Type of attention. - attention_dropout: Dropout probability for attention weights. - softmax_scale: Scale factor for softmax. - k_channels: Number of channels in key tensor. - v_channels: Number of channels in value tensor. - cp_comm_type: Context parallel communication type. - pg_collection: Process group collection for distributed training. + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597 """ def __init__( @@ -502,26 +507,13 @@ def __init__( pg_collection: ProcessGroupCollection = None, ): super().__init__(config=config) - - self.config: TransformerConfig = config - assert ( self.config.context_parallel_size == 1 ), "Currently context parallelism is not supported by SparseAttention!" - self.layer_number = max(1, layer_number) - self.attn_mask_type = attn_mask_type - self.attention_type = attention_type - - if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp']) - else: - assert hasattr( - pg_collection, 'tp' - ), "SparseAttention pg_collection must have tp process group" - self.pg_collection = pg_collection - - world_size = pg_collection.tp.size() + self.indexer = build_module( + submodules.indexer, config=self.config, pg_collection=pg_collection + ) if softmax_scale is None: softmax_scale = 1.0 / math.sqrt( @@ -529,19 +521,6 @@ def __init__( ) self.softmax_scale = softmax_scale - assert submodules.indexer is not None, "Indexer is required for SparseAttention" - self.indexer = build_module( - submodules.indexer, - config=self.config, - dim=self.config.hidden_size, - n_heads=self.config.index_n_heads, - head_dim=self.config.index_head_dim, - rope_head_dim=self.config.qk_pos_emb_head_dim, - index_topk=self.config.index_topk, - q_lora_rank=self.config.q_lora_rank, - pg_collection=self.pg_collection, - ) - def forward( self, query: torch.Tensor, @@ -558,37 +537,31 @@ def forward( Forward pass for Sparse Attention. Args: - query: Query tensor [seqlen_q, bsz, n_heads, head_dim]. - key: Key tensor [seqlen_k, bsz, n_heads, head_dim]. - value: Value tensor [seqlen_k, bsz, n_heads, head_dim_v]. - x: Original hidden states [seqlen_q, bsz, hidden_dim]. - qr: Low-rank query representation [seqlen_q, bsz, q_lora_rank]. - attention_mask: Attention mask tensor. + 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: - context: Output tensor [sq, b, hp] + output: Output tensor [sq, b, hidden_size] """ sq, b, np, hn = query.size() - sk = key.size(0) - # Value head dimension may differ from query/key. - v_hn = value.size(3) + skv = key.size(0) + hnv = value.size(3) # Detach x and qr to prevent gradients of indexer from flowing back to the main model. - # TODO(kunlunl): Should x and qr be detached? x = x.detach() qr = qr.detach() # Get a FP32 mask with -inf for masked positions. - if attention_mask is not None: - mask = attention_mask.squeeze() - float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( - mask, float('-inf') - ) - else: - float_mask = None + # [b, 1, sq, skv] -> [b, sq, skv] + mask = attention_mask.squeeze() + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(mask, float('-inf')) # =================================== # Get index scores and top-k indices @@ -598,80 +571,60 @@ def forward( ) # =================================== - # Raw attention scores [b, np, sq, sk] + # Raw attention scores [b, np, sq, skv] # =================================== # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] - query_reshaped = query.permute(1, 2, 0, 3).reshape(b * np, sq, hn) - # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] - key_reshaped = key.permute(1, 2, 3, 0).reshape(b * np, hn, sk) - # Compute attention scores: [b * np, sq, sk] - attention_scores = torch.bmm( - query_reshaped.float(), key_reshaped.float() - ) * self.softmax_scale - # Reshape to [b, np, sq, sk] - attention_scores = attention_scores.view(b, np, sq, sk) + 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()) * self.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, sk] - index_mask = torch.full((b, sq, sk), float("-inf"), device=x.device) + # 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) - if float_mask is not None: - index_mask += float_mask + index_mask += float_mask + # [b, np, sq, skv] + [b, 1, sq, skv] -> [b, np, sq, skv] attention_scores += index_mask.unsqueeze(1) # =================================== - # Attention probabilities [b, np, sq, sk] + # Attention probabilities [b, np, sq, skv] # =================================== - attention_probs_fp32 = torch.nn.functional.softmax( - attention_scores, dim=-1, dtype=torch.float32 - ) - attention_probs = attention_probs_fp32.to(query.dtype) + attention_probs = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) # =================================== # Output # =================================== - # [sk, b, np, v_hn] -> [b, np, sk, v_hn] -> [b * np, sk, v_hn] - value_reshaped = value.permute(1, 2, 0, 3).reshape(b * np, sk, v_hn) - # Reshape attention_probs: [b, np, sq, sk] -> [b * np, sq, sk] - attention_probs_reshaped = attention_probs.view(b * np, sq, sk) - # Compute output: [b * np, sq, v_hn] - output = torch.bmm(attention_probs_reshaped, value_reshaped) - # Reshape output: [b * np, sq, v_hn] -> [b, np, sq, v_hn] -> [sq, b, np, v_hn] - output = output.view(b, np, sq, v_hn).permute(2, 0, 1, 3).contiguous() - # Flatten: [sq, b, np, v_hn] -> [sq, b, np * v_hn] - output = output.view(sq, b, np * v_hn) + # [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_probs: [b, np, sq, skv] -> [b * np, sq, skv] + attention_probs_reshaped = attention_probs.reshape(b * np, sq, skv) + # Compute output: [b * np, sq, hnv] + output = torch.bmm(attention_probs_reshaped.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) # =================================== # Attach indexer loss # =================================== if self.training and torch.is_grad_enabled(): - # Get indexer loss coefficient from config - indexer_loss_coeff = getattr(self.config, 'indexer_loss_coeff', 0.0) # Compute KL divergence loss between indexer scores and true attention scores indexer_loss = compute_indexer_loss( index_scores, - attention_probs_fp32.detach(), - indexer_loss_coeff, + topk_indices, + attention_probs.detach(), + getattr(self.config, 'indexer_loss_coeff', 0.0), + getattr(self.config, "use_sparse_indexer_loss", False), + self.indexer.pg_collection, ) - # Attach loss to output output (will trigger backward through indexer) + # Attach loss to output output output = IndexerLossAutoScaler.apply(output, indexer_loss) return output - - def sharded_state_dict( - self, - prefix: str = '', - sharded_offsets: Tuple[Tuple[int, int, int]] = (), - metadata: Optional[dict] = None, - ) -> ShardedStateDict: - """Sharded state dict for the learnable softmax offset parameter""" - # TODO(kunlunl): Add checkpointing for indexer. - if self.config.softmax_type == "learnable": - state_dict = self.state_dict(prefix="", keep_vars=True) - else: - state_dict = {} - return make_sharded_tensors_for_checkpoint( - state_dict, prefix, {'softmax_offset': 0}, sharded_offsets - ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 9a050e71416..0a745eace9a 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -231,12 +231,16 @@ class TransformerConfig(ModelParallelConfig): index_head_dim: Optional[int] = None """Dimension per indexer head. If None, defaults to kv_channels.""" - index_topk: int = 256 + index_topk: Optional[int] = None """Number of top-k tokens to select in sparse attention indexer.""" - indexer_loss_coeff: float = 0.0 + indexer_loss_coeff: Optional[float] = None """Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.""" + use_sparse_indexer_loss: Optional[bool] = None + """Whether to use sparse indexer loss. If True, the indexer loss will be computed using the + top-k indices.""" + #################### # linear attention #################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 95d883ad714..486152b3dbe 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3279,6 +3279,8 @@ def _add_sparse_attention_args(parser): help='Number of top-k tokens to select in sparse attention indexer.') group.add_argument('--indexer-loss-coeff', default=0.0, type=float, help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') + group.add_argument('--use-sparse-indexer-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_linear_attention_args(parser): From 6e6fb50f65a7f5dedea32368fc35da6c1c4f8773 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 18 Nov 2025 19:29:14 +0800 Subject: [PATCH 07/28] Fix attn mask and norm type Signed-off-by: kunlunl --- .../models/gpt/sparse_attention_module_specs.py | 3 +-- megatron/core/transformer/sparse_attention.py | 17 ++++++++++++++--- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/megatron/core/models/gpt/sparse_attention_module_specs.py b/megatron/core/models/gpt/sparse_attention_module_specs.py index a81da994518..3323504fb03 100644 --- a/megatron/core/models/gpt/sparse_attention_module_specs.py +++ b/megatron/core/models/gpt/sparse_attention_module_specs.py @@ -24,13 +24,12 @@ def get_indexer_spec_for_backend( Returns: ModuleSpec for Indexer with appropriate submodules. """ - rms_norm = normalization == "RMSNorm" return ModuleSpec( module=Indexer, submodules=IndexerSubmodules( linear_wq_b=backend.linear(), linear_wk=backend.linear(), - k_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=True), + k_norm=backend.layer_norm(rms_norm=False, for_qk=True), linear_weights_proj=backend.linear(), ), ) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index aed7ff1bb29..d578c5e6285 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -559,9 +559,20 @@ def forward( qr = qr.detach() # Get a FP32 mask with -inf for masked positions. - # [b, 1, sq, skv] -> [b, sq, skv] - mask = attention_mask.squeeze() - float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill(mask, float('-inf')) + 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 = 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 From 237164ee66f232b3024524313c8febf77843eff0 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 18 Nov 2025 19:30:05 +0800 Subject: [PATCH 08/28] Format Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index d578c5e6285..399b2f8db2a 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -565,14 +565,16 @@ def forward( # torch.triu with diagonal=1 creates upper triangular matrix (excluding main diagonal) float_mask = torch.triu( torch.full((sq, skv), float('-inf'), dtype=torch.float32, device=x.device), - diagonal=1 + 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')) + float_mask = torch.zeros_like(mask, dtype=torch.float32).masked_fill( + mask, float('-inf') + ) # =================================== # Get index scores and top-k indices From 72ba916d77f6ee697d6ad76c23a52d49b719f432 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Tue, 18 Nov 2025 19:51:52 +0800 Subject: [PATCH 09/28] Resolve minor comments Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 399b2f8db2a..b53f245b3e5 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -1,5 +1,6 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +import copy import math from dataclasses import dataclass from typing import Optional, Tuple, Union @@ -37,7 +38,9 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: Returns: Rotated tensor. """ - assert x.dtype == torch.bfloat16 + 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) @@ -252,7 +255,7 @@ def __init__( f'"yarn"' ) - self.wq_b = build_module( + self.linear_wq_b = build_module( submodules.linear_wq_b, self.q_lora_rank, self.index_n_heads * self.index_head_dim, @@ -264,7 +267,7 @@ def __init__( parallel_mode="duplicated", ) - self.wk = build_module( + self.linear_wk = build_module( submodules.linear_wk, self.hidden_size, self.index_head_dim, @@ -276,15 +279,17 @@ def __init__( parallel_mode="duplicated", ) + k_norm_config = copy.copy(self.config) + k_norm_config.normalization = "LayerNorm" self.k_norm = build_module( submodules.k_norm, - config=self.config, + config=k_norm_config, hidden_size=self.index_head_dim, eps=self.config.layernorm_epsilon, ) # TODO(kunlunl): The dtype of this module should be torch.get_default_dtype(). - self.weights_proj = build_module( + self.linear_weights_proj = build_module( submodules.linear_weights_proj, self.hidden_size, self.index_n_heads, @@ -415,7 +420,7 @@ def forward_with_scores( # q linear and apply rope to q # ========================================= # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim] - q, _ = self.wq_b(qr) + 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) @@ -425,7 +430,7 @@ def forward_with_scores( # k linear and apply rope to k # ========================================= # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] - k, _ = self.wk(x) + 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) @@ -443,7 +448,7 @@ def forward_with_scores( # Compute index scores # ========================================= # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] - weights, _ = self.weights_proj(x) + 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) From 0e773553b43f85a3f105a243f132654b8dc15d9c Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 19 Nov 2025 04:13:40 +0800 Subject: [PATCH 10/28] Add unit test Signed-off-by: kunlunl --- megatron/core/models/gpt/gpt_layer_specs.py | 7 +- megatron/core/transformer/sparse_attention.py | 3 - .../transformer/test_sparse_attention.py | 1185 +++++++++++++++++ .../test_sparse_attention_checkpoint.py | 346 +++++ 4 files changed, 1535 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/transformer/test_sparse_attention.py create mode 100644 tests/unit_tests/transformer/test_sparse_attention_checkpoint.py diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index dfa7d148a09..2ed7432d14d 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -327,19 +327,20 @@ def get_attention_module_spec_for_backend( if mla_down_proj_use_column_parallel else backend.linear() ) + fuse_norm_and_linear = backend.fuse_layernorm_and_linear() and sparse_attention_type is None linear_q_up_proj = ( backend.column_parallel_layer_norm_linear() - if qk_layernorm and backend.fuse_layernorm_and_linear() + if qk_layernorm and fuse_norm_and_linear else backend.column_parallel_linear() ) linear_kv_up_proj = ( backend.column_parallel_layer_norm_linear() - if qk_layernorm and backend.fuse_layernorm_and_linear() + if qk_layernorm and fuse_norm_and_linear else backend.column_parallel_linear() ) qk_norm = ( backend.layer_norm(rms_norm=rms_norm, for_qk=True) - if qk_layernorm and not backend.fuse_layernorm_and_linear() + if qk_layernorm and not fuse_norm_and_linear else IdentityOp ) attention = ModuleSpec( diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index b53f245b3e5..394f1315e8a 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -301,9 +301,6 @@ def __init__( parallel_mode="duplicated", ) - for param in self.parameters(): - setattr(param, 'sequence_parallel', self.config.sequence_parallel) - 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] diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py new file mode 100644 index 00000000000..ca6b7c24519 --- /dev/null +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -0,0 +1,1185 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +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.sparse_attention import ( + Indexer, + IndexerLossAutoScaler, + IndexerSubmodules, + SparseAttention, + SparseAttentionSubmodules, + compute_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 + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + + +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() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") + 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 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") + 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) + + +class TestComputeIndexerLoss: + """Test compute_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_indexer_loss_shape(self): + """Test that indexer loss returns a scalar.""" + batch_size = 2 + seqlen = 16 + num_heads = 4 + index_topk = 8 + + # Create dummy tensors + index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).cuda() + attention_scores = torch.softmax( + torch.randn(batch_size, num_heads, seqlen, seqlen, dtype=torch.float32).cuda(), dim=-1 + ) + + loss = compute_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + attention_scores=attention_scores, + indexer_loss_coeff=1.0, + use_sparse_indexer_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_indexer_loss_sparse(self): + """Test sparse indexer loss computation.""" + batch_size = 2 + seqlen = 16 + num_heads = 4 + index_topk = 8 + + # Create dummy tensors + index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).cuda() + attention_scores = torch.softmax( + torch.randn(batch_size, num_heads, seqlen, seqlen, dtype=torch.float32).cuda(), dim=-1 + ) + + loss_sparse = compute_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + attention_scores=attention_scores, + indexer_loss_coeff=1.0, + use_sparse_indexer_loss=True, + pg_collection=self.pg_collection, + ) + + loss_dense = compute_indexer_loss( + index_scores=index_scores, + topk_indices=topk_indices, + attention_scores=attention_scores, + indexer_loss_coeff=1.0, + use_sparse_indexer_loss=False, + pg_collection=self.pg_collection, + ) + + # Sparse loss should be different from dense loss + assert loss_sparse != loss_dense + assert loss_sparse >= 0 + assert loss_dense >= 0 + + +class TestIndexerLossAutoScaler: + """Test IndexerLossAutoScaler 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 = IndexerLossAutoScaler.apply(output, indexer_loss) + + assert torch.allclose(result, output) + + @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_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() + IndexerLossAutoScaler.set_loss_scale(scale) + + # Apply the autograd function + result = IndexerLossAutoScaler.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()}" + + +class TestIndexer: + """Test 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=self.index_topk, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + ) + + # Create indexer submodules spec + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + indexer_submodules = IndexerSubmodules( + 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 = Indexer(self.config, indexer_submodules, self.pg_collection) + + yield + Utils.destroy_model_parallel() + + def test_indexer_constructor(self): + """Test indexer initialization.""" + assert isinstance(self.indexer, Indexer) + 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_indexer_forward(self): + """Test indexer forward pass.""" + seq_len = 64 + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + + # Forward pass + topk_indices = self.indexer(x, qr) + + # Check output shape + assert topk_indices.shape == (batch_size, seq_len, min(self.config.index_topk, seq_len)) + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_forward_with_scores(self): + """Test indexer forward pass with scores.""" + seq_len = 16 + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + + # Forward pass with scores + index_scores, topk_indices = self.indexer.forward_with_scores(x, qr) + + # Check output shapes + assert index_scores.shape == (batch_size, seq_len, seq_len) + assert topk_indices.shape == (batch_size, seq_len, min(self.config.index_topk, seq_len)) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_with_mask(self): + """Test indexer with attention mask.""" + seq_len = 16 + batch_size = 2 + + self.indexer.cuda() + + # Create input tensors + 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() + mask = torch.triu( + torch.full((batch_size, seq_len, seq_len), 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(seq_len): + assert torch.all(topk_indices[b, i] <= max(self.index_topk, i)) + + +class TestSparseAttention: + """Test SparseAttention 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=32, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + indexer_loss_coeff=0.1, + use_sparse_indexer_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 = IndexerSubmodules( + 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=Indexer, submodules=indexer_submodules) + sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) + + self.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + self.sparse_attention = SparseAttention( + 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_sparse_attention_constructor(self): + """Test sparse attention initialization.""" + assert isinstance(self.sparse_attention, SparseAttention) + assert hasattr(self.sparse_attention, 'indexer') + assert isinstance(self.sparse_attention.indexer, Indexer) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_sparse_attention_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_sparse_attention_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_sparse_attention_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.index_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 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=32, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + ) + + 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 = IndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + + return Indexer(config, indexer_submodules, pg_collection) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_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_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_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 TestSparseAttentionTensorParallel: + """Test SparseAttention 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=32, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + indexer_loss_coeff=0.1, + use_sparse_indexer_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 = IndexerSubmodules( + 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=Indexer, submodules=indexer_submodules) + sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) + + return SparseAttention( + 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_sparse_attention_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): + assert torch.equal( + param_list[0], param_list[i] + ), f"Indexer 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_sparse_attention_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=False + ) # 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().cpu() + 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() + + # 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}" + + # Compare gradients + # 1. Indexer gradients should be identical + for name, param in sparse_attention_tpn.indexer.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=1e-3, atol=1e-3 + ), f"Indexer gradient {name} mismatch between TP=1 and TP={tensor_model_parallel_size}" + + # 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_sparse_attention_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_sparse_attention_checkpoint.py b/tests/unit_tests/transformer/test_sparse_attention_checkpoint.py new file mode 100644 index 00000000000..3337bbea831 --- /dev/null +++ b/tests/unit_tests/transformer/test_sparse_attention_checkpoint.py @@ -0,0 +1,346 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +import os +from types import SimpleNamespace + +import pytest +import torch + +import megatron.core.parallel_state as parallel_state +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 MegatronModule, TransformerConfig +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.sparse_attention import ( + Indexer, + IndexerSubmodules, + SparseAttention, + SparseAttentionSubmodules, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.training.checkpointing import load_checkpoint, save_checkpoint +from megatron.training.global_vars import set_args +from tests.unit_tests.dist_checkpointing import TempNamedDir +from tests.unit_tests.test_utilities import Utils + + +class MockState: + """Mock optimizer/scheduler state for checkpointing.""" + + def __init__(self, state_dict): + self._state_dict = state_dict + + def state_dict(self): + return self._state_dict + + def load_state_dict(self, state_dict): + self._state_dict = state_dict + + +def create_checkpoint_args(save_dir, load_dir=None): + """Create args for Megatron checkpointing.""" + args = SimpleNamespace() + args.save = save_dir + args.load = load_dir if load_dir is not None else save_dir + args.ckpt_format = 'torch' + args.use_distributed_optimizer = True + args.use_dist_ckpt = False + args.finetune = False + args.no_load_optim = False + args.no_load_rng = False + args.perform_initialization = True + args.bf16 = True + args.pipeline_model_parallel_size = 1 + args.tensor_model_parallel_size = 1 + args.num_layers_per_virtual_pipeline_stage = None + return args + + +class TestIndexerCheckpointing: + """Test checkpoint save and load for Indexer.""" + + def _create_config(self): + """Helper to create MLA config.""" + return MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=4, + 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=32, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + ) + + 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 = IndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_wk=ModuleSpec(module=TELinear), + k_norm=ModuleSpec(module=TENorm), + linear_weights_proj=ModuleSpec(module=TELinear), + ) + + return Indexer(config, indexer_submodules, pg_collection) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_save_load_checkpoint(self, tmp_path_dist_ckpt): + """Test that indexer can be saved and loaded using Megatron checkpointing.""" + 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 = self._create_config() + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + # Create indexer directly (it's already a MegatronModule) + indexer = self._create_indexer(config, pg_collection).cuda() + + # Save original state + original_state_dict = {k: v.clone().cpu() for k, v in indexer.state_dict().items()} + + with TempNamedDir(tmp_path_dist_ckpt / 'test_indexer_mcore_checkpoint') as ckpt_dir: + # Setup args for checkpointing + args = create_checkpoint_args(ckpt_dir) + set_args(args) + + # Save checkpoint using Megatron API + iteration = 100 + optimizer = MockState({"optimizer": "state"}) + opt_param_scheduler = MockState({"scheduler": "state"}) + + save_checkpoint(iteration, [indexer], optimizer, opt_param_scheduler, 0) + + # Verify checkpoint file exists + ckpt_path = ckpt_dir / "iter_0000100" / "mp_rank_00" / "model_optim_rng.pt" + assert os.path.exists(ckpt_path), f"Checkpoint file should exist at {ckpt_path}" + + # Create new indexer with different initialization + new_indexer = self._create_indexer(config, pg_collection).cuda() + new_optimizer = MockState({"optimizer": "dummy"}) + new_opt_param_scheduler = MockState({"scheduler": "dummy"}) + + # Load checkpoint using Megatron API + loaded_iter, _ = load_checkpoint( + [new_indexer], new_optimizer, new_opt_param_scheduler, strict=True + ) + + assert loaded_iter == iteration, f"Loaded iteration should be {iteration}" + + # Verify weights match after loading + for key in original_state_dict: + assert torch.allclose( + original_state_dict[key], + new_indexer.state_dict()[key].cpu(), + rtol=1e-5, + atol=1e-5, + ), f"Loaded weights should match original for {key}" + + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.parametrize("tp_size", [1, 2, 4]) + def test_indexer_checkpoint_with_tp(self, tmp_path_dist_ckpt, tp_size): + """Test indexer checkpoint save/load with Megatron API and tensor parallelism.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config() + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + # Create indexer directly (it's already a MegatronModule) + indexer = self._create_indexer(config, pg_collection).cuda() + + # Get original state on all ranks + original_state_dict = {k: v.clone().cpu() for k, v in indexer.state_dict().items()} + + with TempNamedDir(tmp_path_dist_ckpt / f'test_indexer_mcore_tp{tp_size}') as ckpt_dir: + # Setup args for checkpointing + args = create_checkpoint_args(ckpt_dir) + args.tensor_model_parallel_size = tp_size + set_args(args) + + # Save checkpoint using Megatron API + iteration = 200 + optimizer = MockState({"optimizer": "state"}) + opt_param_scheduler = MockState({"scheduler": "state"}) + + save_checkpoint(iteration, [indexer], optimizer, opt_param_scheduler, 0) + + # Create new indexer with different initialization + new_indexer = self._create_indexer(config, pg_collection).cuda() + new_optimizer = MockState({"optimizer": "dummy"}) + new_opt_param_scheduler = MockState({"scheduler": "dummy"}) + + # Load checkpoint using Megatron API + loaded_iter, _ = load_checkpoint( + [new_indexer], new_optimizer, new_opt_param_scheduler, strict=True + ) + + assert loaded_iter == iteration + + # Verify weights match on all ranks + for key in original_state_dict: + assert torch.allclose( + original_state_dict[key], + new_indexer.state_dict()[key].cpu(), + rtol=1e-5, + atol=1e-5, + ), f"Loaded weights should match original for {key} on TP rank {parallel_state.get_tensor_model_parallel_rank()}" + + # Verify weights are identical across all TP ranks (duplicated) + world_size = torch.distributed.get_world_size() + if world_size > 1: + for key, param in new_indexer.state_dict().items(): + param_list = [torch.zeros_like(param) for _ in range(world_size)] + torch.distributed.all_gather(param_list, param) + for i in range(1, world_size): + assert torch.equal( + param_list[0], param_list[i] + ), f"Parameter {key} should be identical across all ranks after loading" + + Utils.destroy_model_parallel() + + +class TestSparseAttentionCheckpointing: + """Test checkpoint save and load for SparseAttention.""" + + def _create_config(self): + """Helper to create MLA config.""" + return MLATransformerConfig( + num_layers=2, + hidden_size=256, + num_attention_heads=4, + 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, + # Sparse attention specific configs + index_n_heads=8, + index_head_dim=64, + index_topk=32, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + indexer_loss_coeff=0.1, + use_sparse_indexer_loss=False, + ) + + 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 = IndexerSubmodules( + 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=Indexer, submodules=indexer_submodules, params={'config': config} + ) + + sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) + + return SparseAttention( + 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") + @pytest.mark.parametrize("tp_size", [1, 2, 4]) + def test_sparse_attention_checkpoint_with_tp(self, tmp_path_dist_ckpt, tp_size): + """Test sparse attention checkpoint save/load with Megatron API and tensor parallelism.""" + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + config = self._create_config() + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + # Create sparse attention directly (it's already a MegatronModule) + sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() + + # Get original indexer state on all ranks + original_indexer_state = { + k: v.clone().cpu() for k, v in sparse_attention.indexer.state_dict().items() + } + + with TempNamedDir(tmp_path_dist_ckpt / f'test_sparse_attn_mcore_tp{tp_size}') as ckpt_dir: + # Setup args for checkpointing + args = create_checkpoint_args(ckpt_dir) + args.tensor_model_parallel_size = tp_size + set_args(args) + + # Save checkpoint using Megatron API + iteration = 300 + optimizer = MockState({"optimizer": "state"}) + opt_param_scheduler = MockState({"scheduler": "state"}) + + save_checkpoint(iteration, [sparse_attention], optimizer, opt_param_scheduler, 0) + + # Create new sparse attention with different initialization + new_sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() + new_optimizer = MockState({"optimizer": "dummy"}) + new_opt_param_scheduler = MockState({"scheduler": "dummy"}) + + # Load checkpoint using Megatron API + loaded_iter, _ = load_checkpoint( + [new_sparse_attention], new_optimizer, new_opt_param_scheduler, strict=True + ) + + assert loaded_iter == iteration + + # Verify indexer weights match on all ranks + for key in original_indexer_state: + assert torch.allclose( + original_indexer_state[key], + new_sparse_attention.indexer.state_dict()[key].cpu(), + rtol=1e-5, + atol=1e-5, + ), f"Loaded indexer weights should match original for {key} on TP rank {parallel_state.get_tensor_model_parallel_rank()}" + + # Verify indexer weights are identical across all TP ranks (duplicated) + world_size = torch.distributed.get_world_size() + if world_size > 1: + for key, param in new_sparse_attention.indexer.state_dict().items(): + param_list = [torch.zeros_like(param) for _ in range(world_size)] + torch.distributed.all_gather(param_list, param) + for i in range(1, world_size): + assert torch.equal( + param_list[0], param_list[i] + ), f"Indexer parameter {key} should be identical across all ranks after loading" + + Utils.destroy_model_parallel() From 8851f4e6aada895a018f5e5643a2fa4ba4a913d2 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 20 Nov 2025 15:19:31 +0800 Subject: [PATCH 11/28] Fix UT Signed-off-by: kunlunl --- .../transformer/test_sparse_attention.py | 31 +- .../test_sparse_attention_checkpoint.py | 346 ------------------ 2 files changed, 24 insertions(+), 353 deletions(-) delete mode 100644 tests/unit_tests/transformer/test_sparse_attention_checkpoint.py diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index ca6b7c24519..03d3ea380a8 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -892,9 +892,7 @@ def test_sparse_attention_weight_consistency( # All weights should be identical across all GPUs for i in range(1, world_size): - assert torch.equal( - param_list[0], param_list[i] - ), f"Indexer parameter {name} differs between rank 0 and rank {i} (world)" + torch.testing.assert_close(param_list[0], param_list[i], rtol=0, atol=0) Utils.destroy_model_parallel() @@ -964,7 +962,7 @@ def test_sparse_attention_forward_consistency( # Save gradients from TP=1 indexer_tp1_grads = { - name: param.grad.clone().cpu() + name: param.grad.clone() for name, param in sparse_attention_tp1.indexer.named_parameters() if param.grad is not None } @@ -987,6 +985,25 @@ def test_sparse_attention_forward_consistency( 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: @@ -1040,9 +1057,9 @@ def test_sparse_attention_forward_consistency( # 1. Indexer gradients should be identical for name, param in sparse_attention_tpn.indexer.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=1e-3, atol=1e-3 - ), f"Indexer gradient {name} mismatch between TP=1 and TP={tensor_model_parallel_size}" + 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] diff --git a/tests/unit_tests/transformer/test_sparse_attention_checkpoint.py b/tests/unit_tests/transformer/test_sparse_attention_checkpoint.py deleted file mode 100644 index 3337bbea831..00000000000 --- a/tests/unit_tests/transformer/test_sparse_attention_checkpoint.py +++ /dev/null @@ -1,346 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -import os -from types import SimpleNamespace - -import pytest -import torch - -import megatron.core.parallel_state as parallel_state -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 MegatronModule, TransformerConfig -from megatron.core.transformer.enums import AttnMaskType -from megatron.core.transformer.sparse_attention import ( - Indexer, - IndexerSubmodules, - SparseAttention, - SparseAttentionSubmodules, -) -from megatron.core.transformer.transformer_config import MLATransformerConfig -from megatron.training.checkpointing import load_checkpoint, save_checkpoint -from megatron.training.global_vars import set_args -from tests.unit_tests.dist_checkpointing import TempNamedDir -from tests.unit_tests.test_utilities import Utils - - -class MockState: - """Mock optimizer/scheduler state for checkpointing.""" - - def __init__(self, state_dict): - self._state_dict = state_dict - - def state_dict(self): - return self._state_dict - - def load_state_dict(self, state_dict): - self._state_dict = state_dict - - -def create_checkpoint_args(save_dir, load_dir=None): - """Create args for Megatron checkpointing.""" - args = SimpleNamespace() - args.save = save_dir - args.load = load_dir if load_dir is not None else save_dir - args.ckpt_format = 'torch' - args.use_distributed_optimizer = True - args.use_dist_ckpt = False - args.finetune = False - args.no_load_optim = False - args.no_load_rng = False - args.perform_initialization = True - args.bf16 = True - args.pipeline_model_parallel_size = 1 - args.tensor_model_parallel_size = 1 - args.num_layers_per_virtual_pipeline_stage = None - return args - - -class TestIndexerCheckpointing: - """Test checkpoint save and load for Indexer.""" - - def _create_config(self): - """Helper to create MLA config.""" - return MLATransformerConfig( - num_layers=2, - hidden_size=256, - num_attention_heads=4, - 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, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=32, - rope_type='rope', - rotary_base=10000, - rotary_percent=1.0, - ) - - 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 = IndexerSubmodules( - linear_wq_b=ModuleSpec(module=TELinear), - linear_wk=ModuleSpec(module=TELinear), - k_norm=ModuleSpec(module=TENorm), - linear_weights_proj=ModuleSpec(module=TELinear), - ) - - return Indexer(config, indexer_submodules, pg_collection) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_save_load_checkpoint(self, tmp_path_dist_ckpt): - """Test that indexer can be saved and loaded using Megatron checkpointing.""" - 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 = self._create_config() - pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) - - # Create indexer directly (it's already a MegatronModule) - indexer = self._create_indexer(config, pg_collection).cuda() - - # Save original state - original_state_dict = {k: v.clone().cpu() for k, v in indexer.state_dict().items()} - - with TempNamedDir(tmp_path_dist_ckpt / 'test_indexer_mcore_checkpoint') as ckpt_dir: - # Setup args for checkpointing - args = create_checkpoint_args(ckpt_dir) - set_args(args) - - # Save checkpoint using Megatron API - iteration = 100 - optimizer = MockState({"optimizer": "state"}) - opt_param_scheduler = MockState({"scheduler": "state"}) - - save_checkpoint(iteration, [indexer], optimizer, opt_param_scheduler, 0) - - # Verify checkpoint file exists - ckpt_path = ckpt_dir / "iter_0000100" / "mp_rank_00" / "model_optim_rng.pt" - assert os.path.exists(ckpt_path), f"Checkpoint file should exist at {ckpt_path}" - - # Create new indexer with different initialization - new_indexer = self._create_indexer(config, pg_collection).cuda() - new_optimizer = MockState({"optimizer": "dummy"}) - new_opt_param_scheduler = MockState({"scheduler": "dummy"}) - - # Load checkpoint using Megatron API - loaded_iter, _ = load_checkpoint( - [new_indexer], new_optimizer, new_opt_param_scheduler, strict=True - ) - - assert loaded_iter == iteration, f"Loaded iteration should be {iteration}" - - # Verify weights match after loading - for key in original_state_dict: - assert torch.allclose( - original_state_dict[key], - new_indexer.state_dict()[key].cpu(), - rtol=1e-5, - atol=1e-5, - ), f"Loaded weights should match original for {key}" - - Utils.destroy_model_parallel() - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - @pytest.mark.parametrize("tp_size", [1, 2, 4]) - def test_indexer_checkpoint_with_tp(self, tmp_path_dist_ckpt, tp_size): - """Test indexer checkpoint save/load with Megatron API and tensor parallelism.""" - Utils.initialize_model_parallel( - tensor_model_parallel_size=tp_size, pipeline_model_parallel_size=1 - ) - torch.manual_seed(123) - model_parallel_cuda_manual_seed(123) - - config = self._create_config() - pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) - - # Create indexer directly (it's already a MegatronModule) - indexer = self._create_indexer(config, pg_collection).cuda() - - # Get original state on all ranks - original_state_dict = {k: v.clone().cpu() for k, v in indexer.state_dict().items()} - - with TempNamedDir(tmp_path_dist_ckpt / f'test_indexer_mcore_tp{tp_size}') as ckpt_dir: - # Setup args for checkpointing - args = create_checkpoint_args(ckpt_dir) - args.tensor_model_parallel_size = tp_size - set_args(args) - - # Save checkpoint using Megatron API - iteration = 200 - optimizer = MockState({"optimizer": "state"}) - opt_param_scheduler = MockState({"scheduler": "state"}) - - save_checkpoint(iteration, [indexer], optimizer, opt_param_scheduler, 0) - - # Create new indexer with different initialization - new_indexer = self._create_indexer(config, pg_collection).cuda() - new_optimizer = MockState({"optimizer": "dummy"}) - new_opt_param_scheduler = MockState({"scheduler": "dummy"}) - - # Load checkpoint using Megatron API - loaded_iter, _ = load_checkpoint( - [new_indexer], new_optimizer, new_opt_param_scheduler, strict=True - ) - - assert loaded_iter == iteration - - # Verify weights match on all ranks - for key in original_state_dict: - assert torch.allclose( - original_state_dict[key], - new_indexer.state_dict()[key].cpu(), - rtol=1e-5, - atol=1e-5, - ), f"Loaded weights should match original for {key} on TP rank {parallel_state.get_tensor_model_parallel_rank()}" - - # Verify weights are identical across all TP ranks (duplicated) - world_size = torch.distributed.get_world_size() - if world_size > 1: - for key, param in new_indexer.state_dict().items(): - param_list = [torch.zeros_like(param) for _ in range(world_size)] - torch.distributed.all_gather(param_list, param) - for i in range(1, world_size): - assert torch.equal( - param_list[0], param_list[i] - ), f"Parameter {key} should be identical across all ranks after loading" - - Utils.destroy_model_parallel() - - -class TestSparseAttentionCheckpointing: - """Test checkpoint save and load for SparseAttention.""" - - def _create_config(self): - """Helper to create MLA config.""" - return MLATransformerConfig( - num_layers=2, - hidden_size=256, - num_attention_heads=4, - 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, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=32, - rope_type='rope', - rotary_base=10000, - rotary_percent=1.0, - indexer_loss_coeff=0.1, - use_sparse_indexer_loss=False, - ) - - 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 = IndexerSubmodules( - 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=Indexer, submodules=indexer_submodules, params={'config': config} - ) - - sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) - - return SparseAttention( - 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") - @pytest.mark.parametrize("tp_size", [1, 2, 4]) - def test_sparse_attention_checkpoint_with_tp(self, tmp_path_dist_ckpt, tp_size): - """Test sparse attention checkpoint save/load with Megatron API and tensor parallelism.""" - Utils.initialize_model_parallel( - tensor_model_parallel_size=tp_size, pipeline_model_parallel_size=1 - ) - torch.manual_seed(123) - model_parallel_cuda_manual_seed(123) - - config = self._create_config() - pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) - - # Create sparse attention directly (it's already a MegatronModule) - sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() - - # Get original indexer state on all ranks - original_indexer_state = { - k: v.clone().cpu() for k, v in sparse_attention.indexer.state_dict().items() - } - - with TempNamedDir(tmp_path_dist_ckpt / f'test_sparse_attn_mcore_tp{tp_size}') as ckpt_dir: - # Setup args for checkpointing - args = create_checkpoint_args(ckpt_dir) - args.tensor_model_parallel_size = tp_size - set_args(args) - - # Save checkpoint using Megatron API - iteration = 300 - optimizer = MockState({"optimizer": "state"}) - opt_param_scheduler = MockState({"scheduler": "state"}) - - save_checkpoint(iteration, [sparse_attention], optimizer, opt_param_scheduler, 0) - - # Create new sparse attention with different initialization - new_sparse_attention = self._create_sparse_attention(config, pg_collection).cuda() - new_optimizer = MockState({"optimizer": "dummy"}) - new_opt_param_scheduler = MockState({"scheduler": "dummy"}) - - # Load checkpoint using Megatron API - loaded_iter, _ = load_checkpoint( - [new_sparse_attention], new_optimizer, new_opt_param_scheduler, strict=True - ) - - assert loaded_iter == iteration - - # Verify indexer weights match on all ranks - for key in original_indexer_state: - assert torch.allclose( - original_indexer_state[key], - new_sparse_attention.indexer.state_dict()[key].cpu(), - rtol=1e-5, - atol=1e-5, - ), f"Loaded indexer weights should match original for {key} on TP rank {parallel_state.get_tensor_model_parallel_rank()}" - - # Verify indexer weights are identical across all TP ranks (duplicated) - world_size = torch.distributed.get_world_size() - if world_size > 1: - for key, param in new_sparse_attention.indexer.state_dict().items(): - param_list = [torch.zeros_like(param) for _ in range(world_size)] - torch.distributed.all_gather(param_list, param) - for i in range(1, world_size): - assert torch.equal( - param_list[0], param_list[i] - ), f"Indexer parameter {key} should be identical across all ranks after loading" - - Utils.destroy_model_parallel() From 7cdedf17da041a961a4099a8da29382b751e83a2 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 24 Nov 2025 20:24:20 +0800 Subject: [PATCH 12/28] Use hidden_size to initialize indexer when q_lora_rank is None Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 6 +++++- megatron/core/transformer/transformer_config.py | 3 ++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 394f1315e8a..338e419e02a 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -219,7 +219,11 @@ def __init__( 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 + 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.index_n_heads self.index_head_dim = self.config.index_head_dim self.index_topk = self.config.index_topk diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 0a745eace9a..1b9cdcc2bcc 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -940,7 +940,8 @@ def __post_init__(self): if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": raise ValueError("DeepEP backend is only supported with flex token dispatcher.") - logging.warning( + self.moe_flex_dispatcher_backend = "deepep" + warnings.warn( "moe_enable_deepep is deprecated." "Please use --moe-flex-dispatcher-backend=deepep instead." ) From c915e4a8c0f07a79d055ff7a974fbbbbc01edc5e Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 26 Nov 2025 00:27:08 +0800 Subject: [PATCH 13/28] Add fused tilelang kernels Signed-off-by: kunlunl --- .../core/fusions/fused_sparse_attention.py | 649 ++++++++++++++++++ megatron/core/transformer/sparse_attention.py | 159 +++-- .../core/transformer/transformer_config.py | 3 + megatron/training/arguments.py | 2 + .../transformer/test_sparse_attention.py | 31 +- 5 files changed, 790 insertions(+), 54 deletions(-) create mode 100644 megatron/core/fusions/fused_sparse_attention.py diff --git a/megatron/core/fusions/fused_sparse_attention.py b/megatron/core/fusions/fused_sparse_attention.py new file mode 100644 index 00000000000..48f975af1ad --- /dev/null +++ b/megatron/core/fusions/fused_sparse_attention.py @@ -0,0 +1,649 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +from unittest.mock import MagicMock + +import torch + +from megatron.core.utils import null_decorator + +try: + import tilelang + from tilelang import language as T + + HAVE_TILELANG = True +except ImportError: + HAVE_TILELANG = False + +if not HAVE_TILELANG: + tilelang = MagicMock() + tilelang.jit = null_decorator + T = MagicMock() + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_attention_fwd( + heads, + dim, + tail_dim, + topk, + kv_group=1, + sm_scale=None, + is_causal=True, + CP0=True, + block_I=64, + num_stages=2, + threads=256, +): + """ + Forward kernel. + """ + assert dim == tilelang.math.next_power_of_2( + dim + ), f"haven't check padding correctness yet, dim={dim}" + assert tail_dim == tilelang.math.next_power_of_2( + tail_dim + ), f"haven't check padding correctness yet, dim={tail_dim}" + assert is_causal == True, "non-casual is not supported" + assert ( + topk % block_I == 0 + ), "otherwise will load some index=0 thus causing wrong kv to be loaded" + if sm_scale is None: + sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) + else: + sm_scale = sm_scale * 1.44269504 # log2(e) + + batch = T.dynamic("batch") + seq_len = T.dynamic("seq_len") + seq_len_kv = T.dynamic("seq_len_kv") + + head_kv = heads // kv_group + q_shape = [batch, seq_len, heads, dim + tail_dim] + k_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] + v_shape = [batch, seq_len_kv, kv_group, dim] + o_shape = [batch, seq_len, heads, dim] + indices_shape = [batch, seq_len, kv_group, topk] + lse_shape = [batch, seq_len, heads] + indices_dtype = "int32" + dtype = "bfloat16" + accum_dtype = "float" + + G = kv_group + H = head_kv + padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) + if padded_H != H: + assert kv_group == 1, ( + "here we solve the H padding automatically, other wise you should handle Q copy and " + "Output copy with your mask (when kv_group == 1, use g_i * padded_H:(g_i+1) * padded_H " + "would be handled automatically)" + ) + BI = block_I + NI = tilelang.cdiv(topk, block_I) + D = dim + D_tail = tail_dim + + if head_kv > 64: + assert head_kv % 64 == 0, "head_kv should be a multiple of 64" + REPLICATE_H = head_kv // 64 + else: + REPLICATE_H = 1 + + H_per_block = padded_H if REPLICATE_H == 1 else 64 + + @T.prim_func + def main( + Q: T.Tensor(q_shape, dtype), # type: ignore + K: T.Tensor(k_shape, dtype), # type: ignore + V: T.Tensor(v_shape, dtype), # type: ignore + Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore + Output: T.Tensor(o_shape, dtype), # type: ignore + Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore + ): + with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as (bx, by, bz): + Q_shared = T.alloc_shared([H_per_block, D], dtype) + Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) + K_shared = T.alloc_shared([BI, D], dtype) + V_shared = T.alloc_shared([BI, D], dtype) + K_tail_shared = T.alloc_shared([BI, D_tail], dtype) + O_shared = T.alloc_shared([H_per_block, D], dtype) + Lse_shared = T.alloc_shared([H_per_block], accum_dtype) + mask = T.alloc_fragment([BI], "bool") + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) + S_shared = T.alloc_shared([H_per_block, BI], dtype) + sumexp = T.alloc_fragment([H_per_block], accum_dtype) + sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) + alpha = T.alloc_fragment([H_per_block], accum_dtype) + m_i = T.alloc_fragment([H_per_block], accum_dtype) + m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) + + T.fill(acc_o, 0) + T.fill(sumexp, 0) + T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan + + b_i, g_i = by, bz + s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) + q_i = s_i + max_kv_i = q_i + + H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) + H1 = H0 + H_per_block + + T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) + T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) + + for i_i in T.Pipelined(NI, num_stages=num_stages): + + for bi_i in T.Parallel(BI): + mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] <= max_kv_i + + for bi_i, d_i in T.Parallel(BI, D): + K_shared[bi_i, d_i] = K[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] + for bi_i, d_i in T.Parallel(BI, D_tail): + K_tail_shared[bi_i, d_i] = K[ + b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i + ] + + for bi_i, d_i in T.Parallel(BI, D): + V_shared[bi_i, d_i] = V[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) + T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_s, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(m_i, m_i_prev) + T.reduce_max(acc_s, m_i, dim=1, clear=False) + for h_i in T.Parallel(H_per_block): + alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) + T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] + + T.copy(acc_s, S_shared) + T.gemm(S_shared, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + # Rescale + for h_i, d_i in T.Parallel(H_per_block, D): + acc_o[h_i, d_i] /= sumexp[h_i] + for h_i in T.Parallel(H_per_block): + sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale + + T.copy(acc_o, O_shared) + T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) + T.copy(sumexp, Lse_shared) + T.copy(sumexp, Lse[b_i, s_i, H0:H1]) + + return main + + +def sparse_attention_fwd_interface( + q, + k, + v, + indices, + sm_scale=None, + return_p_sum: bool = False, + d_v=512, + block_I=64, + num_stages=2, + threads=256, +): + """ + Forward kernel interface. + """ + assert HAVE_TILELANG, "tilelang is not installed" + + is_casual = True + assert return_p_sum == False, "This kernel file is for fwd only" + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() and indices.is_contiguous() + batch, seq_len, heads, dim_plus_tail_dim = q.shape + _, seq_len_kv, kv_group, _ = k.shape + + dim = d_v + + assert k.shape[-1] == dim_plus_tail_dim + tail_dim = dim_plus_tail_dim - dim + assert k.shape[0] == batch + _, _, _, topk = indices.shape + assert indices.shape == (batch, seq_len, kv_group, topk) + + kernel = sparse_attention_fwd( + heads, + dim, + tail_dim, + topk, + kv_group, + sm_scale, + is_casual, + block_I=block_I, + num_stages=num_stages, + threads=threads, + ) + out, lse = kernel(q, k, v, indices) + return out, lse + + +@tilelang.jit(out_idx=[-1]) +def preprocess(B, S, H, D, block_ND=32, num_stages=5, dtype="bfloat16", accum_dtype="float"): + """ + Preprocess kernel for backward. + """ + assert dtype == "bfloat16" + assert accum_dtype == "float" + shape = [B, S, H, D] + + @T.prim_func + def preprocess_kernel( + O: T.Tensor(shape, dtype), + dO: T.Tensor(shape, dtype), + Delta: T.Tensor([B, S, H], accum_dtype), + ): + with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz): + o = T.alloc_fragment([block_ND, block_ND], accum_dtype) + do = T.alloc_fragment([block_ND, block_ND], accum_dtype) + delta = T.alloc_fragment([block_ND], accum_dtype) + acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) + T.clear(acc) + for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): + T.copy( + O[ + bz, + by * block_ND : (by + 1) * block_ND, + bx, + k * block_ND : (k + 1) * block_ND, + ], + o, + ) + T.copy( + dO[ + bz, + by * block_ND : (by + 1) * block_ND, + bx, + k * block_ND : (k + 1) * block_ND, + ], + do, + ) + for i, j in T.Parallel(block_ND, block_ND): + acc[i, j] += o[i, j] * do[i, j] + T.reduce_sum(acc, delta, 1) + T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) + + return preprocess_kernel + + +@tilelang.jit(out_idx=[-1]) +def postprocess( + B, S_kv, D, D_tail, kv_group=1, block_N=64, threads=128, dtype="bfloat16", accum_dtype="float" +): + """ + Postprocess kernel for backward. + """ + assert dtype == "bfloat16" + assert accum_dtype == "float" + dkv_shape = [B, S_kv, kv_group, D + D_tail] + + @T.prim_func + def postprocess_kernel( + dKV: T.Tensor(dkv_shape, accum_dtype), dKV_out: T.Tensor(dkv_shape, dtype) + ): + with T.Kernel(T.ceildiv(S_kv, block_N), kv_group, B, threads=threads) as (bx, by, bz): + T.copy( + dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], + dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], + ) + + return postprocess_kernel + + +@tilelang.jit( + out_idx=[-3], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def sparse_attention_bwd( + B, + S, + S_kv, + H, + D, + D_tail, + topk, + kv_group=1, + sm_scale=None, + is_causal=True, + block_size=32, + num_stages=0, + threads=256, + indices_dtype="int32", + dtype="bfloat16", + accum_dtype="float", +): + """ + Backward kernel. + """ + assert is_causal == True, 'non-casual is not supported now' + assert ( + topk % block_size == 0 + ), 'otherwise will load some index=0 thus causing wrong kv to be loaded' + assert dtype == "bfloat16" + assert accum_dtype == "float" + assert indices_dtype == "int32" + + if sm_scale is None: + sm_scale = (D + D_tail) ** (-0.5) + sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) + + H_kv = H // kv_group + q_shape = [B, S, H, D + D_tail] + k_shape = [B, S_kv, kv_group, D + D_tail] + v_shape = [B, S_kv, kv_group, D] + o_shape = [B, S, H, D] + indices_shape = [B, S, kv_group, topk] + delta_shape = [B, S, H] + lse_shape = [B, S, H] + assert indices_dtype == "int32" + assert dtype == "bfloat16" + assert accum_dtype == "float" + + H = H_kv + padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) + BS = block_size + NS = tilelang.cdiv(topk, block_size) + + split_store = 2 + + @T.prim_func + def sparse_mla_bwd_kernel( + Q: T.Tensor(q_shape, dtype), + K: T.Tensor(k_shape, dtype), + V: T.Tensor(v_shape, dtype), + dO: T.Tensor(o_shape, dtype), + Indices: T.Tensor(indices_shape, indices_dtype), + Lse: T.Tensor(lse_shape, accum_dtype), + Delta: T.Tensor(delta_shape, accum_dtype), + dQ: T.Tensor(q_shape, dtype), + dK: T.Tensor(k_shape, accum_dtype), + dV: T.Tensor(v_shape, accum_dtype), + ): + with T.Kernel(S, B, kv_group, threads=threads) as (s_i, by, bz): + Q_shared = T.alloc_shared([padded_H, D], dtype) + Q_tail_shared = T.alloc_shared([padded_H, D_tail], dtype) + K_shared = T.alloc_shared([BS, D], dtype) + K_tail_shared = T.alloc_shared([BS, D_tail], dtype) + V_shared = T.alloc_shared([BS, D], dtype) + dO_shared = T.alloc_shared([padded_H, D], dtype) + mask = T.alloc_fragment([BS], "bool") + + P_shared_cast = T.alloc_shared([padded_H, BS], dtype) + dP_shared_cast = T.alloc_shared([padded_H, BS], dtype) + dQ_shared = T.alloc_shared([padded_H, D], dtype) + dQ_tail_shared = T.alloc_shared([padded_H, D_tail], dtype) + + acc_p = T.alloc_fragment([padded_H, BS], accum_dtype) + acc_dp = T.alloc_fragment([padded_H, BS], accum_dtype) + acc_dq = T.alloc_fragment([padded_H, D], accum_dtype) + acc_dq_tail = T.alloc_fragment([padded_H, D_tail], accum_dtype) + acc_dk = T.alloc_fragment([BS, D], accum_dtype) + acc_dk_tail = T.alloc_fragment([BS, D_tail], accum_dtype) + acc_dv = T.alloc_fragment([BS, D], accum_dtype) + acc_dk_shared = T.view(K_shared, shape=[BS // split_store, D], dtype=accum_dtype) + acc_dv_shared = T.view(V_shared, shape=[BS // split_store, D], dtype=accum_dtype) + acc_dk_tail_shared = T.view( + K_tail_shared, shape=[BS // split_store, D_tail], dtype=accum_dtype + ) + + max_kv_i = s_i + + T.copy(Q[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D], Q_shared) + T.copy(Q[by, s_i, bz * padded_H : (bz + 1) * padded_H, D:], Q_tail_shared) + T.copy(dO[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D], dO_shared) + + T.clear(acc_dq) + T.clear(acc_dq_tail) + + T.annotate_layout( + { + dQ_shared: tilelang.layout.make_swizzled_layout(dQ_shared), + dQ_tail_shared: tilelang.layout.make_swizzled_layout(dQ_tail_shared), + } + ) + + # Process each block of indices + for i_i in T.Pipelined(NS, num_stages=num_stages): + # Check which indices are valid + for bi_i in T.Parallel(BS): + mask[bi_i] = Indices[by, s_i, bz, i_i * BS + bi_i] <= max_kv_i + + # Compute attention scores + for h_i, bi_i in T.Parallel(padded_H, BS): + acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) + + # Load KV, V for this block of indices + for bi_i, d_i in T.Parallel(BS, D): + K_shared[bi_i, d_i] = K[by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, d_i] + + T.gemm(Q_shared, K_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + for bi_i, d_i in T.Parallel(BS, D_tail): + K_tail_shared[bi_i, d_i] = K[ + by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, D + d_i + ] + T.gemm( + Q_tail_shared, + K_tail_shared, + acc_p, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + for h_i, bi_i in T.Parallel(padded_H, BS): + acc_p[h_i, bi_i] = T.exp2( + acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 + - Lse[by, s_i, bz * padded_H + h_i] + ) + + T.copy(acc_p, P_shared_cast) + + # Load KV, V for this block of indices + for bi_i, d_i in T.Parallel(BS, D): + V_shared[bi_i, d_i] = V[by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, d_i] + + T.gemm( + dO_shared, + V_shared, + acc_dp, + transpose_B=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + + for h_i, bi_i in T.Parallel(padded_H, BS): + acc_dp[h_i, bi_i] = ( + acc_p[h_i, bi_i] + * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * padded_H + h_i]) + * sm_scale + ) + + T.copy(acc_dp, dP_shared_cast) + T.gemm(dP_shared_cast, K_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) + T.gemm(dP_shared_cast, K_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) + + T.gemm( + dP_shared_cast, + Q_shared, + acc_dk, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + + T.gemm( + P_shared_cast, + dO_shared, + acc_dv, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + clear_accum=True, + ) + + T.clear(acc_dk_tail) + T.gemm( + dP_shared_cast, + Q_tail_shared, + acc_dk_tail, + transpose_A=True, + policy=T.GemmWarpPolicy.FullCol, + ) + + for s in range(split_store): + for bi_i, d_i in T.Parallel(BS, D): + if bi_i < BS // split_store: + acc_dk_shared[bi_i, d_i] = acc_dk[bi_i + s * (BS // split_store), d_i] + if bi_i < BS // split_store: + acc_dv_shared[bi_i, d_i] = acc_dv[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS, D_tail): + if bi_i < BS // split_store: + acc_dk_tail_shared[bi_i, d_i] = acc_dk_tail[ + bi_i + s * (BS // split_store), d_i + ] + + for bi_i, d_i in T.Parallel(BS // split_store, D // 4): + T.atomic_addx4( + dK[ + by, + Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], + bz, + d_i * 4, + ], + acc_dk_shared[bi_i, d_i * 4], + ) + T.atomic_addx4( + dV[ + by, + Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], + bz, + d_i * 4, + ], + acc_dv_shared[bi_i, d_i * 4], + ) + + # Atomically update dKV, dKV_tail tensors + for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): + T.atomic_addx4( + dK[ + by, + Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], + bz, + D + d_i * 4, + ], + acc_dk_tail_shared[bi_i, d_i * 4], + ) + + # Store the accumulated dQ + T.copy(acc_dq, dQ_shared) + T.copy(acc_dq_tail, dQ_tail_shared) + + T.copy(dQ_shared, dQ[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D]) + T.copy(dQ_tail_shared, dQ[by, s_i, bz * padded_H : (bz + 1) * padded_H, D:]) + + return sparse_mla_bwd_kernel + + +def sparse_attention_bwd_interface( + q, + k, + v, + o, + do, + indices, + lse, + sm_scale=None, + is_casual=True, + return_kernel=False, + delta=None, + d_v=512, +): + """ + Backward kernel interface. + """ + assert HAVE_TILELANG, "tilelang is not installed" + + assert q.is_contiguous() + assert k.is_contiguous() + assert v.is_contiguous() + assert indices.is_contiguous() + assert lse.is_contiguous() + B, S, H, dim_plus_tail_dim = q.shape + _, S_kv, kv_group, _ = k.shape + assert k.shape[-1] == dim_plus_tail_dim + assert k.shape[0] == B + # dim should be assigned + D = d_v + + D_tail = dim_plus_tail_dim - D + topk = indices.shape[-1] + assert indices.shape == (B, S, kv_group, topk) + assert lse.shape == (B, S, H) + + # Get kernels + preprocess_kernel = preprocess(B, S, H, D) + + bwd_kernel = sparse_attention_bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual) + postprocess_kernel_k = postprocess(B, S_kv, D, D_tail, kv_group) + postprocess_kernel_v = postprocess(B, S_kv, D, 0, kv_group) + + if delta is None: + delta = preprocess_kernel(o, do) + dk = torch.zeros_like(k, dtype=torch.float32) + dv = torch.zeros_like(v, dtype=torch.float32) + dq = bwd_kernel(q, k, v, do, indices, lse, delta, dk, dv) + dk = postprocess_kernel_k(dk) + dv = postprocess_kernel_v(dv) + + return dq, dk, dv + + +class FusedSparseAttention(torch.autograd.Function): + """ + Fused sparse attention kernel wrapper. + """ + + @staticmethod + def forward(ctx, q, k, v, indices, d_v, sm_scale=None): + """ + Forward pass. + """ + out, lse = sparse_attention_fwd_interface(q, k, v, indices, sm_scale=sm_scale, d_v=d_v) + ctx.save_for_backward(q, k, v, indices, out, lse) + ctx.sm_scale = sm_scale + ctx.d_v = d_v + return out + + @staticmethod + def backward(ctx, do): + """ + Backward pass. + """ + q, k, v, indices, out, lse = ctx.saved_tensors + dq, dk, dv = sparse_attention_bwd_interface( + q, k, v, out, do, indices, lse, sm_scale=ctx.sm_scale, d_v=ctx.d_v + ) + return dq, dk, dv, None, None diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 338e419e02a..b407cc71ed6 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -7,6 +7,7 @@ import torch +from megatron.core.fusions.fused_sparse_attention import FusedSparseAttention from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -20,7 +21,6 @@ from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig -# TODO(kunlunl): Add third-party fused kernels. try: from fast_hadamard_transform import hadamard_transform except ImportError: @@ -49,7 +49,9 @@ def rotate_activation(x: torch.Tensor) -> torch.Tensor: def compute_indexer_loss( index_scores: torch.Tensor, topk_indices: torch.Tensor, - attention_scores: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + softmax_scale: float, indexer_loss_coeff: float, use_sparse_indexer_loss: bool, pg_collection: ProcessGroupCollection, @@ -66,7 +68,9 @@ def compute_indexer_loss( Args: index_scores: Scores predicted by indexer [batch, seqlen_q, seqlen_k]. topk_indices: Top-k indices [batch, seqlen_q, index_topk]. - attention_scores: True attention scores from q @ k^T [batch, heads, seqlen_q, seqlen_k]. + query: Query tensor [seqlen_q, batch, heads, dim]. + key: Key tensor [seqlen_k, batch, heads, dim]. + softmax_scale: Scale coefficient after q @ k^T. indexer_loss_coeff: Coefficient for the indexer KL divergence loss. use_sparse_indexer_loss: bool, whether to use sparse indexer loss. If True, only the topk indices will be used to compute the loss. @@ -75,6 +79,28 @@ def compute_indexer_loss( 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, + ) + # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] + attention_scores += causal_mask.view(1, 1, sq, sk) + # [b, np, sq, sk] -> [b, np, sq, sk] + attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) + # Sum attention scores across heads. # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] target_scores = attention_scores.sum(dim=1) @@ -489,6 +515,83 @@ def forward( return topk_indices +def unfused_sparse_attention_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 + + +def fused_sparse_attention_fn(query, key, value, topk_indices, softmax_scale): + """ + Fused sparse attention implementation. + """ + sq, b, np, _ = query.size() + skv = key.size(0) + hnv = value.size(3) + # [sq, b, ...] -> [b, sq, ...] + query = query.transpose(0, 1).contiguous() + # [skv, b, ...] -> [b, skv, ...] + key = key.transpose(0, 1).contiguous() + # [skv, b, ...] -> [b, skv, ...] + value = value.transpose(0, 1).contiguous() + # [b, s, index_topk] -> [b, s, head, index_topk] + topk_indices = topk_indices.unsqueeze(2).repeat(1, 1, np, 1) + # output [sq, b, np, hnv] + output = FusedSparseAttention.apply(query, key, value, topk_indices, hnv, softmax_scale) + # [b, sq, np, hnv] -> [s, b, np, hnv] + output = output.transpose(0, 1).contiguous() + # [sq, b, np, hnv] -> [sq, b, np * hnv] + output = output.view(sq, b, np * hnv) + return output + + class SparseAttention(MegatronModule): """ This module implements sparse attention mechanism using an Indexer to compute top-k attention @@ -569,6 +672,7 @@ def forward( 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, @@ -590,45 +694,14 @@ def forward( ) # =================================== - # 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()) * self.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) - index_mask += float_mask - # [b, np, sq, skv] + [b, 1, sq, skv] -> [b, np, sq, skv] - attention_scores += index_mask.unsqueeze(1) - - # =================================== - # Attention probabilities [b, np, sq, skv] + # Run sparse attention kernel # =================================== - attention_probs = 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_probs: [b, np, sq, skv] -> [b * np, sq, skv] - attention_probs_reshaped = attention_probs.reshape(b * np, sq, skv) - # Compute output: [b * np, sq, hnv] - output = torch.bmm(attention_probs_reshaped.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) + if getattr(self.config, "use_fused_sparse_attention", False): + output = fused_sparse_attention_fn(query, key, value, topk_indices, self.softmax_scale) + else: + output = unfused_sparse_attention_fn( + query, key, value, topk_indices, self.softmax_scale + ) # =================================== # Attach indexer loss @@ -638,7 +711,9 @@ def forward( indexer_loss = compute_indexer_loss( index_scores, topk_indices, - attention_probs.detach(), + query.detach(), + key.detach(), + self.softmax_scale, getattr(self.config, 'indexer_loss_coeff', 0.0), getattr(self.config, "use_sparse_indexer_loss", False), self.indexer.pg_collection, diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1b9cdcc2bcc..3bd04c2347a 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -241,6 +241,9 @@ class TransformerConfig(ModelParallelConfig): """Whether to use sparse indexer loss. If True, the indexer loss will be computed using the top-k indices.""" + use_fused_sparse_attention: Optional[bool] = None + """Whether to use fused sparse attention implementation.""" + #################### # linear attention #################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 486152b3dbe..49284d65e9d 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3281,6 +3281,8 @@ def _add_sparse_attention_args(parser): help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') group.add_argument('--use-sparse-indexer-loss', action='store_true', help='Use sparse indexer loss. If set, the indexer loss will be computed using the top-k indices.') + group.add_argument('--use-fused-sparse-attention', action='store_true', + help='Use fused sparse attention. If set, the sparse attention will be computed using the fused sparse attention kernel.') return parser def _add_linear_attention_args(parser): diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index 03d3ea380a8..87314752f38 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -82,19 +82,22 @@ def test_indexer_loss_shape(self): batch_size = 2 seqlen = 16 num_heads = 4 + head_dim = 128 index_topk = 8 # Create dummy tensors index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).cuda() - attention_scores = torch.softmax( - torch.randn(batch_size, num_heads, seqlen, seqlen, dtype=torch.float32).cuda(), dim=-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_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, - attention_scores=attention_scores, + query=query, + key=key, + softmax_scale=softmax_scale, indexer_loss_coeff=1.0, use_sparse_indexer_loss=False, pg_collection=self.pg_collection, @@ -110,19 +113,22 @@ def test_indexer_loss_sparse(self): batch_size = 2 seqlen = 16 num_heads = 4 + head_dim = 128 index_topk = 8 # Create dummy tensors index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).cuda() - attention_scores = torch.softmax( - torch.randn(batch_size, num_heads, seqlen, seqlen, dtype=torch.float32).cuda(), dim=-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_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, - attention_scores=attention_scores, + query=query, + key=key, + softmax_scale=softmax_scale, indexer_loss_coeff=1.0, use_sparse_indexer_loss=True, pg_collection=self.pg_collection, @@ -131,7 +137,9 @@ def test_indexer_loss_sparse(self): loss_dense = compute_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, - attention_scores=attention_scores, + query=query, + key=key, + softmax_scale=softmax_scale, indexer_loss_coeff=1.0, use_sparse_indexer_loss=False, pg_collection=self.pg_collection, @@ -909,7 +917,7 @@ def test_sparse_attention_forward_consistency( model_parallel_cuda_manual_seed(123) config_tp1 = self._create_config( - sequence_parallel=False, use_sparse_indexer_loss=False + 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() @@ -1053,8 +1061,7 @@ def test_sparse_attention_forward_consistency( 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}" - # Compare gradients - # 1. Indexer gradients should be identical + # 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( From de64232ce9a7f6c74dde57942fadfe46b4b55140 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 26 Nov 2025 01:33:44 +0800 Subject: [PATCH 14/28] Add indexer loss tracker Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 119 +++++++++++++++++- megatron/training/training.py | 11 ++ 2 files changed, 128 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index b407cc71ed6..0ab1963296c 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -27,6 +27,109 @@ hadamard_transform = None +class IndexerLossLoggingHelper: + """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 = IndexerLossLoggingHelper.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 = IndexerLossLoggingHelper.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 = IndexerLossLoggingHelper.tracker + if "values" not in tracker: + return + values = tracker["values"] + # 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 + ) + + @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. + """ + IndexerLossLoggingHelper.reduce_loss_in_tracker() + tracker = IndexerLossLoggingHelper.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) + + IndexerLossLoggingHelper.clean_loss_in_tracker() + + def rotate_activation(x: torch.Tensor) -> torch.Tensor: """Apply Hadamard rotation activation. Reference: @@ -620,6 +723,9 @@ def __init__( self.config.context_parallel_size == 1 ), "Currently context parallelism is not supported by SparseAttention!" + self.layer_number = layer_number + self.pg_collection = pg_collection + self.indexer = build_module( submodules.indexer, config=self.config, pg_collection=pg_collection ) @@ -708,17 +814,26 @@ def forward( # =================================== 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, 'indexer_loss_coeff', 0.0) indexer_loss = compute_indexer_loss( index_scores, topk_indices, query.detach(), key.detach(), self.softmax_scale, - getattr(self.config, 'indexer_loss_coeff', 0.0), + indexer_loss_coeff, getattr(self.config, "use_sparse_indexer_loss", False), self.indexer.pg_collection, ) - # Attach loss to output output + # Save indexer loss for logging + if indexer_loss_coeff > 0: + IndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + avg_group=self.pg_collection.dp_cp, + ) + # Attach loss to output output = IndexerLossAutoScaler.apply(output, indexer_loss) return output diff --git a/megatron/training/training.py b/megatron/training/training.py index f805dab0f15..83f949408f8 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -90,6 +90,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.sparse_attention import IndexerLossLoggingHelper from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper from megatron.core.parallel_state import ( destroy_global_memory_buffer, @@ -1676,6 +1677,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.indexer_loss_coeff is not None and args.indexer_loss_coeff > 0: + indexer_loss_scale = 1 / get_num_microbatches() + IndexerLossLoggingHelper.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(): snapshot = torch.cuda.memory._snapshot() From 1d6f4d36abfecde5538f56fcce3669822646a673 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Wed, 26 Nov 2025 21:22:34 +0800 Subject: [PATCH 15/28] Fix sparse indexer loss Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 46 +++++++++++++------ .../transformer/test_sparse_attention.py | 36 +++++++++++++-- 2 files changed, 63 insertions(+), 19 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 0ab1963296c..25996fc7aec 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -7,6 +7,7 @@ import torch +from megatron.core import parallel_state from megatron.core.fusions.fused_sparse_attention import FusedSparseAttention from megatron.core.models.common.embeddings import ( RotaryEmbedding, @@ -76,6 +77,10 @@ def reduce_loss_in_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')) @@ -83,6 +88,11 @@ def reduce_loss_in_tracker(): 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( @@ -199,34 +209,42 @@ def compute_indexer_loss( 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 use_sparse_indexer_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] - target_scores = attention_scores.sum(dim=1) + 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(target_scores.contiguous(), group=pg_collection.tp) - + 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. - target_probs = target_scores / target_scores.sum(dim=-1, keepdim=True) - - # Convert index_scores to probabilities with softmax. - index_probs = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + 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 = target_probs * ( - torch.log(target_probs + 1e-10) - torch.log(index_probs + 1e-10) + # kl_per_element [b, sq, sk] + kl_per_element = attention_scores * ( + torch.log(attention_scores + 1e-10) - torch.log(index_scores + 1e-10) ) - if use_sparse_indexer_loss: - sparse_mask = torch.zeros_like(kl_per_element).scatter_(-1, topk_indices, 1) - kl_per_element = kl_per_element * sparse_mask - + # [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. @@ -724,7 +742,6 @@ def __init__( ), "Currently context parallelism is not supported by SparseAttention!" self.layer_number = layer_number - self.pg_collection = pg_collection self.indexer = build_module( submodules.indexer, config=self.config, pg_collection=pg_collection @@ -831,7 +848,6 @@ def forward( loss=indexer_loss, layer_number=self.layer_number, num_layers=self.config.num_layers, - avg_group=self.pg_collection.dp_cp, ) # Attach loss to output output = IndexerLossAutoScaler.apply(output, indexer_loss) diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index 87314752f38..67d2c248081 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -85,9 +85,23 @@ def test_indexer_loss_shape(self): head_dim = 128 index_topk = 8 - # Create dummy tensors + # Create dummy index scores index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() - topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).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 @@ -116,9 +130,23 @@ def test_indexer_loss_sparse(self): head_dim = 128 index_topk = 8 - # Create dummy tensors + # Create dummy index scores index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() - topk_indices = torch.randint(0, seqlen, (batch_size, seqlen, index_topk)).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 From 8495e03e1a4077eebc0ded4c24c21451bc22c156 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 27 Nov 2025 16:09:44 +0800 Subject: [PATCH 16/28] Address minor comments Signed-off-by: kunlunl --- megatron/core/transformer/sparse_attention.py | 4 -- .../core/transformer/transformer_config.py | 6 ++ .../transformer/test_sparse_attention.py | 66 +++++++++++-------- 3 files changed, 46 insertions(+), 30 deletions(-) diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index 25996fc7aec..e13661bf8f9 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -538,7 +538,6 @@ def forward_with_scores( topk_indices: Top-k indices [batch, seqlen, index_topk]. """ assert packed_seq_params is None, "Packed sequence is not supported for SparseAttention" - assert not self.config.apply_rope_fusion, "RoPE fusion is not supported for SparseAttention" # ========================================= # Prepare RoPE params @@ -737,9 +736,6 @@ def __init__( pg_collection: ProcessGroupCollection = None, ): super().__init__(config=config) - assert ( - self.config.context_parallel_size == 1 - ), "Currently context parallelism is not supported by SparseAttention!" self.layer_number = layer_number diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 3bd04c2347a..40d1d3e1eab 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1682,6 +1682,12 @@ def __post_init__(self): f"the number of layers ({self.num_layers})" ) + if self.sparse_attention_type is not None: + assert ( + self.context_parallel_size == 1 + ), "Currently context parallelism is not supported by SparseAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for SparseAttention" + @dataclass class MLATransformerConfig(TransformerConfig): diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index 67d2c248081..ec95a2af08b 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -64,6 +64,7 @@ def test_rotate_activation_dtype_check(self): rotate_activation(x) +@pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) class TestComputeIndexerLoss: """Test compute_indexer_loss function.""" @@ -77,13 +78,13 @@ def setup_method(self): Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_loss_shape(self): + def test_indexer_loss_shape(self, seqlen_and_topk): """Test that indexer loss returns a scalar.""" batch_size = 2 - seqlen = 16 + seqlen = seqlen_and_topk[0] num_heads = 4 head_dim = 128 - index_topk = 8 + index_topk = seqlen_and_topk[1] # Create dummy index scores index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() @@ -122,13 +123,13 @@ def test_indexer_loss_shape(self): assert loss >= 0 # KL divergence should be non-negative @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_loss_sparse(self): + def test_indexer_loss_sparse(self, seqlen_and_topk): """Test sparse indexer loss computation.""" batch_size = 2 - seqlen = 16 + seqlen = seqlen_and_topk[0] num_heads = 4 head_dim = 128 - index_topk = 8 + index_topk = seqlen_and_topk[1] # Create dummy index scores index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() @@ -174,7 +175,10 @@ def test_indexer_loss_sparse(self): ) # Sparse loss should be different from dense loss - assert loss_sparse != loss_dense + if seqlen > index_topk: + assert loss_sparse != loss_dense + else: + assert loss_sparse == loss_dense assert loss_sparse >= 0 assert loss_dense >= 0 @@ -200,7 +204,7 @@ def test_forward_pass(self): result = IndexerLossAutoScaler.apply(output, indexer_loss) - assert torch.allclose(result, output) + 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): @@ -241,6 +245,7 @@ def test_backward_pass(self): ), 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 TestIndexer: """Test Indexer module basic functionality with TP=1.""" @@ -295,7 +300,7 @@ def setup_method(self): yield Utils.destroy_model_parallel() - def test_indexer_constructor(self): + def test_indexer_constructor(self, seqlen): """Test indexer initialization.""" assert isinstance(self.indexer, Indexer) assert self.indexer.hidden_size == 256 @@ -304,58 +309,67 @@ def test_indexer_constructor(self): assert self.indexer.index_topk == 32 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_forward(self): + def test_indexer_forward(self, seqlen): """Test indexer forward pass.""" - seq_len = 64 batch_size = 2 self.indexer.cuda() # Create input tensors - 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() + 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, seq_len, min(self.config.index_topk, seq_len)) + assert topk_indices.shape == (batch_size, seqlen, min(self.config.index_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_indexer_forward_with_scores(self): + def test_indexer_forward_with_scores(self, seqlen): """Test indexer forward pass with scores.""" - seq_len = 16 batch_size = 2 self.indexer.cuda() # Create input tensors - 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() + 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, seq_len, seq_len) - assert topk_indices.shape == (batch_size, seq_len, min(self.config.index_topk, seq_len)) + assert index_scores.shape == (batch_size, seqlen, seqlen) + assert topk_indices.shape == (batch_size, seqlen, min(self.config.index_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_indexer_with_mask(self): + def test_indexer_with_mask(self, seqlen): """Test indexer with attention mask.""" - seq_len = 16 batch_size = 2 self.indexer.cuda() # Create input tensors - 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() + 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, seq_len, seq_len), float('-inf'), dtype=torch.float32).cuda(), + torch.full((batch_size, seqlen, seqlen), float('-inf'), dtype=torch.float32).cuda(), diagonal=1, ) @@ -366,7 +380,7 @@ def test_indexer_with_mask(self): # 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(seq_len): + for i in range(seqlen): assert torch.all(topk_indices[b, i] <= max(self.index_topk, i)) From ad04d583843994d0c9952ef1c469ed5f57fc5ce4 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 27 Nov 2025 16:27:18 +0800 Subject: [PATCH 17/28] Temporarily delete fused kernels Signed-off-by: kunlunl --- .../core/fusions/fused_sparse_attention.py | 649 ------------------ megatron/core/transformer/sparse_attention.py | 32 +- .../core/transformer/transformer_config.py | 3 - megatron/training/arguments.py | 2 - 4 files changed, 1 insertion(+), 685 deletions(-) delete mode 100644 megatron/core/fusions/fused_sparse_attention.py diff --git a/megatron/core/fusions/fused_sparse_attention.py b/megatron/core/fusions/fused_sparse_attention.py deleted file mode 100644 index 48f975af1ad..00000000000 --- a/megatron/core/fusions/fused_sparse_attention.py +++ /dev/null @@ -1,649 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -from unittest.mock import MagicMock - -import torch - -from megatron.core.utils import null_decorator - -try: - import tilelang - from tilelang import language as T - - HAVE_TILELANG = True -except ImportError: - HAVE_TILELANG = False - -if not HAVE_TILELANG: - tilelang = MagicMock() - tilelang.jit = null_decorator - T = MagicMock() - - -@tilelang.jit( - out_idx=[-2, -1], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def sparse_attention_fwd( - heads, - dim, - tail_dim, - topk, - kv_group=1, - sm_scale=None, - is_causal=True, - CP0=True, - block_I=64, - num_stages=2, - threads=256, -): - """ - Forward kernel. - """ - assert dim == tilelang.math.next_power_of_2( - dim - ), f"haven't check padding correctness yet, dim={dim}" - assert tail_dim == tilelang.math.next_power_of_2( - tail_dim - ), f"haven't check padding correctness yet, dim={tail_dim}" - assert is_causal == True, "non-casual is not supported" - assert ( - topk % block_I == 0 - ), "otherwise will load some index=0 thus causing wrong kv to be loaded" - if sm_scale is None: - sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e) - else: - sm_scale = sm_scale * 1.44269504 # log2(e) - - batch = T.dynamic("batch") - seq_len = T.dynamic("seq_len") - seq_len_kv = T.dynamic("seq_len_kv") - - head_kv = heads // kv_group - q_shape = [batch, seq_len, heads, dim + tail_dim] - k_shape = [batch, seq_len_kv, kv_group, dim + tail_dim] - v_shape = [batch, seq_len_kv, kv_group, dim] - o_shape = [batch, seq_len, heads, dim] - indices_shape = [batch, seq_len, kv_group, topk] - lse_shape = [batch, seq_len, heads] - indices_dtype = "int32" - dtype = "bfloat16" - accum_dtype = "float" - - G = kv_group - H = head_kv - padded_H = max(tilelang.math.next_power_of_2(head_kv), 16) - if padded_H != H: - assert kv_group == 1, ( - "here we solve the H padding automatically, other wise you should handle Q copy and " - "Output copy with your mask (when kv_group == 1, use g_i * padded_H:(g_i+1) * padded_H " - "would be handled automatically)" - ) - BI = block_I - NI = tilelang.cdiv(topk, block_I) - D = dim - D_tail = tail_dim - - if head_kv > 64: - assert head_kv % 64 == 0, "head_kv should be a multiple of 64" - REPLICATE_H = head_kv // 64 - else: - REPLICATE_H = 1 - - H_per_block = padded_H if REPLICATE_H == 1 else 64 - - @T.prim_func - def main( - Q: T.Tensor(q_shape, dtype), # type: ignore - K: T.Tensor(k_shape, dtype), # type: ignore - V: T.Tensor(v_shape, dtype), # type: ignore - Indices: T.Tensor(indices_shape, indices_dtype), # type: ignore - Output: T.Tensor(o_shape, dtype), # type: ignore - Lse: T.Tensor(lse_shape, accum_dtype), # type: ignore - ): - with T.Kernel(seq_len * REPLICATE_H, batch, kv_group, threads=threads) as (bx, by, bz): - Q_shared = T.alloc_shared([H_per_block, D], dtype) - Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype) - K_shared = T.alloc_shared([BI, D], dtype) - V_shared = T.alloc_shared([BI, D], dtype) - K_tail_shared = T.alloc_shared([BI, D_tail], dtype) - O_shared = T.alloc_shared([H_per_block, D], dtype) - Lse_shared = T.alloc_shared([H_per_block], accum_dtype) - mask = T.alloc_fragment([BI], "bool") - - acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) - acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) - S_shared = T.alloc_shared([H_per_block, BI], dtype) - sumexp = T.alloc_fragment([H_per_block], accum_dtype) - sumexp_i = T.alloc_fragment([H_per_block], accum_dtype) - alpha = T.alloc_fragment([H_per_block], accum_dtype) - m_i = T.alloc_fragment([H_per_block], accum_dtype) - m_i_prev = T.alloc_fragment([H_per_block], accum_dtype) - - T.fill(acc_o, 0) - T.fill(sumexp, 0) - T.fill(m_i, -(2**30)) # avoid -inf - inf to cause nan - - b_i, g_i = by, bz - s_i = bx if REPLICATE_H == 1 else (bx // REPLICATE_H) - q_i = s_i - max_kv_i = q_i - - H0 = g_i * padded_H + (0 if REPLICATE_H == 1 else (bx % REPLICATE_H) * 64) - H1 = H0 + H_per_block - - T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared) - T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared) - - for i_i in T.Pipelined(NI, num_stages=num_stages): - - for bi_i in T.Parallel(BI): - mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] <= max_kv_i - - for bi_i, d_i in T.Parallel(BI, D): - K_shared[bi_i, d_i] = K[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] - for bi_i, d_i in T.Parallel(BI, D_tail): - K_tail_shared[bi_i, d_i] = K[ - b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i - ] - - for bi_i, d_i in T.Parallel(BI, D): - V_shared[bi_i, d_i] = V[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] - - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) - T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_s, - transpose_B=True, - policy=T.GemmWarpPolicy.FullRow, - ) - T.copy(m_i, m_i_prev) - T.reduce_max(acc_s, m_i, dim=1, clear=False) - for h_i in T.Parallel(H_per_block): - alpha[h_i] = T.exp2((m_i_prev[h_i] - m_i[h_i]) * sm_scale) - for h_i, bi_i in T.Parallel(H_per_block, BI): - acc_s[h_i, bi_i] = T.exp2(acc_s[h_i, bi_i] * sm_scale - m_i[h_i] * sm_scale) - T.reduce_sum(acc_s, sumexp_i, dim=1) # is this a accumulate operator? - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = sumexp[h_i] * alpha[h_i] + sumexp_i[h_i] - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] = acc_o[h_i, d_i] * alpha[h_i] - - T.copy(acc_s, S_shared) - T.gemm(S_shared, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) - - # Rescale - for h_i, d_i in T.Parallel(H_per_block, D): - acc_o[h_i, d_i] /= sumexp[h_i] - for h_i in T.Parallel(H_per_block): - sumexp[h_i] = T.log2(sumexp[h_i]) + m_i[h_i] * sm_scale - - T.copy(acc_o, O_shared) - T.copy(acc_o, Output[b_i, s_i, H0:H1, :]) - T.copy(sumexp, Lse_shared) - T.copy(sumexp, Lse[b_i, s_i, H0:H1]) - - return main - - -def sparse_attention_fwd_interface( - q, - k, - v, - indices, - sm_scale=None, - return_p_sum: bool = False, - d_v=512, - block_I=64, - num_stages=2, - threads=256, -): - """ - Forward kernel interface. - """ - assert HAVE_TILELANG, "tilelang is not installed" - - is_casual = True - assert return_p_sum == False, "This kernel file is for fwd only" - assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() and indices.is_contiguous() - batch, seq_len, heads, dim_plus_tail_dim = q.shape - _, seq_len_kv, kv_group, _ = k.shape - - dim = d_v - - assert k.shape[-1] == dim_plus_tail_dim - tail_dim = dim_plus_tail_dim - dim - assert k.shape[0] == batch - _, _, _, topk = indices.shape - assert indices.shape == (batch, seq_len, kv_group, topk) - - kernel = sparse_attention_fwd( - heads, - dim, - tail_dim, - topk, - kv_group, - sm_scale, - is_casual, - block_I=block_I, - num_stages=num_stages, - threads=threads, - ) - out, lse = kernel(q, k, v, indices) - return out, lse - - -@tilelang.jit(out_idx=[-1]) -def preprocess(B, S, H, D, block_ND=32, num_stages=5, dtype="bfloat16", accum_dtype="float"): - """ - Preprocess kernel for backward. - """ - assert dtype == "bfloat16" - assert accum_dtype == "float" - shape = [B, S, H, D] - - @T.prim_func - def preprocess_kernel( - O: T.Tensor(shape, dtype), - dO: T.Tensor(shape, dtype), - Delta: T.Tensor([B, S, H], accum_dtype), - ): - with T.Kernel(H, T.ceildiv(S, block_ND), B) as (bx, by, bz): - o = T.alloc_fragment([block_ND, block_ND], accum_dtype) - do = T.alloc_fragment([block_ND, block_ND], accum_dtype) - delta = T.alloc_fragment([block_ND], accum_dtype) - acc = T.alloc_fragment([block_ND, block_ND], accum_dtype) - T.clear(acc) - for k in T.Pipelined(T.ceildiv(D, block_ND), num_stages=num_stages): - T.copy( - O[ - bz, - by * block_ND : (by + 1) * block_ND, - bx, - k * block_ND : (k + 1) * block_ND, - ], - o, - ) - T.copy( - dO[ - bz, - by * block_ND : (by + 1) * block_ND, - bx, - k * block_ND : (k + 1) * block_ND, - ], - do, - ) - for i, j in T.Parallel(block_ND, block_ND): - acc[i, j] += o[i, j] * do[i, j] - T.reduce_sum(acc, delta, 1) - T.copy(delta, Delta[bz, by * block_ND : (by + 1) * block_ND, bx]) - - return preprocess_kernel - - -@tilelang.jit(out_idx=[-1]) -def postprocess( - B, S_kv, D, D_tail, kv_group=1, block_N=64, threads=128, dtype="bfloat16", accum_dtype="float" -): - """ - Postprocess kernel for backward. - """ - assert dtype == "bfloat16" - assert accum_dtype == "float" - dkv_shape = [B, S_kv, kv_group, D + D_tail] - - @T.prim_func - def postprocess_kernel( - dKV: T.Tensor(dkv_shape, accum_dtype), dKV_out: T.Tensor(dkv_shape, dtype) - ): - with T.Kernel(T.ceildiv(S_kv, block_N), kv_group, B, threads=threads) as (bx, by, bz): - T.copy( - dKV[bz, bx * block_N : (bx + 1) * block_N, by, :], - dKV_out[bz, bx * block_N : (bx + 1) * block_N, by, :], - ) - - return postprocess_kernel - - -@tilelang.jit( - out_idx=[-3], - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - }, -) -def sparse_attention_bwd( - B, - S, - S_kv, - H, - D, - D_tail, - topk, - kv_group=1, - sm_scale=None, - is_causal=True, - block_size=32, - num_stages=0, - threads=256, - indices_dtype="int32", - dtype="bfloat16", - accum_dtype="float", -): - """ - Backward kernel. - """ - assert is_causal == True, 'non-casual is not supported now' - assert ( - topk % block_size == 0 - ), 'otherwise will load some index=0 thus causing wrong kv to be loaded' - assert dtype == "bfloat16" - assert accum_dtype == "float" - assert indices_dtype == "int32" - - if sm_scale is None: - sm_scale = (D + D_tail) ** (-0.5) - sm_scale_mul_reciprocal_log2 = sm_scale * 1.44269504 # log2(e) - - H_kv = H // kv_group - q_shape = [B, S, H, D + D_tail] - k_shape = [B, S_kv, kv_group, D + D_tail] - v_shape = [B, S_kv, kv_group, D] - o_shape = [B, S, H, D] - indices_shape = [B, S, kv_group, topk] - delta_shape = [B, S, H] - lse_shape = [B, S, H] - assert indices_dtype == "int32" - assert dtype == "bfloat16" - assert accum_dtype == "float" - - H = H_kv - padded_H = max(tilelang.math.next_power_of_2(H_kv), 16) - BS = block_size - NS = tilelang.cdiv(topk, block_size) - - split_store = 2 - - @T.prim_func - def sparse_mla_bwd_kernel( - Q: T.Tensor(q_shape, dtype), - K: T.Tensor(k_shape, dtype), - V: T.Tensor(v_shape, dtype), - dO: T.Tensor(o_shape, dtype), - Indices: T.Tensor(indices_shape, indices_dtype), - Lse: T.Tensor(lse_shape, accum_dtype), - Delta: T.Tensor(delta_shape, accum_dtype), - dQ: T.Tensor(q_shape, dtype), - dK: T.Tensor(k_shape, accum_dtype), - dV: T.Tensor(v_shape, accum_dtype), - ): - with T.Kernel(S, B, kv_group, threads=threads) as (s_i, by, bz): - Q_shared = T.alloc_shared([padded_H, D], dtype) - Q_tail_shared = T.alloc_shared([padded_H, D_tail], dtype) - K_shared = T.alloc_shared([BS, D], dtype) - K_tail_shared = T.alloc_shared([BS, D_tail], dtype) - V_shared = T.alloc_shared([BS, D], dtype) - dO_shared = T.alloc_shared([padded_H, D], dtype) - mask = T.alloc_fragment([BS], "bool") - - P_shared_cast = T.alloc_shared([padded_H, BS], dtype) - dP_shared_cast = T.alloc_shared([padded_H, BS], dtype) - dQ_shared = T.alloc_shared([padded_H, D], dtype) - dQ_tail_shared = T.alloc_shared([padded_H, D_tail], dtype) - - acc_p = T.alloc_fragment([padded_H, BS], accum_dtype) - acc_dp = T.alloc_fragment([padded_H, BS], accum_dtype) - acc_dq = T.alloc_fragment([padded_H, D], accum_dtype) - acc_dq_tail = T.alloc_fragment([padded_H, D_tail], accum_dtype) - acc_dk = T.alloc_fragment([BS, D], accum_dtype) - acc_dk_tail = T.alloc_fragment([BS, D_tail], accum_dtype) - acc_dv = T.alloc_fragment([BS, D], accum_dtype) - acc_dk_shared = T.view(K_shared, shape=[BS // split_store, D], dtype=accum_dtype) - acc_dv_shared = T.view(V_shared, shape=[BS // split_store, D], dtype=accum_dtype) - acc_dk_tail_shared = T.view( - K_tail_shared, shape=[BS // split_store, D_tail], dtype=accum_dtype - ) - - max_kv_i = s_i - - T.copy(Q[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D], Q_shared) - T.copy(Q[by, s_i, bz * padded_H : (bz + 1) * padded_H, D:], Q_tail_shared) - T.copy(dO[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D], dO_shared) - - T.clear(acc_dq) - T.clear(acc_dq_tail) - - T.annotate_layout( - { - dQ_shared: tilelang.layout.make_swizzled_layout(dQ_shared), - dQ_tail_shared: tilelang.layout.make_swizzled_layout(dQ_tail_shared), - } - ) - - # Process each block of indices - for i_i in T.Pipelined(NS, num_stages=num_stages): - # Check which indices are valid - for bi_i in T.Parallel(BS): - mask[bi_i] = Indices[by, s_i, bz, i_i * BS + bi_i] <= max_kv_i - - # Compute attention scores - for h_i, bi_i in T.Parallel(padded_H, BS): - acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) - - # Load KV, V for this block of indices - for bi_i, d_i in T.Parallel(BS, D): - K_shared[bi_i, d_i] = K[by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, d_i] - - T.gemm(Q_shared, K_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) - - for bi_i, d_i in T.Parallel(BS, D_tail): - K_tail_shared[bi_i, d_i] = K[ - by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, D + d_i - ] - T.gemm( - Q_tail_shared, - K_tail_shared, - acc_p, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - ) - - for h_i, bi_i in T.Parallel(padded_H, BS): - acc_p[h_i, bi_i] = T.exp2( - acc_p[h_i, bi_i] * sm_scale_mul_reciprocal_log2 - - Lse[by, s_i, bz * padded_H + h_i] - ) - - T.copy(acc_p, P_shared_cast) - - # Load KV, V for this block of indices - for bi_i, d_i in T.Parallel(BS, D): - V_shared[bi_i, d_i] = V[by, Indices[by, s_i, bz, i_i * BS + bi_i], bz, d_i] - - T.gemm( - dO_shared, - V_shared, - acc_dp, - transpose_B=True, - policy=T.GemmWarpPolicy.FullCol, - clear_accum=True, - ) - - for h_i, bi_i in T.Parallel(padded_H, BS): - acc_dp[h_i, bi_i] = ( - acc_p[h_i, bi_i] - * (acc_dp[h_i, bi_i] - Delta[by, s_i, bz * padded_H + h_i]) - * sm_scale - ) - - T.copy(acc_dp, dP_shared_cast) - T.gemm(dP_shared_cast, K_shared, acc_dq, policy=T.GemmWarpPolicy.FullCol) - T.gemm(dP_shared_cast, K_tail_shared, acc_dq_tail, policy=T.GemmWarpPolicy.FullCol) - - T.gemm( - dP_shared_cast, - Q_shared, - acc_dk, - transpose_A=True, - policy=T.GemmWarpPolicy.FullCol, - clear_accum=True, - ) - - T.gemm( - P_shared_cast, - dO_shared, - acc_dv, - transpose_A=True, - policy=T.GemmWarpPolicy.FullCol, - clear_accum=True, - ) - - T.clear(acc_dk_tail) - T.gemm( - dP_shared_cast, - Q_tail_shared, - acc_dk_tail, - transpose_A=True, - policy=T.GemmWarpPolicy.FullCol, - ) - - for s in range(split_store): - for bi_i, d_i in T.Parallel(BS, D): - if bi_i < BS // split_store: - acc_dk_shared[bi_i, d_i] = acc_dk[bi_i + s * (BS // split_store), d_i] - if bi_i < BS // split_store: - acc_dv_shared[bi_i, d_i] = acc_dv[bi_i + s * (BS // split_store), d_i] - - for bi_i, d_i in T.Parallel(BS, D_tail): - if bi_i < BS // split_store: - acc_dk_tail_shared[bi_i, d_i] = acc_dk_tail[ - bi_i + s * (BS // split_store), d_i - ] - - for bi_i, d_i in T.Parallel(BS // split_store, D // 4): - T.atomic_addx4( - dK[ - by, - Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], - bz, - d_i * 4, - ], - acc_dk_shared[bi_i, d_i * 4], - ) - T.atomic_addx4( - dV[ - by, - Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], - bz, - d_i * 4, - ], - acc_dv_shared[bi_i, d_i * 4], - ) - - # Atomically update dKV, dKV_tail tensors - for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): - T.atomic_addx4( - dK[ - by, - Indices[by, s_i, bz, i_i * BS + bi_i + s * (BS // split_store)], - bz, - D + d_i * 4, - ], - acc_dk_tail_shared[bi_i, d_i * 4], - ) - - # Store the accumulated dQ - T.copy(acc_dq, dQ_shared) - T.copy(acc_dq_tail, dQ_tail_shared) - - T.copy(dQ_shared, dQ[by, s_i, bz * padded_H : (bz + 1) * padded_H, :D]) - T.copy(dQ_tail_shared, dQ[by, s_i, bz * padded_H : (bz + 1) * padded_H, D:]) - - return sparse_mla_bwd_kernel - - -def sparse_attention_bwd_interface( - q, - k, - v, - o, - do, - indices, - lse, - sm_scale=None, - is_casual=True, - return_kernel=False, - delta=None, - d_v=512, -): - """ - Backward kernel interface. - """ - assert HAVE_TILELANG, "tilelang is not installed" - - assert q.is_contiguous() - assert k.is_contiguous() - assert v.is_contiguous() - assert indices.is_contiguous() - assert lse.is_contiguous() - B, S, H, dim_plus_tail_dim = q.shape - _, S_kv, kv_group, _ = k.shape - assert k.shape[-1] == dim_plus_tail_dim - assert k.shape[0] == B - # dim should be assigned - D = d_v - - D_tail = dim_plus_tail_dim - D - topk = indices.shape[-1] - assert indices.shape == (B, S, kv_group, topk) - assert lse.shape == (B, S, H) - - # Get kernels - preprocess_kernel = preprocess(B, S, H, D) - - bwd_kernel = sparse_attention_bwd(B, S, S_kv, H, D, D_tail, topk, kv_group, sm_scale, is_casual) - postprocess_kernel_k = postprocess(B, S_kv, D, D_tail, kv_group) - postprocess_kernel_v = postprocess(B, S_kv, D, 0, kv_group) - - if delta is None: - delta = preprocess_kernel(o, do) - dk = torch.zeros_like(k, dtype=torch.float32) - dv = torch.zeros_like(v, dtype=torch.float32) - dq = bwd_kernel(q, k, v, do, indices, lse, delta, dk, dv) - dk = postprocess_kernel_k(dk) - dv = postprocess_kernel_v(dv) - - return dq, dk, dv - - -class FusedSparseAttention(torch.autograd.Function): - """ - Fused sparse attention kernel wrapper. - """ - - @staticmethod - def forward(ctx, q, k, v, indices, d_v, sm_scale=None): - """ - Forward pass. - """ - out, lse = sparse_attention_fwd_interface(q, k, v, indices, sm_scale=sm_scale, d_v=d_v) - ctx.save_for_backward(q, k, v, indices, out, lse) - ctx.sm_scale = sm_scale - ctx.d_v = d_v - return out - - @staticmethod - def backward(ctx, do): - """ - Backward pass. - """ - q, k, v, indices, out, lse = ctx.saved_tensors - dq, dk, dv = sparse_attention_bwd_interface( - q, k, v, out, do, indices, lse, sm_scale=ctx.sm_scale, d_v=ctx.d_v - ) - return dq, dk, dv, None, None diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index e13661bf8f9..c3a4f75fc8e 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -8,7 +8,6 @@ import torch from megatron.core import parallel_state -from megatron.core.fusions.fused_sparse_attention import FusedSparseAttention from megatron.core.models.common.embeddings import ( RotaryEmbedding, YarnRotaryEmbedding, @@ -688,30 +687,6 @@ def unfused_sparse_attention_fn(query, key, value, topk_indices, softmax_scale): return output -def fused_sparse_attention_fn(query, key, value, topk_indices, softmax_scale): - """ - Fused sparse attention implementation. - """ - sq, b, np, _ = query.size() - skv = key.size(0) - hnv = value.size(3) - # [sq, b, ...] -> [b, sq, ...] - query = query.transpose(0, 1).contiguous() - # [skv, b, ...] -> [b, skv, ...] - key = key.transpose(0, 1).contiguous() - # [skv, b, ...] -> [b, skv, ...] - value = value.transpose(0, 1).contiguous() - # [b, s, index_topk] -> [b, s, head, index_topk] - topk_indices = topk_indices.unsqueeze(2).repeat(1, 1, np, 1) - # output [sq, b, np, hnv] - output = FusedSparseAttention.apply(query, key, value, topk_indices, hnv, softmax_scale) - # [b, sq, np, hnv] -> [s, b, np, hnv] - output = output.transpose(0, 1).contiguous() - # [sq, b, np, hnv] -> [sq, b, np * hnv] - output = output.view(sq, b, np * hnv) - return output - - class SparseAttention(MegatronModule): """ This module implements sparse attention mechanism using an Indexer to compute top-k attention @@ -815,12 +790,7 @@ def forward( # =================================== # Run sparse attention kernel # =================================== - if getattr(self.config, "use_fused_sparse_attention", False): - output = fused_sparse_attention_fn(query, key, value, topk_indices, self.softmax_scale) - else: - output = unfused_sparse_attention_fn( - query, key, value, topk_indices, self.softmax_scale - ) + output = unfused_sparse_attention_fn(query, key, value, topk_indices, self.softmax_scale) # =================================== # Attach indexer loss diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 40d1d3e1eab..25fbab1c0ff 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -241,9 +241,6 @@ class TransformerConfig(ModelParallelConfig): """Whether to use sparse indexer loss. If True, the indexer loss will be computed using the top-k indices.""" - use_fused_sparse_attention: Optional[bool] = None - """Whether to use fused sparse attention implementation.""" - #################### # linear attention #################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 49284d65e9d..486152b3dbe 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3281,8 +3281,6 @@ def _add_sparse_attention_args(parser): help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') group.add_argument('--use-sparse-indexer-loss', action='store_true', help='Use sparse indexer loss. If set, the indexer loss will be computed using the top-k indices.') - group.add_argument('--use-fused-sparse-attention', action='store_true', - help='Use fused sparse attention. If set, the sparse attention will be computed using the fused sparse attention kernel.') return parser def _add_linear_attention_args(parser): From 9dae1ce4228723a751744d12145d451c79aad2d7 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 27 Nov 2025 16:40:48 +0800 Subject: [PATCH 18/28] Fix lint error Signed-off-by: kunlunl --- megatron/core/transformer/attention.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 1daca372fe1..ce916c64052 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -48,9 +48,7 @@ rearrange = None try: - from flashattn_hopper.flash_attn_interface import ( - _flash_attn_forward, - ) + from flashattn_hopper.flash_attn_interface import _flash_attn_forward from flashattn_hopper.flash_attn_interface import ( flash_attn_with_kvcache as flash_attn3_with_kvcache, ) From 9cc697bc249e636c61ad7f842a50612a4f6f3cac Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 27 Nov 2025 18:10:54 +0800 Subject: [PATCH 19/28] Make variable/class names more specific (SparseAttention -> DSA) Signed-off-by: kunlunl --- .../gpt/sparse_attention_module_specs.py | 24 +-- megatron/core/transformer/sparse_attention.py | 134 ++++++------- .../core/transformer/transformer_config.py | 24 +-- megatron/training/arguments.py | 10 +- megatron/training/training.py | 6 +- .../transformer/test_sparse_attention.py | 177 +++++++++--------- 6 files changed, 189 insertions(+), 186 deletions(-) diff --git a/megatron/core/models/gpt/sparse_attention_module_specs.py b/megatron/core/models/gpt/sparse_attention_module_specs.py index 3323504fb03..596690011d1 100644 --- a/megatron/core/models/gpt/sparse_attention_module_specs.py +++ b/megatron/core/models/gpt/sparse_attention_module_specs.py @@ -4,29 +4,29 @@ from megatron.core.models.backends import BackendSpecProvider from megatron.core.transformer.sparse_attention import ( - Indexer, - IndexerSubmodules, - SparseAttention, - SparseAttentionSubmodules, + DSAIndexer, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, ) from megatron.core.transformer.spec_utils import ModuleSpec -def get_indexer_spec_for_backend( +def get_dsa_indexer_spec_for_backend( backend: BackendSpecProvider, normalization: Optional[str] = None ) -> ModuleSpec: - """Helper function to get Indexer module spec for a given backend. + """Helper function to get DSA Indexer module spec for a given backend. Args: backend: Backend specification provider (TE or Local). normalization: Normalization type ("RMSNorm" or None for LayerNorm). Returns: - ModuleSpec for Indexer with appropriate submodules. + ModuleSpec for DSA Indexer with appropriate submodules. """ return ModuleSpec( - module=Indexer, - submodules=IndexerSubmodules( + module=DSAIndexer, + submodules=DSAIndexerSubmodules( linear_wq_b=backend.linear(), linear_wk=backend.linear(), k_norm=backend.layer_norm(rms_norm=False, for_qk=True), @@ -52,9 +52,9 @@ def get_sparse_attention_module_spec_for_backend( # Because TransformerEngine does not support sparse attention yet, we use local # implementation whether the backend is TransformerEngine or not. return ModuleSpec( - module=SparseAttention, - submodules=SparseAttentionSubmodules( - indexer=get_indexer_spec_for_backend(backend, normalization=normalization) + module=DSAttention, + submodules=DSAttentionSubmodules( + indexer=get_dsa_indexer_spec_for_backend(backend, normalization=normalization) ), ) else: diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/sparse_attention.py index c3a4f75fc8e..95e012121a8 100644 --- a/megatron/core/transformer/sparse_attention.py +++ b/megatron/core/transformer/sparse_attention.py @@ -27,7 +27,26 @@ hadamard_transform = None -class IndexerLossLoggingHelper: +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 = {} @@ -53,7 +72,7 @@ def save_loss_to_tracker( if layer_number is None: return - tracker = IndexerLossLoggingHelper.tracker + 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() @@ -63,7 +82,7 @@ def save_loss_to_tracker( @staticmethod def clean_loss_in_tracker(): """Clear the indexer losses.""" - tracker = IndexerLossLoggingHelper.tracker + tracker = DSAIndexerLossLoggingHelper.tracker if "values" in tracker: tracker["values"].zero_() tracker["reduce_group"] = None @@ -72,7 +91,7 @@ def clean_loss_in_tracker(): @staticmethod def reduce_loss_in_tracker(): """Collect and reduce the indexer losses across ranks.""" - tracker = IndexerLossLoggingHelper.tracker + tracker = DSAIndexerLossLoggingHelper.tracker if "values" not in tracker: return values = tracker["values"] @@ -112,8 +131,8 @@ def track_indexer_metrics( total_loss_dict: Dictionary to accumulate total losses. per_layer_logging: Whether to log per-layer losses. """ - IndexerLossLoggingHelper.reduce_loss_in_tracker() - tracker = IndexerLossLoggingHelper.tracker + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + tracker = DSAIndexerLossLoggingHelper.tracker if "values" not in tracker: return @@ -136,36 +155,17 @@ def track_indexer_metrics( if wandb_writer is not None: wandb_writer.log({"indexer loss": avg_indexer_loss}, iteration) - IndexerLossLoggingHelper.clean_loss_in_tracker() - - -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) + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() -def compute_indexer_loss( +def compute_dsa_indexer_loss( index_scores: torch.Tensor, topk_indices: torch.Tensor, query: torch.Tensor, key: torch.Tensor, softmax_scale: float, - indexer_loss_coeff: float, - use_sparse_indexer_loss: bool, + loss_coeff: float, + sparse_loss: bool, pg_collection: ProcessGroupCollection, ) -> torch.Tensor: """ @@ -183,8 +183,8 @@ def compute_indexer_loss( query: Query tensor [seqlen_q, batch, heads, dim]. key: Key tensor [seqlen_k, batch, heads, dim]. softmax_scale: Scale coefficient after q @ k^T. - indexer_loss_coeff: Coefficient for the indexer KL divergence loss. - use_sparse_indexer_loss: bool, whether to use sparse indexer loss. If True, only the topk + 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. @@ -215,7 +215,7 @@ def compute_indexer_loss( # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] attention_scores += causal_mask.view(1, 1, sq, sk) - if use_sparse_indexer_loss: + 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] @@ -247,12 +247,12 @@ def compute_indexer_loss( kl_div = kl_per_element.sum(dim=-1).mean() # Scale by coefficient. - indexer_loss = kl_div * indexer_loss_coeff + indexer_loss = kl_div * loss_coeff return indexer_loss -class IndexerLossAutoScaler(torch.autograd.Function): +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 @@ -287,11 +287,11 @@ def backward(ctx, grad_output: torch.Tensor): gradient. """ (indexer_loss,) = ctx.saved_tensors - if IndexerLossAutoScaler.main_loss_backward_scale is None: - IndexerLossAutoScaler.main_loss_backward_scale = torch.tensor( + 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 = IndexerLossAutoScaler.main_loss_backward_scale + 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 @@ -302,16 +302,16 @@ def set_loss_scale(scale: torch.Tensor): Args: scale: The scale value to set. """ - if IndexerLossAutoScaler.main_loss_backward_scale is None: - IndexerLossAutoScaler.main_loss_backward_scale = scale + if DSAIndexerLossAutoScaler.main_loss_backward_scale is None: + DSAIndexerLossAutoScaler.main_loss_backward_scale = scale else: - IndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) + DSAIndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) @dataclass -class IndexerSubmodules: +class DSAIndexerSubmodules: """ - Configuration class for specifying the submodules of an Indexer. + Configuration class for specifying the submodules of an DSA Indexer. Args: linear_wq_b: Linear projection for query bottleneck expansion. @@ -327,20 +327,20 @@ class IndexerSubmodules: @dataclass -class SparseAttentionSubmodules: +class DSAttentionSubmodules: """ - Configuration class for specifying the submodules of SparseAttention. + Configuration class for specifying the submodules of DSAttention. Args: - indexer: Indexer module for computing sparse attention indices. + indexer: DSA Indexer module for computing sparse attention indices. """ indexer: Union[ModuleSpec, type] = None -class Indexer(MegatronModule): +class DSAIndexer(MegatronModule): """ - Lightning Indexer for DeepSeek Sparse Attention. + 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. @@ -352,14 +352,14 @@ class Indexer(MegatronModule): def __init__( self, config: TransformerConfig, - submodules: IndexerSubmodules, + submodules: DSAIndexerSubmodules, pg_collection: Optional[ProcessGroupCollection] = None, ) -> None: """Initialize the indexer. Args: config (TransformerConfig): The configuration for the transformer model. - submodules (IndexerSubmodules): Indexer submodules specification. + submodules (DSAIndexerSubmodules): Indexer submodules specification. pg_collection (ProcessGroupCollection, optional): Process groups for the indexer. """ super().__init__(config=config) @@ -370,9 +370,11 @@ def __init__( if self.config.q_lora_rank is not None else self.config.hidden_size ) - self.index_n_heads = self.config.index_n_heads - self.index_head_dim = self.config.index_head_dim - self.index_topk = self.config.index_topk + + 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: @@ -522,7 +524,7 @@ def forward_with_scores( packed_seq_params: Optional[PackedSeqParams] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ - Forward pass for Indexer that returns both index scores and top-k indices. + 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. @@ -536,7 +538,7 @@ def forward_with_scores( 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 SparseAttention" + assert packed_seq_params is None, "Packed sequence is not supported for DSAttention" # ========================================= # Prepare RoPE params @@ -619,7 +621,7 @@ def forward( packed_seq_params: Optional[PackedSeqParams] = None, ): """ - Forward pass for Indexer. + Forward pass for DSA Indexer. Args: x: hidden states [seqlen, batch, hidden_size]. @@ -634,7 +636,7 @@ def forward( return topk_indices -def unfused_sparse_attention_fn(query, key, value, topk_indices, softmax_scale): +def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): """ Unfused sparse attention implementation. """ @@ -687,10 +689,10 @@ def unfused_sparse_attention_fn(query, key, value, topk_indices, softmax_scale): return output -class SparseAttention(MegatronModule): +class DSAttention(MegatronModule): """ - This module implements sparse attention mechanism using an Indexer to compute top-k attention - indices for reducing computational complexity. + 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 @@ -699,7 +701,7 @@ class SparseAttention(MegatronModule): def __init__( self, config: TransformerConfig, - submodules: SparseAttentionSubmodules, + submodules: DSAttentionSubmodules, layer_number: int, attn_mask_type: AttnMaskType, attention_type: str, @@ -790,32 +792,32 @@ def forward( # =================================== # Run sparse attention kernel # =================================== - output = unfused_sparse_attention_fn(query, key, value, topk_indices, self.softmax_scale) + 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, 'indexer_loss_coeff', 0.0) - indexer_loss = compute_indexer_loss( + 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, "use_sparse_indexer_loss", False), + getattr(self.config, "dsa_indexer_use_sparse_loss", False), self.indexer.pg_collection, ) # Save indexer loss for logging if indexer_loss_coeff > 0: - IndexerLossLoggingHelper.save_loss_to_tracker( + DSAIndexerLossLoggingHelper.save_loss_to_tracker( loss=indexer_loss, layer_number=self.layer_number, num_layers=self.config.num_layers, ) # Attach loss to output - output = IndexerLossAutoScaler.apply(output, indexer_loss) + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) return output diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 25fbab1c0ff..a514a0446fe 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -225,20 +225,20 @@ class TransformerConfig(ModelParallelConfig): sparse_attention_type: Optional[str] = None """Type of sparse attention to use. Currently only supports dsa (DeepSeek Sparse Attention).""" - index_n_heads: Optional[int] = None - """Number of indexer heads for sparse attention. If None, defaults to num_attention_heads.""" + dsa_indexer_n_heads: Optional[int] = None + """Number of DSA indexer heads for DSA.""" - index_head_dim: Optional[int] = None - """Dimension per indexer head. If None, defaults to kv_channels.""" + dsa_indexer_head_dim: Optional[int] = None + """Dimension per DSA indexer head.""" - index_topk: Optional[int] = None - """Number of top-k tokens to select in sparse attention indexer.""" + dsa_indexer_topk: Optional[int] = None + """Number of top-k tokens to select in DSA indexer.""" - indexer_loss_coeff: Optional[float] = None - """Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.""" + dsa_indexer_loss_coeff: Optional[float] = None + """Coefficient for the DSA indexer KL divergence loss. Set to 0 to disable indexer loss.""" - use_sparse_indexer_loss: Optional[bool] = None - """Whether to use sparse indexer loss. If True, the indexer loss will be computed using the + dsa_indexer_use_sparse_loss: Optional[bool] = None + """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices.""" #################### @@ -1682,8 +1682,8 @@ def __post_init__(self): if self.sparse_attention_type is not None: assert ( self.context_parallel_size == 1 - ), "Currently context parallelism is not supported by SparseAttention!" - assert not self.apply_rope_fusion, "RoPE fusion is not supported for SparseAttention" + ), "Currently context parallelism is not supported by DSAttention!" + assert not self.apply_rope_fusion, "RoPE fusion is not supported for DSAttention" @dataclass diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 486152b3dbe..41f2ffae3cc 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3271,15 +3271,15 @@ def _add_sparse_attention_args(parser): group = parser.add_argument_group(title="sparse_attention") group.add_argument('--sparse-attention-type', default=None, choices=['dsa'], type=str, help="Type of sparse attention to use. Currently support dsa (DeepSeek Sparse Attention).") - group.add_argument('--index-n-heads', default=None, type=int, + 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('--index-head-dim', default=None, type=int, + 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('--index-topk', default=None, type=int, + 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('--indexer-loss-coeff', default=0.0, type=float, + group.add_argument('--dsa-indexer-loss-coeff', default=0.0, type=float, help='Coefficient for the indexer KL divergence loss. Set to 0 to disable indexer loss.') - group.add_argument('--use-sparse-indexer-loss', action='store_true', + 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 diff --git a/megatron/training/training.py b/megatron/training/training.py index 83f949408f8..6334a4396d9 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -90,7 +90,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.sparse_attention import IndexerLossLoggingHelper +from megatron.core.transformer.sparse_attention import DSAIndexerLossLoggingHelper from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper from megatron.core.parallel_state import ( destroy_global_memory_buffer, @@ -1678,9 +1678,9 @@ def training_log( mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict ) # Track sparse attention indexer loss - if args.indexer_loss_coeff is not None and args.indexer_loss_coeff > 0: + if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: indexer_loss_scale = 1 / get_num_microbatches() - IndexerLossLoggingHelper.track_indexer_metrics( + DSAIndexerLossLoggingHelper.track_indexer_metrics( loss_scale=indexer_loss_scale, iteration=iteration, writer=writer, diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index ec95a2af08b..96a90cef150 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -10,12 +10,12 @@ from megatron.core.transformer import TransformerConfig from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.sparse_attention import ( - Indexer, - IndexerLossAutoScaler, - IndexerSubmodules, - SparseAttention, - SparseAttentionSubmodules, - compute_indexer_loss, + DSAIndexer, + DSAIndexerLossAutoScaler, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, + compute_dsa_indexer_loss, rotate_activation, ) from megatron.core.transformer.transformer_config import MLATransformerConfig @@ -29,6 +29,7 @@ HAVE_HADAMARD = False +@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") class TestRotateActivation: """Test rotate_activation function.""" @@ -40,8 +41,6 @@ def setup_method(self): yield Utils.destroy_model_parallel() - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - @pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") def test_rotate_activation_shape(self): """Test that rotate_activation preserves shape.""" batch_size = 2 @@ -54,8 +53,6 @@ def test_rotate_activation_shape(self): assert output.shape == x.shape assert output.dtype == torch.bfloat16 - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - @pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") def test_rotate_activation_dtype_check(self): """Test that rotate_activation only accepts bfloat16.""" x = torch.randn(16, 2, 128, dtype=torch.float32).cuda() @@ -65,8 +62,8 @@ def test_rotate_activation_dtype_check(self): @pytest.mark.parametrize("seqlen_and_topk", [[16, 32], [64, 32]]) -class TestComputeIndexerLoss: - """Test compute_indexer_loss function.""" +class TestComputeDSAIndexerLoss: + """Test compute_dsa_indexer_loss function.""" @pytest.fixture(scope='function', autouse=True) def setup_method(self): @@ -78,7 +75,7 @@ def setup_method(self): Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_loss_shape(self, seqlen_and_topk): + 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] @@ -107,14 +104,14 @@ def test_indexer_loss_shape(self, seqlen_and_topk): key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() softmax_scale = head_dim**-0.5 - loss = compute_indexer_loss( + loss = compute_dsa_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, query=query, key=key, softmax_scale=softmax_scale, - indexer_loss_coeff=1.0, - use_sparse_indexer_loss=False, + loss_coeff=1.0, + sparse_loss=False, pg_collection=self.pg_collection, ) @@ -123,7 +120,7 @@ def test_indexer_loss_shape(self, seqlen_and_topk): assert loss >= 0 # KL divergence should be non-negative @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_loss_sparse(self, seqlen_and_topk): + def test_dsa_indexer_loss_sparse(self, seqlen_and_topk): """Test sparse indexer loss computation.""" batch_size = 2 seqlen = seqlen_and_topk[0] @@ -152,25 +149,25 @@ def test_indexer_loss_sparse(self, seqlen_and_topk): key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() softmax_scale = head_dim**-0.5 - loss_sparse = compute_indexer_loss( + loss_sparse = compute_dsa_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, query=query, key=key, softmax_scale=softmax_scale, - indexer_loss_coeff=1.0, - use_sparse_indexer_loss=True, + loss_coeff=1.0, + sparse_loss=True, pg_collection=self.pg_collection, ) - loss_dense = compute_indexer_loss( + loss_dense = compute_dsa_indexer_loss( index_scores=index_scores, topk_indices=topk_indices, query=query, key=key, softmax_scale=softmax_scale, - indexer_loss_coeff=1.0, - use_sparse_indexer_loss=False, + loss_coeff=1.0, + sparse_loss=False, pg_collection=self.pg_collection, ) @@ -183,8 +180,8 @@ def test_indexer_loss_sparse(self, seqlen_and_topk): assert loss_dense >= 0 -class TestIndexerLossAutoScaler: - """Test IndexerLossAutoScaler autograd function.""" +class TestDSAIndexerLossAutoScaler: + """Test DSAIndexerLossAutoScaler autograd function.""" @pytest.fixture(scope='function', autouse=True) def setup_method(self): @@ -202,7 +199,7 @@ def test_forward_pass(self): indexer_loss = torch.tensor(0.5).cuda() indexer_loss.requires_grad_(True) - result = IndexerLossAutoScaler.apply(output, indexer_loss) + result = DSAIndexerLossAutoScaler.apply(output, indexer_loss) assert torch.allclose(result, output, atol=0, rtol=0) @@ -213,17 +210,17 @@ def test_backward_pass(self): output.requires_grad_(True) # Create indexer_loss with computation graph - # This simulates compute_indexer_loss which computes KL divergence + # 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() - IndexerLossAutoScaler.set_loss_scale(scale) + DSAIndexerLossAutoScaler.set_loss_scale(scale) # Apply the autograd function - result = IndexerLossAutoScaler.apply(output, indexer_loss) + result = DSAIndexerLossAutoScaler.apply(output, indexer_loss) # Trigger backward main_loss = result.sum() @@ -245,9 +242,10 @@ def test_backward_pass(self): ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}" +@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @pytest.mark.parametrize("seqlen", [16, 64]) -class TestIndexer: - """Test Indexer module basic functionality with TP=1.""" +class TestDSAIndexer: + """Test DSA Indexer module basic functionality with TP=1.""" @pytest.fixture(scope='function', autouse=True) def setup_method(self): @@ -272,20 +270,20 @@ def setup_method(self): qk_head_dim=64, qk_pos_emb_head_dim=32, v_head_dim=64, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=self.index_topk, 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 = IndexerSubmodules( + indexer_submodules = DSAIndexerSubmodules( linear_wq_b=ModuleSpec(module=TELinear), linear_wk=ModuleSpec(module=TELinear), k_norm=ModuleSpec(module=TENorm), @@ -295,21 +293,21 @@ def setup_method(self): self.pg_collection = ProcessGroupCollection.use_mpu_process_groups( required_pgs=['tp', 'cp'] ) - self.indexer = Indexer(self.config, indexer_submodules, self.pg_collection) + self.indexer = DSAIndexer(self.config, indexer_submodules, self.pg_collection) yield Utils.destroy_model_parallel() - def test_indexer_constructor(self, seqlen): + def test_dsa_indexer_constructor(self, seqlen): """Test indexer initialization.""" - assert isinstance(self.indexer, Indexer) + 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_indexer_forward(self, seqlen): + def test_dsa_indexer_forward(self, seqlen): """Test indexer forward pass.""" batch_size = 2 @@ -323,7 +321,7 @@ def test_indexer_forward(self, seqlen): topk_indices = self.indexer(x, qr) # Check output shape - assert topk_indices.shape == (batch_size, seqlen, min(self.config.index_topk, seqlen)) + 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 @@ -333,7 +331,7 @@ def test_indexer_forward(self, seqlen): ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_forward_with_scores(self, seqlen): + def test_dsa_indexer_forward_with_scores(self, seqlen): """Test indexer forward pass with scores.""" batch_size = 2 @@ -348,7 +346,7 @@ def test_indexer_forward_with_scores(self, seqlen): # Check output shapes assert index_scores.shape == (batch_size, seqlen, seqlen) - assert topk_indices.shape == (batch_size, seqlen, min(self.config.index_topk, 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)) @@ -359,7 +357,7 @@ def test_indexer_forward_with_scores(self, seqlen): ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_with_mask(self, seqlen): + def test_dsa_indexer_with_mask(self, seqlen): """Test indexer with attention mask.""" batch_size = 2 @@ -384,8 +382,9 @@ def test_indexer_with_mask(self, seqlen): assert torch.all(topk_indices[b, i] <= max(self.index_topk, i)) -class TestSparseAttention: - """Test SparseAttention module basic functionality with TP=1.""" +@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") +class TestDSAttention: + """Test DSAttention module basic functionality with TP=1.""" @pytest.fixture(scope='function', autouse=True) def setup_method(self): @@ -409,35 +408,35 @@ def setup_method(self): qk_head_dim=64, qk_pos_emb_head_dim=32, v_head_dim=64, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=32, rope_type='rope', rotary_base=10000, rotary_percent=1.0, - indexer_loss_coeff=0.1, - use_sparse_indexer_loss=False, + # 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 = IndexerSubmodules( + 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=Indexer, submodules=indexer_submodules) - sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) + 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 = SparseAttention( + self.sparse_attention = DSAttention( config=self.config, submodules=sparse_attention_submodules, layer_number=1, @@ -449,14 +448,14 @@ def setup_method(self): yield Utils.destroy_model_parallel() - def test_sparse_attention_constructor(self): + def test_dsa_constructor(self): """Test sparse attention initialization.""" - assert isinstance(self.sparse_attention, SparseAttention) + assert isinstance(self.sparse_attention, DSAttention) assert hasattr(self.sparse_attention, 'indexer') - assert isinstance(self.sparse_attention.indexer, Indexer) + assert isinstance(self.sparse_attention.indexer, DSAIndexer) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_sparse_attention_forward(self): + def test_dsa_forward(self): """Test sparse attention forward pass.""" seq_len = 16 batch_size = 2 @@ -506,7 +505,7 @@ def test_sparse_attention_forward(self): assert output.dtype == torch.bfloat16 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_sparse_attention_backward(self): + def test_dsa_backward(self): """Test sparse attention backward pass with indexer loss.""" seq_len = 16 batch_size = 2 @@ -567,7 +566,7 @@ def test_sparse_attention_backward(self): 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_sparse_attention_topk_selection(self): + def test_dsa_topk_selection(self): """Test that sparse attention correctly selects top-k indices.""" seq_len = 16 batch_size = 2 @@ -608,7 +607,7 @@ def test_sparse_attention_topk_selection(self): # 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.index_topk, seq_len) + assert topk_indices.shape[2] == min(self.config.dsa_indexer_topk, seq_len) # ====================================================================================== @@ -616,10 +615,11 @@ def test_sparse_attention_topk_selection(self): # ====================================================================================== +@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @pytest.mark.parametrize("tensor_model_parallel_size", [2, 4, 8]) @pytest.mark.parametrize("sequence_parallel", [False, True]) class TestIndexerTensorParallel: - """Test Indexer with different TP sizes and SP settings, compare with TP=1 baseline.""" + """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.""" @@ -641,13 +641,13 @@ def _create_config(self, sequence_parallel=False): qk_head_dim=64, qk_pos_emb_head_dim=32, v_head_dim=64, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=32, 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): @@ -655,17 +655,17 @@ def _create_indexer(self, config, pg_collection): from megatron.core.extensions.transformer_engine import TELinear, TENorm from megatron.core.transformer.spec_utils import ModuleSpec - indexer_submodules = IndexerSubmodules( + 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 Indexer(config, indexer_submodules, pg_collection) + return DSAIndexer(config, indexer_submodules, pg_collection) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_weight_consistency(self, tensor_model_parallel_size, sequence_parallel): + 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 @@ -696,7 +696,7 @@ def test_indexer_weight_consistency(self, tensor_model_parallel_size, sequence_p Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_forward_consistency(self, tensor_model_parallel_size, sequence_parallel): + 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( @@ -791,7 +791,7 @@ def test_indexer_forward_consistency(self, tensor_model_parallel_size, sequence_ Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_indexer_gradient_sync(self, tensor_model_parallel_size, sequence_parallel): + 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 @@ -853,11 +853,12 @@ def test_indexer_gradient_sync(self, tensor_model_parallel_size, sequence_parall Utils.destroy_model_parallel() +@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @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 TestSparseAttentionTensorParallel: - """Test SparseAttention with different TP sizes, SP settings, and sparse indexer loss.""" +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.""" @@ -879,15 +880,15 @@ def _create_config(self, sequence_parallel=False, use_sparse_indexer_loss=False) qk_head_dim=64, qk_pos_emb_head_dim=32, v_head_dim=64, - # Sparse attention specific configs - index_n_heads=8, - index_head_dim=64, - index_topk=32, rope_type='rope', rotary_base=10000, rotary_percent=1.0, - indexer_loss_coeff=0.1, - use_sparse_indexer_loss=use_sparse_indexer_loss, + # 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): @@ -895,16 +896,16 @@ def _create_sparse_attention(self, config, pg_collection): from megatron.core.extensions.transformer_engine import TELinear, TENorm from megatron.core.transformer.spec_utils import ModuleSpec - indexer_submodules = IndexerSubmodules( + 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=Indexer, submodules=indexer_submodules) - sparse_attention_submodules = SparseAttentionSubmodules(indexer=indexer_spec) + indexer_spec = ModuleSpec(module=DSAIndexer, submodules=indexer_submodules) + sparse_attention_submodules = DSAttentionSubmodules(indexer=indexer_spec) - return SparseAttention( + return DSAttention( config=config, submodules=sparse_attention_submodules, layer_number=1, @@ -914,7 +915,7 @@ def _create_sparse_attention(self, config, pg_collection): ) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_sparse_attention_weight_consistency( + 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.""" @@ -947,7 +948,7 @@ def test_sparse_attention_weight_consistency( Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_sparse_attention_forward_consistency( + 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.""" @@ -1146,7 +1147,7 @@ def test_sparse_attention_forward_consistency( Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_sparse_attention_gradient_sync( + 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.""" From 0d7e1d11cd38367baa9a00703509166e6c6b8634 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 27 Nov 2025 18:21:33 +0800 Subject: [PATCH 20/28] Fix lint error Signed-off-by: kunlunl --- megatron/core/models/gpt/gpt_layer_specs.py | 4 +++- megatron/core/transformer/transformer_config.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index ecaec8c0360..7fbe836b20c 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -320,7 +320,9 @@ def get_attention_module_spec_for_backend( normalization=normalization, ) else: - core_attention = backend.core_attention() if not fallback_to_eager_attn else DotProductAttention + core_attention = ( + backend.core_attention() if not fallback_to_eager_attn else DotProductAttention + ) if multi_latent_attention: assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index fbd589b1811..2445f266be5 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1895,7 +1895,7 @@ def __post_init__(self): f"fallback_to_eager_attn only supports all_gather communication type " f"for context parallelism, but got {self.cp_comm_type=} instead." ) - + if self.sparse_attention_type is not None: assert ( self.context_parallel_size == 1 From 6630c7f3b345de83827fd9e15a98b7d095a2518a Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 28 Nov 2025 15:32:43 +0800 Subject: [PATCH 21/28] Merge linear-attention-type and sparse-attention-type to experimental-attention-variant Signed-off-by: kunlunl --- gpt_builders.py | 9 +- ...rimental_attention_variant_module_specs.py | 129 ++++++++++++++++++ megatron/core/models/gpt/gpt_layer_specs.py | 87 +++++------- .../gpt/linear_attention_module_specs.py | 27 ---- .../gpt/sparse_attention_module_specs.py | 61 --------- .../dsa.py} | 0 .../transformer/multi_latent_attention.py | 23 +++- .../core/transformer/transformer_config.py | 66 ++++----- megatron/training/arguments.py | 50 ++++--- megatron/training/training.py | 7 +- tests/unit_tests/ssm/test_gated_delta_net.py | 4 +- .../transformer/test_sparse_attention.py | 2 +- 12 files changed, 250 insertions(+), 215 deletions(-) create mode 100644 megatron/core/models/gpt/experimental_attention_variant_module_specs.py delete mode 100644 megatron/core/models/gpt/linear_attention_module_specs.py delete mode 100644 megatron/core/models/gpt/sparse_attention_module_specs.py rename megatron/core/transformer/{sparse_attention.py => experimental_attention_variant/dsa.py} (100%) diff --git a/gpt_builders.py b/gpt_builders.py index d81a36f021c..61d159b9967 100644 --- a/gpt_builders.py +++ b/gpt_builders.py @@ -42,7 +42,8 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None): else: use_te = args.transformer_impl == "transformer_engine" - if args.num_experts or (args.linear_attention_type is not None): + linear_attention_variants = ["gated_delta_net"] + if args.num_experts or args.experimental_attention_variant in linear_attention_variants: # Define the decoder block spec transformer_layer_spec = get_gpt_decoder_block_spec( config, @@ -114,8 +115,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, - args.sparse_attention_type, - args.linear_attention_type, + 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, @@ -127,8 +127,7 @@ def _get_transformer_layer_spec(use_te, config): args.moe_grouped_gemm, args.qk_layernorm, args.multi_latent_attention, - args.sparse_attention_type, - args.linear_attention_type, + 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..6ef274989bf --- /dev/null +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +from typing import Optional + +from megatron.core.models.backends import BackendSpecProvider +from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexer, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, +) +from megatron.core.transformer.multi_latent_attention import ( + MLASelfAttention, + MLASelfAttentionSubmodules, +) +from megatron.core.transformer.spec_utils import ModuleSpec + + +def get_gated_delta_net_module_spec_for_backend( + backend: BackendSpecProvider, normalization: Optional[str] = None +) -> ModuleSpec: + """Helper function to get module spec for Linear Attention""" + rms_norm = normalization == "RMSNorm" + attention = ModuleSpec( + module=GatedDeltaNet, + submodules=GatedDeltaNetSubmodules( + in_proj=backend.column_parallel_layer_norm_linear(), + out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), + out_proj=backend.row_parallel_linear(), + ), + metainfo={"fuse_input_layernorm": True}, + ) + return attention + + +def get_dsa_module_spec_for_backend( + backend: BackendSpecProvider, + qk_l2_norm: Optional[bool] = False, + multi_latent_attention: Optional[bool] = False, + mla_down_proj_use_column_parallel: Optional[bool] = False, + normalization: Optional[str] = None, + fallback_to_eager_attn: Optional[bool] = False, +) -> 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." + assert fallback_to_eager_attn is False, "Fallback to eager attention is not supported with DSA." + + # Adjust for RMS norm. + rms_norm = normalization == "RMSNorm" + qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) + + linear_q_down_proj = ( + backend.column_parallel_linear() if mla_down_proj_use_column_parallel else backend.linear() + ) + linear_kv_down_proj = ( + backend.column_parallel_linear() if mla_down_proj_use_column_parallel else backend.linear() + ) + linear_q_up_proj = backend.column_parallel_linear() + linear_kv_up_proj = backend.column_parallel_linear() + qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) + + # 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=linear_q_down_proj, + linear_q_up_proj=linear_q_up_proj, + linear_kv_down_proj=linear_kv_down_proj, + linear_kv_up_proj=linear_kv_up_proj, + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) + + return attention + + +def get_experimental_attention_variant_module_spec_for_backend( + backend: BackendSpecProvider, + sharded_state_dict_keys_map: dict, + experimental_attention_variant: Optional[str] = None, + qk_layernorm: Optional[bool] = False, + qk_l2_norm: Optional[bool] = False, + multi_latent_attention: Optional[bool] = False, + mla_down_proj_use_column_parallel: Optional[bool] = False, + normalization: Optional[str] = None, + fallback_to_eager_attn: Optional[bool] = False, +) -> ModuleSpec: + """Helper function to get module spec for Attention""" + if experimental_attention_variant == "gated_delta_net": + return get_gated_delta_net_module_spec_for_backend( + backend=backend, normalization=normalization + ) + elif experimental_attention_variant == "dsa": + return get_dsa_module_spec_for_backend( + backend=backend, + qk_l2_norm=qk_l2_norm, + multi_latent_attention=multi_latent_attention, + mla_down_proj_use_column_parallel=mla_down_proj_use_column_parallel, + normalization=normalization, + fallback_to_eager_attn=fallback_to_eager_attn, + ) + 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 7fbe836b20c..35283582eee 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -5,13 +5,10 @@ from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add from megatron.core.models.backends import BackendSpecProvider, LocalSpecProvider -from megatron.core.models.gpt.linear_attention_module_specs import ( - get_linear_attention_module_spec_for_backend, +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.models.gpt.sparse_attention_module_specs import ( - get_sparse_attention_module_spec_for_backend, -) from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules from megatron.core.transformer.dot_product_attention import DotProductAttention from megatron.core.transformer.enums import AttnMaskType, LayerType @@ -81,8 +78,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, - sparse_attention_type: Optional[str] = None, - linear_attention_type: Optional[str] = None, + 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, @@ -100,8 +96,8 @@ def get_gpt_layer_with_transformer_engine_spec( 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 multi-latent attention. Defaults to False. - sparse_attention_type (str, optional): The type of sparse attention. Defaults to None. - linear_attention_type (str, optional): The type of linear attention. Defaults to None. + 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. @@ -138,8 +134,7 @@ def get_gpt_layer_with_transformer_engine_spec( attention = get_attention_module_spec_for_backend( backend=backend, sharded_state_dict_keys_map=sharded_state_dict_keys_map, - sparse_attention_type=sparse_attention_type, - linear_attention_type=linear_attention_type, + experimental_attention_variant=experimental_attention_variant, qk_layernorm=qk_layernorm, qk_l2_norm=qk_l2_norm, multi_latent_attention=multi_latent_attention, @@ -172,8 +167,7 @@ def get_gpt_layer_local_spec( moe_grouped_gemm: Optional[bool] = False, qk_layernorm: Optional[bool] = False, multi_latent_attention: Optional[bool] = False, - sparse_attention_type: Optional[str] = None, - linear_attention_type: Optional[str] = None, + 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, @@ -188,8 +182,8 @@ def get_gpt_layer_local_spec( 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 multi-latent attention. Defaults to False. - sparse_attention_type (str, optional): The type of sparse attention. Defaults to None. - linear_attention_type (str, optional): The type of linear attention. Defaults to None. + 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. @@ -213,19 +207,17 @@ def get_gpt_layer_local_spec( " and will be removed soon. Please update your code accordingly." ) - if sparse_attention_type is not None: - raise NotImplementedError("Sparse attention is not supported with local spec yet.") - - if linear_attention_type is not None: - raise NotImplementedError("Linear attention is not supported with local spec yet.") + if experimental_attention_variant is not None: + raise NotImplementedError( + "Experimental attention variant is not supported with local spec yet." + ) sharded_state_dict_keys_map = {} attention = get_attention_module_spec_for_backend( backend=backend, sharded_state_dict_keys_map=sharded_state_dict_keys_map, - sparse_attention_type=sparse_attention_type, - linear_attention_type=linear_attention_type, + experimental_attention_variant=experimental_attention_variant, qk_layernorm=qk_layernorm, qk_l2_norm=qk_l2_norm, multi_latent_attention=multi_latent_attention, @@ -290,8 +282,7 @@ def get_transformer_layer_spec_for_backend( def get_attention_module_spec_for_backend( backend: BackendSpecProvider, sharded_state_dict_keys_map: dict, - sparse_attention_type: Optional[str] = None, - linear_attention_type: Optional[str] = None, + experimental_attention_variant: Optional[str] = None, qk_layernorm: Optional[bool] = False, qk_l2_norm: Optional[bool] = False, multi_latent_attention: Optional[bool] = False, @@ -301,28 +292,24 @@ def get_attention_module_spec_for_backend( ) -> ModuleSpec: """Helper function to get module spec for Attention""" - if linear_attention_type is not None: - return get_linear_attention_module_spec_for_backend( - backend=backend, - linear_attention_type=linear_attention_type, - normalization=normalization, + if experimental_attention_variant is not None: + return get_experimental_attention_variant_module_spec_for_backend( + backend, + sharded_state_dict_keys_map, + experimental_attention_variant, + qk_layernorm, + qk_l2_norm, + multi_latent_attention, + mla_down_proj_use_column_parallel, + normalization, + fallback_to_eager_attn, ) # Adjust for RMS norm. rms_norm = normalization == "RMSNorm" qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) - if sparse_attention_type is not None: - assert multi_latent_attention, "Currently only MLA supports sparse attention." - core_attention = get_sparse_attention_module_spec_for_backend( - backend=backend, - sparse_attention_type=sparse_attention_type, - normalization=normalization, - ) - else: - core_attention = ( - backend.core_attention() if not fallback_to_eager_attn else DotProductAttention - ) + core_attention = backend.core_attention() if not fallback_to_eager_attn else DotProductAttention if multi_latent_attention: assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." @@ -336,20 +323,19 @@ def get_attention_module_spec_for_backend( if mla_down_proj_use_column_parallel else backend.linear() ) - fuse_norm_and_linear = backend.fuse_layernorm_and_linear() and sparse_attention_type is None linear_q_up_proj = ( backend.column_parallel_layer_norm_linear() - if qk_layernorm and fuse_norm_and_linear + if qk_layernorm and backend.fuse_layernorm_and_linear() else backend.column_parallel_linear() ) linear_kv_up_proj = ( backend.column_parallel_layer_norm_linear() - if qk_layernorm and fuse_norm_and_linear + if qk_layernorm and backend.fuse_layernorm_and_linear() else backend.column_parallel_linear() ) qk_norm = ( backend.layer_norm(rms_norm=rms_norm, for_qk=True) - if qk_layernorm and not fuse_norm_and_linear + if qk_layernorm and not backend.fuse_layernorm_and_linear() else IdentityOp ) attention = ModuleSpec( @@ -551,24 +537,20 @@ def get_gpt_decoder_layer_specs( num_experts = None moe_grouped_gemm = None if attention_type == "linear_attention": - if config.linear_attention_type is None: + linear_attention_variants = ["gated_delta_net"] + if config.experimental_attention_variant not in linear_attention_variants: # Skip if there is no linear attention layer in the model. continue - linear_attention_type = config.linear_attention_type multi_latent_attention = None - sparse_attention_type = None else: - linear_attention_type = None multi_latent_attention = config.multi_latent_attention - sparse_attention_type = config.sparse_attention_type layer_spec_key = f"{mlp_type}_{attention_type}" layer_spec_dict[layer_spec_key] = get_layer_spec_fn( num_experts=num_experts, moe_grouped_gemm=moe_grouped_gemm, multi_latent_attention=multi_latent_attention, - sparse_attention_type=sparse_attention_type, - linear_attention_type=linear_attention_type, + experimental_attention_variant=config.experimental_attention_variant, **get_layer_spec_kwargs, ) @@ -611,7 +593,8 @@ def get_gpt_decoder_layer_specs( f"current linear attention pattern: {config.linear_attention_freq}" ) elif config.linear_attention_freq is None: - if config.linear_attention_type is None: + linear_attention_variants = ["gated_delta_net"] + if config.experimental_attention_variant not in linear_attention_variants: linear_attention_pattern = [0] * config.num_layers else: linear_attention_pattern = [1] * config.num_layers diff --git a/megatron/core/models/gpt/linear_attention_module_specs.py b/megatron/core/models/gpt/linear_attention_module_specs.py deleted file mode 100644 index 7e76d845cff..00000000000 --- a/megatron/core/models/gpt/linear_attention_module_specs.py +++ /dev/null @@ -1,27 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -from typing import Optional - -from megatron.core.models.backends import BackendSpecProvider -from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules -from megatron.core.transformer.spec_utils import ModuleSpec - - -def get_linear_attention_module_spec_for_backend( - backend: BackendSpecProvider, linear_attention_type: str, normalization: Optional[str] = None -) -> ModuleSpec: - """Helper function to get module spec for Linear Attention""" - rms_norm = normalization == "RMSNorm" - if linear_attention_type == "gated_delta_net": - attention = ModuleSpec( - module=GatedDeltaNet, - submodules=GatedDeltaNetSubmodules( - in_proj=backend.column_parallel_layer_norm_linear(), - out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False), - out_proj=backend.row_parallel_linear(), - ), - metainfo={"fuse_input_layernorm": True}, - ) - else: - raise ValueError(f"Invalid linear attention type: {linear_attention_type}") - return attention diff --git a/megatron/core/models/gpt/sparse_attention_module_specs.py b/megatron/core/models/gpt/sparse_attention_module_specs.py deleted file mode 100644 index 596690011d1..00000000000 --- a/megatron/core/models/gpt/sparse_attention_module_specs.py +++ /dev/null @@ -1,61 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. - -from typing import Optional - -from megatron.core.models.backends import BackendSpecProvider -from megatron.core.transformer.sparse_attention import ( - DSAIndexer, - DSAIndexerSubmodules, - DSAttention, - DSAttentionSubmodules, -) -from megatron.core.transformer.spec_utils import ModuleSpec - - -def get_dsa_indexer_spec_for_backend( - backend: BackendSpecProvider, normalization: Optional[str] = None -) -> ModuleSpec: - """Helper function to get DSA Indexer module spec for a given backend. - - Args: - backend: Backend specification provider (TE or Local). - normalization: Normalization type ("RMSNorm" or None for LayerNorm). - - Returns: - ModuleSpec for DSA Indexer with appropriate submodules. - """ - return 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(), - ), - ) - - -def get_sparse_attention_module_spec_for_backend( - backend: BackendSpecProvider, sparse_attention_type: str, normalization: Optional[str] = None -) -> ModuleSpec: - """Helper function to get module spec for Sparse Attention. - - Args: - backend: Backend specification provider (TE or Local). - sparse_attention_type: Type of sparse attention. - normalization: Normalization type ("RMSNorm" or None for LayerNorm). - - Returns: - ModuleSpec for the sparse attention implementation with appropriate submodules. - """ - if sparse_attention_type == "dsa": - # Because TransformerEngine does not support sparse attention yet, we use local - # implementation whether the backend is TransformerEngine or not. - return ModuleSpec( - module=DSAttention, - submodules=DSAttentionSubmodules( - indexer=get_dsa_indexer_spec_for_backend(backend, normalization=normalization) - ), - ) - else: - raise ValueError(f"Invalid sparse attention type: {sparse_attention_type}") diff --git a/megatron/core/transformer/sparse_attention.py b/megatron/core/transformer/experimental_attention_variant/dsa.py similarity index 100% rename from megatron/core/transformer/sparse_attention.py rename to megatron/core/transformer/experimental_attention_variant/dsa.py diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index a0e906e3a5f..d24be883d11 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -243,7 +243,7 @@ def forward( # Get the query, key and value tensors based on the type of attention - # self or cross attn. # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128] - if self.config.sparse_attention_type is None: + if self.config.experimental_attention_variant is None: query, key, value = self.get_query_key_value_tensors( hidden_states, key_value_states, @@ -251,8 +251,7 @@ def forward( packed_seq_params, inference_context=inference_context, ) - else: - # TODO(kunlunl): Is this a universal usage of sparse attention? + elif self.config.experimental_attention_variant == "dsa": query, key, value, q_compressed, _ = self.get_query_key_value_and_compressed_tensors( hidden_states, key_value_states, @@ -260,6 +259,11 @@ def forward( packed_seq_params, inference_context=inference_context, ) + else: + raise ValueError( + f"Unsupported experimental attention variant: " + f"{self.config.experimental_attention_variant}" + ) # =================================================== # Adjust key, value for inference @@ -291,7 +295,7 @@ def forward( if inference_context is None or inference_context.is_static_batching(): with get_fine_grained_offloading_context(self.offload_core_attention): - if self.config.sparse_attention_type is None: + if self.config.experimental_attention_variant is None: core_attn_out = self.core_attention( query, key, @@ -300,9 +304,9 @@ def forward( packed_seq_params=packed_seq_params, attn_mask_type=attn_mask_type, ) - else: - # For sparse attention, use a specialized forward. - # TODO(kunlunl): Is there a unified interface for sparse attention? + elif self.config.experimental_attention_variant == "dsa": + # For dsa we need to pass in the original hidden states and the compressed + # query representation. core_attn_out = self.core_attention( query, key, @@ -314,6 +318,11 @@ def forward( attention_bias=None, packed_seq_params=packed_seq_params, ) + else: + raise ValueError( + f"Unsupported attention variant: " + f"{self.config.experimental_attention_variant}" + ) elif self.cache_mla_latents: # Dynamic batching attention kernel. q, k, v = (query, key, value) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2445f266be5..656699ea2a2 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -233,33 +233,14 @@ class TransformerConfig(ModelParallelConfig): 16 SMs can generally achieve good bandwidth.""" #################### - # sparse attention + # attention variant #################### - sparse_attention_type: Optional[str] = None - """Type of sparse attention to use. Currently only supports dsa (DeepSeek Sparse Attention).""" - - dsa_indexer_n_heads: Optional[int] = None - """Number of DSA indexer heads for DSA.""" - - 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: Optional[bool] = None - """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the - top-k indices.""" + experimental_attention_variant: Optional[str] = None + """Type of attention variant to use. Currently support gated_delta_net and dsa.""" #################### - # linear attention + # attention variant: gated_delta_net #################### - linear_attention_type: Optional[str] = None - """Type of linear attention to use. Currently support gated_delta_net.""" - linear_attention_freq: Optional[Union[int, List[int]]] = None """Frequency between LA (linear attention) layers and SDPA (scaled dot-product attention) layers. @@ -282,6 +263,25 @@ class TransformerConfig(ModelParallelConfig): linear_num_value_heads: Optional[int] = None """Number of value and gate heads for the gated delta net.""" + #################### + # attention variant: 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: Optional[bool] = None + """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the + top-k indices.""" + #################### # initialization #################### @@ -877,17 +877,12 @@ def __post_init__(self): f"tensor_model_parallel_size ({self.tensor_model_parallel_size})." ) - if self.linear_attention_type is not None: - supported_la_types = ["gated_delta_net"] - assert self.linear_attention_type in supported_la_types, ( - f"linear_attention_type ({self.linear_attention_type}) only support" - f" one of {supported_la_types}." - ) + if self.experimental_attention_variant in ["gated_delta_net"]: assert ( self.linear_attention_freq is not None ), f"linear_attention_freq must be set for linear attention." - if self.linear_attention_type == "gated_delta_net": + if self.experimental_attention_variant == "gated_delta_net": # Check required parameters assert ( self.linear_conv_kernel_dim is not None @@ -922,6 +917,11 @@ def __post_init__(self): f"Gated delta net does not support context parallel for now," f" but got {self.context_parallel_size=}." ) + elif 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.fp8: # cannot support first last layer bf16 with delayed scaling @@ -1896,12 +1896,6 @@ def __post_init__(self): f"for context parallelism, but got {self.cp_comm_type=} instead." ) - if self.sparse_attention_type is not None: - 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" - @dataclass class MLATransformerConfig(TransformerConfig): diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 69c9ba25d40..c7ee0426b78 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -69,8 +69,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_sparse_attention_args(parser) - parser = _add_linear_attention_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) @@ -1201,7 +1200,7 @@ def validate_args(args, defaults={}): assert not args.use_torch_fsdp2, "Muon optimizer does not support Torch-FSDP2 for now." assert not args.use_megatron_fsdp, "Muon optimizer does not support Megatron-FSDP for now." assert args.ckpt_format in ["torch", "torch_dist"], "Muon optimizer supports torch and torch_dist checkpoint format." - assert args.linear_attention_type is None, "Muon optimizer does not support linear attention type for now." + assert args.experimental_attention_variant is None, "Muon optimizer does not support attention variant for now." assert not args.attention_output_gate, "Muon optimizer does not support attention output gate for now." # Optimizer CPU offload check @@ -1270,6 +1269,14 @@ def validate_args(args, defaults={}): if args.multi_latent_attention: assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." + if args.linear_attention_type is not None: + print_rank_0( + '--linear-attention-type is deprecated, use --experimental-attention-variant instead.', + args.rank, + ) + args.experimental_attention_variant = args.linear_attention_type + del args.linear_attention_type + # Print arguments. _print_args("arguments", args) @@ -3352,26 +3359,14 @@ def _add_mla_args(parser): return parser -def _add_sparse_attention_args(parser): - group = parser.add_argument_group(title="sparse_attention") - group.add_argument('--sparse-attention-type', default=None, choices=['dsa'], type=str, - help="Type of sparse attention to use. Currently support dsa (DeepSeek Sparse Attention).") - 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=0.0, 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_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.') -def _add_linear_attention_args(parser): - group = parser.add_argument_group(title="la") + # Linear attention group.add_argument('--linear-attention-type', default=None, choices=['gated_delta_net'], type=str, - help='Type of linear attention to use. Currently support gated_delta_net.') + help='(Deprecated, use --experimental-attention-variant instead) Type of linear attention to use. Currently support gated_delta_net.') group.add_argument('--linear-attention-freq', type=la_freq_type, default=None, help='Frequency between LA (linear attention) layers and' ' SDPA (scaled dot-product attention) layers. Accepts either: ' @@ -3391,6 +3386,19 @@ def _add_linear_attention_args(parser): help='Number of query and key heads for the gated delta net.') group.add_argument('--linear-num-value-heads', default=32, type=int, help='Number of value and gate heads for the gated delta net.') + + # 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=0.0, 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): diff --git a/megatron/training/training.py b/megatron/training/training.py index a02715c636a..a8e6da33184 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -91,7 +91,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.sparse_attention import DSAIndexerLossLoggingHelper +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, @@ -376,7 +376,8 @@ def transformer_flops(): ) ) - if args.linear_attention_type is not None: + linear_attention_variants = ["gated_delta_net"] + if args.experimental_attention_variant in linear_attention_variants: # Calculate number of dense and MoE Transformer MLPs. if isinstance(args.linear_attention_freq, int): linear_attention_pattern = [ @@ -401,7 +402,7 @@ def transformer_flops(): num_linear_attention_layers = sum(linear_attention_pattern) num_standard_attention_layers = num_layers - num_linear_attention_layers - if args.linear_attention_type == "gated_delta_net": + if args.experimental_attention_variant == "gated_delta_net": # Calculate the FLOPs for the gated delta net attention. qk_head_dim = args.linear_key_head_dim v_head_dim = args.linear_value_head_dim diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index dbf8d203634..89a185e3755 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -88,7 +88,7 @@ def setup_method(self, tp_size, sp, cp_size): context_parallel_size=cp_size, ) gdn_submodules = get_gpt_layer_with_transformer_engine_spec( - linear_attention_type="gated_delta_net", normalization="RMSNorm" + experimental_attention_variant="gated_delta_net", normalization="RMSNorm" ).submodules.self_attention.submodules self.gdn = GatedDeltaNet( @@ -157,7 +157,7 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, tp, sp, cp): # Model initialization function def initialize_gpt_model(config, pre_process=True, post_process=True, vp_stage=None): layer_spec = get_gpt_layer_with_transformer_engine_spec( - linear_attention_type="gated_delta_net", normalization=normalization + experimental_attention_variant="gated_delta_net", normalization=normalization ) gpt_model = GPTModel( config=config, diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_sparse_attention.py index 96a90cef150..6f45862c7c6 100644 --- a/tests/unit_tests/transformer/test_sparse_attention.py +++ b/tests/unit_tests/transformer/test_sparse_attention.py @@ -9,7 +9,7 @@ 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.sparse_attention import ( +from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerLossAutoScaler, DSAIndexerSubmodules, From b1e309e99ff980fb86df3c57a9a864944a8db7db Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 28 Nov 2025 16:08:45 +0800 Subject: [PATCH 22/28] Rename test_sparse_attention to test_attention_variant_dsa Signed-off-by: kunlunl --- .../models/gpt/experimental_attention_variant_module_specs.py | 1 + megatron/core/models/gpt/gpt_layer_specs.py | 1 - .../{test_sparse_attention.py => test_attention_variant_dsa.py} | 0 3 files changed, 1 insertion(+), 1 deletion(-) rename tests/unit_tests/transformer/{test_sparse_attention.py => test_attention_variant_dsa.py} (100%) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 6ef274989bf..cd8cd1fefd1 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -4,6 +4,7 @@ from megatron.core.models.backends import BackendSpecProvider from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules +from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py index a8d86267f0f..a650bd5fe75 100755 --- a/megatron/core/models/gpt/gpt_layer_specs.py +++ b/megatron/core/models/gpt/gpt_layer_specs.py @@ -410,7 +410,6 @@ def get_attention_module_spec_for_backend( qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) core_attention = backend.core_attention() if not fallback_to_eager_attn else DotProductAttention - if multi_latent_attention: assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." linear_q_down_proj = ( diff --git a/tests/unit_tests/transformer/test_sparse_attention.py b/tests/unit_tests/transformer/test_attention_variant_dsa.py similarity index 100% rename from tests/unit_tests/transformer/test_sparse_attention.py rename to tests/unit_tests/transformer/test_attention_variant_dsa.py From 6fe2f352d24e38e7b395d5ca0ffda4112db1029f Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 28 Nov 2025 16:47:08 +0800 Subject: [PATCH 23/28] Minor fixes Signed-off-by: kunlunl --- .../experimental_attention_variant/dsa.py | 1 - .../transformer/multi_latent_attention.py | 21 +++++++------------ pyproject.toml | 1 + 3 files changed, 9 insertions(+), 14 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 95e012121a8..fc994490b1b 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -440,7 +440,6 @@ def __init__( eps=self.config.layernorm_epsilon, ) - # TODO(kunlunl): The dtype of this module should be torch.get_default_dtype(). self.linear_weights_proj = build_module( submodules.linear_weights_proj, self.hidden_size, diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index d24be883d11..3953d933b45 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -252,12 +252,13 @@ def forward( inference_context=inference_context, ) elif self.config.experimental_attention_variant == "dsa": - query, key, value, q_compressed, _ = self.get_query_key_value_and_compressed_tensors( + query, key, value, q_compressed, _ = self.get_query_key_value_tensors( hidden_states, key_value_states, position_ids, packed_seq_params, inference_context=inference_context, + return_compressed_tensors=True, ) else: raise ValueError( @@ -519,7 +520,7 @@ def __init__( eps=self.config.layernorm_epsilon, ) - def get_query_key_value_and_compressed_tensors( + def get_query_key_value_tensors( self, hidden_states, key_value_states=None, @@ -528,6 +529,7 @@ def get_query_key_value_and_compressed_tensors( inference_context=None, *, inference_params=None, + return_compressed_tensors=False, ): """ Derives `query`, `key` and `value` tensors from `hidden_states`. @@ -862,17 +864,10 @@ 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, q_compressed, kv_compressed - - def get_query_key_value_tensors(self, *args, **kwargs): - """ - Derives `query`, `key` and `value` tensors from `hidden_states`. - """ - query, key, value, q_compressed, kv_compressed = ( - self.get_query_key_value_and_compressed_tensors(*args, **kwargs) - ) - # Only return query, key and value. - return query, key, value + if return_compressed_tensors: + return query, key, value, q_compressed, kv_compressed + else: + return query, key, value def uncompress_kv_from_cache(self, kv_cached): """ diff --git a/pyproject.toml b/pyproject.toml index 7f734927c1a..c784f829e06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,6 +87,7 @@ dev = [ "onnxscript", "flash-linear-attention~=0.3.2", "emerging_optimizers", + "fast_hadamard_transform==1.0.4post1", ] lts = [ From 33684d29a3c0a1e5ac12e2f146e3afaa3b7035d6 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 28 Nov 2025 18:21:30 +0800 Subject: [PATCH 24/28] Fix qk norm spec in dsa spec Signed-off-by: kunlunl --- .../experimental_attention_variant_module_specs.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index cd8cd1fefd1..cbe59618baf 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -11,6 +11,7 @@ DSAttention, DSAttentionSubmodules, ) +from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.multi_latent_attention import ( MLASelfAttention, MLASelfAttentionSubmodules, @@ -37,6 +38,7 @@ def get_gated_delta_net_module_spec_for_backend( 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, mla_down_proj_use_column_parallel: Optional[bool] = False, @@ -48,10 +50,6 @@ def get_dsa_module_spec_for_backend( assert qk_l2_norm is False, "qk_l2_norm is not supported with MLA." assert fallback_to_eager_attn is False, "Fallback to eager attention is not supported with DSA." - # Adjust for RMS norm. - rms_norm = normalization == "RMSNorm" - qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) - linear_q_down_proj = ( backend.column_parallel_linear() if mla_down_proj_use_column_parallel else backend.linear() ) @@ -60,7 +58,6 @@ def get_dsa_module_spec_for_backend( ) linear_q_up_proj = backend.column_parallel_linear() linear_kv_up_proj = backend.column_parallel_linear() - qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) # Because TransformerEngine does not support sparse attention yet, we use local # implementation whether the backend is TransformerEngine or not. @@ -79,6 +76,10 @@ def get_dsa_module_spec_for_backend( ), ) + # Adjust for RMS norm. + rms_norm = normalization == "RMSNorm" + qk_norm = backend.layer_norm(rms_norm=rms_norm, for_qk=True) if qk_layernorm else IdentityOp + attention = ModuleSpec( module=MLASelfAttention, params={"attn_mask_type": AttnMaskType.causal}, @@ -118,6 +119,7 @@ def get_experimental_attention_variant_module_spec_for_backend( elif 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, mla_down_proj_use_column_parallel=mla_down_proj_use_column_parallel, From a0b6fd995ccf52606725bb8702b84f3c8cc50c5d Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 28 Nov 2025 19:00:15 +0800 Subject: [PATCH 25/28] Add fast-hadamard-transform to pyproject.toml Signed-off-by: kunlunl --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c784f829e06..265fcad3b26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ dev = [ "onnxscript", "flash-linear-attention~=0.3.2", "emerging_optimizers", - "fast_hadamard_transform==1.0.4post1", + "fast_hadamard_transform~=1.0.4post1", ] lts = [ @@ -153,6 +153,7 @@ no-build-isolation-package = [ "mamba-ssm", "transformer-engine", "transformer-engine-torch", + "fast-hadamard-transform", ] link-mode = "copy" conflicts = [[{ extra = "lts" }, { extra = "dev" }]] From f22344ae3ee6d2efd12c15f92793e555d01b9dde Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 1 Dec 2025 10:55:28 +0800 Subject: [PATCH 26/28] Remove fast-hadamard-transform Signed-off-by: kunlunl --- pyproject.toml | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 265fcad3b26..7f734927c1a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,6 @@ dev = [ "onnxscript", "flash-linear-attention~=0.3.2", "emerging_optimizers", - "fast_hadamard_transform~=1.0.4post1", ] lts = [ @@ -153,7 +152,6 @@ no-build-isolation-package = [ "mamba-ssm", "transformer-engine", "transformer-engine-torch", - "fast-hadamard-transform", ] link-mode = "copy" conflicts = [[{ extra = "lts" }, { extra = "dev" }]] From ce99e9f8c579fd4cb829f49c9e9c03e26639de61 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 1 Dec 2025 11:01:52 +0800 Subject: [PATCH 27/28] Minor fix for args Signed-off-by: kunlunl --- megatron/training/arguments.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index ffdc7446d1e..b93a7597429 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1205,6 +1205,14 @@ def validate_args(args, defaults={}): args.no_load_rng = True print('Warning: disabling --no-load-rng for upcycling.') + if args.linear_attention_type is not None: + print_rank_0( + '--linear-attention-type is deprecated, use --experimental-attention-variant instead.', + args.rank, + ) + args.experimental_attention_variant = args.linear_attention_type + del args.linear_attention_type + # Muon optimizercheck if 'muon' in args.optimizer: assert not args.use_distributed_optimizer, "Muon optimizer does not support distributed optimizer for now." @@ -1280,14 +1288,6 @@ def validate_args(args, defaults={}): if args.multi_latent_attention: assert not args.group_query_attention, "Group query attention is mutually exclusive with multi latent attention." - if args.linear_attention_type is not None: - print_rank_0( - '--linear-attention-type is deprecated, use --experimental-attention-variant instead.', - args.rank, - ) - args.experimental_attention_variant = args.linear_attention_type - del args.linear_attention_type - # Print arguments. _print_args("arguments", args) From cbfa053e8b97082ec27324c47f7724696e5554f1 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 1 Dec 2025 11:55:30 +0800 Subject: [PATCH 28/28] Add mock hadamard_transformer Signed-off-by: kunlunl --- .../transformer/test_attention_variant_dsa.py | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/transformer/test_attention_variant_dsa.py b/tests/unit_tests/transformer/test_attention_variant_dsa.py index 6f45862c7c6..bd106aa6f0e 100644 --- a/tests/unit_tests/transformer/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/test_attention_variant_dsa.py @@ -1,5 +1,7 @@ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +from unittest.mock import patch + import pytest import torch @@ -22,14 +24,35 @@ from tests.unit_tests.test_utilities import Utils try: - from fast_hadamard_transform import hadamard_transform + 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 -@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") class TestRotateActivation: """Test rotate_activation function.""" @@ -242,7 +265,6 @@ def test_backward_pass(self): ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}" -@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @pytest.mark.parametrize("seqlen", [16, 64]) class TestDSAIndexer: """Test DSA Indexer module basic functionality with TP=1.""" @@ -382,7 +404,6 @@ def test_dsa_indexer_with_mask(self, seqlen): assert torch.all(topk_indices[b, i] <= max(self.index_topk, i)) -@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") class TestDSAttention: """Test DSAttention module basic functionality with TP=1.""" @@ -615,7 +636,6 @@ def test_dsa_topk_selection(self): # ====================================================================================== -@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @pytest.mark.parametrize("tensor_model_parallel_size", [2, 4, 8]) @pytest.mark.parametrize("sequence_parallel", [False, True]) class TestIndexerTensorParallel: @@ -853,7 +873,6 @@ def test_dsa_indexer_gradient_sync(self, tensor_model_parallel_size, sequence_pa Utils.destroy_model_parallel() -@pytest.mark.skipif(not HAVE_HADAMARD, reason="fast_hadamard_transform not installed") @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])