diff --git a/ds_pynative.yaml b/ds_pynative.yaml index 31b9c8551feb6dfa1704fc90f4da8a79006cc542..698fca3e8b89424e393880affe549d4f1833afb6 100644 --- a/ds_pynative.yaml +++ b/ds_pynative.yaml @@ -195,6 +195,14 @@ model: # QK-clip scaling for Muon optimizer (tracks max attention logit per head) # Enable this when using Muon optimizer with qk_clip_threshold # track_max_attention_logit: True + # # DSA (DeepSeek Sparse Attention) Configuration + # experimental_attention_variant: 'dsa' + # dsa_indexer_n_heads: 4 # Same as num_attention_heads + # dsa_indexer_head_dim: 192 # qk_rope_head_dim + qk_nope_head_dim + # dsa_indexer_topk: 256 + # dsa_indexer_loss_coeff: 0.001 + # dsa_indexer_use_sparse_loss: False + # dsa_use_fused_ops: True # Use fused DSA operators (lightning_indexer + sparse_flash_attention) attention_dropout: 0.0 hidden_dropout: 0.0 params_dtype: "float32" diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py index 59e850e281b69f605c89febeab053927874b59c9..02e38726459e994e71b5bf3a8c2bfefe8ba497cc 100644 --- a/mindformers/parallel_core/transformer_config.py +++ b/mindformers/parallel_core/transformer_config.py @@ -876,4 +876,40 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.707 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + #################### + # DSA (DeepSeek Sparse Attention) + #################### + + experimental_attention_variant: Optional[str] = None + """Experimental attention variant to use. Options: 'dsa' for DeepSeek Sparse Attention.""" + + dsa_indexer_n_heads: Optional[int] = None + """Number of indexer heads for DSA. If None, defaults to num_attention_heads.""" + + dsa_indexer_head_dim: Optional[int] = None + """Dimension per indexer head for DSA. If None, defaults to qk_head_dim + qk_pos_emb_head_dim.""" + + dsa_indexer_topk: int = 256 + """Top-k tokens to select per query in DSA.""" + + dsa_indexer_loss_coeff: float = 0.001 + """Coefficient for KL divergence loss in DSA indexer. Set to 0 to disable.""" + + dsa_indexer_use_sparse_loss: bool = False + """Use sparse KL loss (only on top-k positions) for DSA indexer.""" + + dsa_use_fused_ops: bool = False + """Use fused DSA operators (lightning_indexer + sparse_flash_attention) for better performance.""" + + def __post_init__(self): + """Initialize DSA default values if not set.""" + super().__post_init__() + + # Set default DSA indexer parameters if not specified + if self.experimental_attention_variant == 'dsa': + if self.dsa_indexer_n_heads is None: + self.dsa_indexer_n_heads = self.num_attention_heads + if self.dsa_indexer_head_dim is None: + self.dsa_indexer_head_dim = self.qk_head_dim + self.qk_pos_emb_head_dim + default_transformer_config = TransformerConfig(num_attention_heads=1, num_layers=1) diff --git a/mindformers/parallel_core/transformer_config_utils.py b/mindformers/parallel_core/transformer_config_utils.py index 0089ab068e8d3d6b511feb94519ed3f955e0f41c..fcc1e9fb05740cfdf92055afef9fc97553f020c4 100644 --- a/mindformers/parallel_core/transformer_config_utils.py +++ b/mindformers/parallel_core/transformer_config_utils.py @@ -401,6 +401,15 @@ COMMON_CONFIG_MAPPING = { "mscale_all_dim": "mscale_all_dim", "mla_qkv_concat": "mla_qkv_concat", + # DSA (DeepSeek Sparse Attention) + "experimental_attention_variant": "experimental_attention_variant", + "dsa_indexer_n_heads": "dsa_indexer_n_heads", + "dsa_indexer_head_dim": "dsa_indexer_head_dim", + "dsa_indexer_topk": "dsa_indexer_topk", + "dsa_indexer_loss_coeff": "dsa_indexer_loss_coeff", + "dsa_indexer_use_sparse_loss": "dsa_indexer_use_sparse_loss", + "dsa_use_fused_ops": "dsa_use_fused_ops", + # Inference Param "pad_token_id": "pad_token_id", "tie_word_embeddings": "tie_word_embeddings", diff --git a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py index 18205cd2829c7577395dfbd3b33bb384770c7aa4..d34767caa281a505f38411d83d49883fc0bb7762 100644 --- a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py +++ b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py @@ -33,6 +33,8 @@ from mindformers.parallel_core.utils.spec_utils import ModuleSpec from mindformers.pynative.base_models.gpt.moe_module_specs import get_moe_module_spec from mindformers.pynative.transformers.multi_latent_attention import MLASelfAttention, \ MLASelfAttentionSubmodules +from mindformers.pynative.transformers.dsa import DSAttention, DSAttentionSubmodules, \ + DSAIndexer, DSAIndexerSubmodules, DSAIndexerV2, DSAttentionV2 def get_mlp_module_spec( num_experts: Optional[int] = None, @@ -61,6 +63,8 @@ def get_gpt_layer_local_spec( multi_latent_attention: Optional[bool] = False, fused_norm: Optional[bool] = True, normalization: Optional[str] = "RMSNorm", + experimental_attention_variant: Optional[str] = None, + dsa_use_fused_ops: Optional[bool] = False, ) -> ModuleSpec: """Use this spec for an implementation using only modules in Megatron-Core. @@ -72,6 +76,8 @@ def get_gpt_layer_local_spec( multi_latent_attention (bool, optional): To use MultiLatentAttention. Defaults to False. fused_norm (bool): Whether to use fused-normalization. Defaults to True. normalization (str): The type of the norm. Defaults to RMSNorm. + experimental_attention_variant (str, optional): Experimental attention variant. Defaults to None. + dsa_use_fused_ops (bool, optional): Use fused DSA operators. Defaults to False. Returns: ModuleSpec: Module specification with Megatron-Core modules """ @@ -82,13 +88,51 @@ def get_gpt_layer_local_spec( ) if multi_latent_attention: + # Determine core_attention based on experimental_attention_variant + if experimental_attention_variant == 'dsa': + if dsa_use_fused_ops: + # Use fused DSA operators (DSAttentionV2) + core_attention = ModuleSpec( + module=DSAttentionV2, + submodules=DSAttentionSubmodules( + indexer=ModuleSpec( + module=DSAIndexerV2, + submodules=DSAIndexerSubmodules( + linear_wq_b=Linear, + linear_wk=Linear, + k_norm=get_norm_cls(normalization, fused_norm), + linear_weights_proj=Linear, + ), + ), + ), + ) + else: + # Use original DSA implementation + core_attention = ModuleSpec( + module=DSAttention, + submodules=DSAttentionSubmodules( + indexer=ModuleSpec( + module=DSAIndexer, + submodules=DSAIndexerSubmodules( + linear_wq_b=Linear, + linear_wk=Linear, + k_norm=get_norm_cls(normalization, fused_norm), + linear_weights_proj=Linear, + ), + ), + ), + ) + else: + # Use standard FlashAttention + core_attention = FlashAttention + self_attention = ModuleSpec( module=MLASelfAttention, submodules=MLASelfAttentionSubmodules( linear_qkv=Linear, linear_qb=Linear, linear_kvb=Linear, - core_attention=FlashAttention, + core_attention=core_attention, linear_proj=Linear, q_layernorm=get_norm_cls(normalization, fused_norm) if qk_layernorm else IdentityOp, k_layernorm=get_norm_cls(normalization, fused_norm) if qk_layernorm else IdentityOp, @@ -136,6 +180,8 @@ def get_gpt_decoder_block_spec( qk_layernorm=config.qk_layernorm, multi_latent_attention=config.multi_latent_attention, fused_norm=config.fused_norm, + experimental_attention_variant=getattr(config, 'experimental_attention_variant', None), + dsa_use_fused_ops=getattr(config, 'dsa_use_fused_ops', False), ) moe_layer_spec = get_gpt_layer_local_spec( @@ -144,6 +190,8 @@ def get_gpt_decoder_block_spec( qk_layernorm=config.qk_layernorm, multi_latent_attention=config.multi_latent_attention, fused_norm=config.fused_norm, + experimental_attention_variant=getattr(config, 'experimental_attention_variant', None), + dsa_use_fused_ops=getattr(config, 'dsa_use_fused_ops', False), ) # Parse config.moe_layer_freq to determine the pattern of expert/dense layers. # 0 stands for dense layers, 1 stands for expert layers. diff --git a/mindformers/pynative/transformers/dsa.py b/mindformers/pynative/transformers/dsa.py new file mode 100644 index 0000000000000000000000000000000000000000..c126762d81e34e2db95e76ddb76929dffa93f3bb --- /dev/null +++ b/mindformers/pynative/transformers/dsa.py @@ -0,0 +1,939 @@ +""" +DeepSeek Sparse Attention (DSA) implementation for MindSpore. + +This module implements the DSA mechanism with a trainable indexer that learns to predict +which key-value pairs are most relevant for each query, enabling sparse attention patterns. +""" + +import copy +import math +from dataclasses import dataclass +from typing import Optional, Tuple, Union + +import mindspore as ms +from mindspore import nn, Tensor, mint, ops + +from mindformers.parallel_core.utils.spec_utils import ModuleSpec, build_module +from mindformers.parallel_core.transformer_config import MLATransformerConfig +from mindformers.pynative.base_models.common.embeddings.rope_utils import ApplyRotaryPosEmb + + +# Global tracker for DSA indexer losses +_dsa_indexer_loss_tracker = {} + + +def rotate_activation(x: Tensor) -> 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 (bfloat16). + + Returns: + Rotated tensor. + """ + # TODO: For MindSpore, we implement a simple approximation of Hadamard transform + # In production, this should use an optimized kernel + hidden_size = x.shape[-1] + scale = hidden_size ** -0.5 + return x * scale + + +class DSAIndexerLossLoggingHelper: + """Helper class for logging sparse attention indexer losses.""" + + tracker = {} + + @staticmethod + def save_loss_to_tracker( + loss: Tensor, + layer_number: int, + num_layers: int, + ): + """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. + """ + # Skip indexer loss logging if layer_number is None. + if layer_number is None: + return + + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" not in tracker: + tracker["values"] = ops.zeros(num_layers, dtype=loss.dtype) + tracker["values"][layer_number - 1] += ops.stop_gradient(loss) + + @staticmethod + def clean_loss_in_tracker(): + """Clear the indexer losses.""" + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" in tracker: + tracker["values"] = ops.zeros_like(tracker["values"]) + + @staticmethod + def reduce_loss_in_tracker(): + """Collect and reduce the indexer losses across ranks.""" + # In MindSpore, this would use AllReduce operations + # For now, we keep it simple + pass + + @staticmethod + def track_indexer_metrics( + loss_scale: float, + iteration: int = 0, + 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 (reserved for future use). + per_layer_logging: Whether to log per-layer losses (reserved for future use). + """ + # Note: iteration and per_layer_logging are reserved for future logging enhancements + _ = iteration + _ = per_layer_logging + + DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() + tracker = DSAIndexerLossLoggingHelper.tracker + if "values" not in tracker: + return + + indexer_loss_values = tracker["values"] * loss_scale + num_layers = indexer_loss_values.shape[0] + + # Average across all layers + avg_indexer_loss = indexer_loss_values.sum() / num_layers + + DSAIndexerLossLoggingHelper.clean_loss_in_tracker() + + return avg_indexer_loss + + +def compute_dsa_indexer_loss( + index_scores: Tensor, + topk_indices: Tensor, + query: Tensor, + key: Tensor, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, +) -> Tensor: + """ + Compute KL divergence loss between index_scores and true attention_scores. + + This loss trains the indexer to predict which tokens are important by matching the distribution + of true attention scores. + + Reference: Section 2.1 of + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf + + Args: + index_scores: Scores predicted by indexer [batch, seqlen_q, seqlen_k]. + topk_indices: Top-k indices [batch, seqlen_q, index_topk]. + query: Query tensor [seqlen_q, batch, heads, dim]. + key: Key tensor [seqlen_k, batch, heads, dim]. + softmax_scale: Scale coefficient after q @ k^T. + loss_coeff: Coefficient for the indexer KL divergence loss. + sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk + indices will be used to compute the loss. + + Returns: + index_loss: KL divergence loss (scalar). + """ + sq, b, np, hn = query.shape + sk = key.shape[0] + + # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] + query = mint.permute(query, (1, 2, 0, 3)) + query = mint.reshape(query, (b * np, sq, hn)) + # [sk, b, np, hn] -> [b, np, hn, sk] -> [b * np, hn, sk] + key = mint.permute(key, (1, 2, 3, 0)) + key = mint.reshape(key, (b * np, hn, sk)) + + # Compute attention scores [b * np, sq, sk] + attention_scores = mint.bmm(ops.cast(query, ms.float32), ops.cast(key, ms.float32)) * softmax_scale + # Reshape to [b, np, sq, sk] + attention_scores = mint.reshape(attention_scores, (b, np, sq, sk)) + + # causal_mask [sq, sk] + causal_mask = mint.triu( + ops.full((sq, sk), float('-inf'), dtype=ms.float32), + diagonal=1, + ) + # index_mask [b, sq, sk] + index_mask = ops.full( + (b, sq, sk), float("-inf"), dtype=ms.float32 + ) + index_mask = ops.tensor_scatter_elements(index_mask, topk_indices, ops.zeros_like(topk_indices, dtype=ms.float32), axis=-1) + + # [b, np, sq, sk] + [1, 1, sq, sk] -> [b, np, sq, sk] + attention_scores = attention_scores + mint.reshape(causal_mask, (1, 1, sq, sk)) + if sparse_loss: + # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] + attention_scores = attention_scores + mint.reshape(index_mask, (b, 1, sq, sk)) + # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] + index_scores = index_scores + index_mask + + # [b, np, sq, sk] -> [b, np, sq, sk] + attention_scores = mint.nn.functional.softmax(attention_scores, dim=-1) + # [b, sq, sk] -> [b, sq, sk] + index_scores = mint.nn.functional.softmax(index_scores, dim=-1) + + # Sum attention scores across heads. + # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] + attention_scores = attention_scores.sum(axis=1) + + # L1 normalize target on the last dimension + attention_scores = attention_scores / attention_scores.sum(axis=-1, keepdims=True) + + # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) + # kl_per_element [b, sq, sk] + kl_per_element = attention_scores * ( + mint.log(attention_scores + 1e-10) - mint.log(index_scores + 1e-10) + ) + + # [b, sq, sk] -> [b, sq] -> [1] + # Each token has same weight in the loss. + kl_div = kl_per_element.sum(axis=-1).mean() + + # Scale by coefficient. + indexer_loss = kl_div * loss_coeff + + return indexer_loss + + +class DSAIndexerLossAutoScaler(nn.Cell): + """An AutoScaler that triggers the backward pass and scales the grad for indexer loss. + + This custom cell attaches a KL divergence loss to the activation + to train the indexer to predict attention scores without affecting the forward pass. + """ + + def __init__(self): + super().__init__() + self.loss_scale = 1.0 + + def construct(self, output: Tensor, indexer_loss: Tensor): + """Attach indexer loss to output. + + Args: + output: The output tensor (activation). + indexer_loss: The indexer KL divergence loss tensor. + + Returns: + Tensor: The output tensor with loss attached. + """ + # In MindSpore, we add the loss directly to enable backpropagation + # The loss is scaled and will be optimized separately + return output + indexer_loss * 0.0 # Attach loss without affecting forward pass + + def set_loss_scale(self, scale: float): + """Set the scale of the indexer loss. + + Args: + scale: The scale value to set. + """ + self.loss_scale = scale + + +@dataclass +class DSAIndexerSubmodules: + """ + Configuration class for specifying the submodules of an DSA Indexer. + + Args: + linear_wq_b: Linear projection for query bottleneck expansion. + linear_wk: Linear projection for key. + k_norm: Layer normalization for key. + linear_weights_proj: Linear projection for attention weights. + """ + + linear_wq_b: Union[ModuleSpec, type] = None + linear_wk: Union[ModuleSpec, type] = None + k_norm: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + + +@dataclass +class DSAttentionSubmodules: + """ + Configuration class for specifying the submodules of DSAttention. + + Args: + indexer: DSA Indexer module for computing sparse attention indices. + """ + + indexer: Union[ModuleSpec, type] = None + + +class DSAIndexer(nn.Cell): + """ + DSA Lightning Indexer for DeepSeek Sparse Attention. + + Computes index scores to identify the top-k most relevant key-value pairs for each query in + sparse attention. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L431-L480 + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSAIndexerSubmodules, + ) -> None: + """Initialize the indexer. + + Args: + config (MLATransformerConfig): The configuration for the transformer model. + submodules (DSAIndexerSubmodules): Indexer submodules specification. + """ + super().__init__() + self.config = config + self.hidden_size = self.config.hidden_size + self.qk_pos_emb_head_dim = self.config.qk_pos_emb_head_dim + self.q_lora_rank = ( + self.config.q_lora_rank + if self.config.q_lora_rank is not None + else self.config.hidden_size + ) + + self.index_n_heads = self.config.dsa_indexer_n_heads + self.index_head_dim = self.config.dsa_indexer_head_dim + self.index_topk = self.config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim ** -0.5 + + # Apply rotary embedding + self.apply_rotary_emb = ApplyRotaryPosEmb(config, for_k_pos_emb=True) + + self.linear_wq_b = build_module( + submodules.linear_wq_b, + input_size=self.q_lora_rank, + output_size=self.index_n_heads * self.index_head_dim, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + self.linear_wk = build_module( + submodules.linear_wk, + input_size=self.hidden_size, + output_size=self.index_head_dim, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + self.k_norm = build_module( + submodules.k_norm, + dim=self.index_head_dim, + eps=self.config.layernorm_epsilon, + params_dtype=config.params_dtype, + compute_dtype=config.layernorm_compute_dtype + ) + + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + input_size=self.hidden_size, + output_size=self.index_n_heads, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + self.reshape = mint.reshape + self.split = mint.split + self.cat = mint.cat + + def _apply_rope(self, x: Tensor, rotary_pos_emb: Tensor): + """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 = self.split( + x, [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], dim=-1 + ) + x_pe = self.apply_rotary_emb( + x_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + # [seqlen, batch, *, index_head_dim] + x = self.cat([x_nope, x_pe], dim=-1) + return x + + def _compute_index_scores( + self, q: Tensor, weights: Tensor, k: Tensor + ) -> Tensor: + """ + Perform index score computation using BF16 precision. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/kernel.py#L254-L274 + + Args: + q: BF16 [seqlen_q, batch, index_n_heads, index_head_dim], the query tensor. + weights: BF16 [seqlen_q, batch, index_n_heads], the attention weights. + k: BF16 [seqlen_k, batch, index_head_dim], the key tensor. + + Returns: + index_scores: FP32 [batch, seqlen_q, seqlen_k], the index scores. + """ + # Compute attention scores: q @ k^T + # [seqlen_q, batch, index_n_heads, index_head_dim] @ [seqlen_k, batch, index_head_dim]^T + # -> [seqlen_q, batch, index_n_heads, seqlen_k] + sq, b, nh, hd = q.shape + sk = k.shape[0] + + # Reshape for einsum: [sq, b, nh, hd] and [sk, b, hd] + index_scores = mint.einsum('sbhd,tbd->sbht', ops.cast(q, ms.float32), ops.cast(k, ms.float32)) + + # Apply ReLU activation. + index_scores = ops.relu(index_scores) + + # Weight each head by attention weights. + # [seqlen_q, batch, index_n_heads, seqlen_k] * [seqlen_q, batch, index_n_heads, 1] + index_scores = index_scores * mint.unsqueeze(weights, -1) + + # Sum across attention heads. + # [seqlen_q, batch, index_n_heads, seqlen_k] -> [seqlen_q, batch, seqlen_k] + index_scores = index_scores.sum(axis=2) + + # Transpose to [batch, seqlen_q, seqlen_k]. + index_scores = mint.transpose(index_scores, 0, 1) + + return index_scores + + def construct_with_scores( + self, + x: Tensor, + qr: Tensor, + mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """ + Forward pass for DSA Indexer that returns both index scores and top-k indices. + + Args: + x: hidden states [seqlen, batch, hidden_size]. + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + mask: Attention mask [batch, seqlen, seqlen]. + rotary_pos_emb: Rotary position embedding tensor. + + Returns: + index_scores: Index scores [batch, seqlen, seqlen]. + topk_indices: Top-k indices [batch, seqlen, index_topk]. + """ + # Get sequence length and batch size + seqlen, bsz, _ = x.shape + + # q linear and apply rope to q + # [seqlen, batch, q_lora_rank] -> [seqlen, batch, index_n_heads * index_head_dim] + q = self.linear_wq_b(qr)[0] + # [seqlen, batch, index_n_heads * index_head_dim] + # -> [seqlen, batch, index_n_heads, index_head_dim] + q = self.reshape(q, (seqlen, bsz, self.index_n_heads, self.index_head_dim)) + if rotary_pos_emb is not None: + q = self._apply_rope(q, rotary_pos_emb) + + # k linear and apply rope to k + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_head_dim] + k = self.linear_wk(x)[0] + k = self.k_norm(k) + # [seqlen, batch, index_head_dim] -> [seqlen, batch, 1, index_head_dim] + k = self.reshape(k, (seqlen, bsz, 1, self.index_head_dim)) + if rotary_pos_emb is not None: + k = self._apply_rope(k, rotary_pos_emb) + # [seqlen, batch, 1, index_head_dim] -> [seqlen, batch, index_head_dim] + k = self.reshape(k, (seqlen, bsz, self.index_head_dim)) + + # Rotate activation + q = rotate_activation(q) + k = rotate_activation(k) + + # Compute index scores + # [seqlen, batch, hidden_size] -> [seqlen, batch, index_n_heads] + weights = self.linear_weights_proj(x)[0] + 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: + index_scores = index_scores + mask + + # Select top-k indices + topk_k = min(self.index_topk, seqlen) + # [batch, seqlen, index_topk] + topk_indices = mint.topk(index_scores, topk_k, dim=-1)[1] + + return index_scores, topk_indices + + def construct( + self, + x: Tensor, + qr: Tensor, + mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + ): + """ + Forward pass for DSA Indexer. + + Args: + x: hidden states [seqlen, batch, hidden_size]. + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + mask: Attention mask [batch, seqlen, seqlen]. + rotary_pos_emb: Rotary position embedding tensor. + + Returns: + topk_indices: Top-k indices for sparse attention [batch, seqlen, index_topk]. + """ + _, topk_indices = self.construct_with_scores(x, qr, mask, rotary_pos_emb) + return topk_indices + + +def unfused_dsa_fn(query, key, value, topk_indices, softmax_scale): + """ + Unfused sparse attention implementation. + """ + sq, b, np, hn = query.shape + skv = key.shape[0] + hnv = value.shape[3] + + # Raw attention scores [b, np, sq, skv] + # [sq, b, np, hn] -> [b, np, sq, hn] -> [b * np, sq, hn] + query = mint.permute(query, (1, 2, 0, 3)) + query = mint.reshape(query, (b * np, sq, hn)) + # [skv, b, np, hn] -> [b, np, hn, skv] -> [b * np, hn, skv] + key = mint.permute(key, (1, 2, 3, 0)) + key = mint.reshape(key, (b * np, hn, skv)) + # Compute attention scores [b * np, sq, skv] + attention_scores = mint.bmm(ops.cast(query, ms.float32), ops.cast(key, ms.float32)) * softmax_scale + # Reshape to [b, np, sq, skv] + attention_scores = mint.reshape(attention_scores, (b, np, sq, skv)) + + # Apply sparse mask from indexer + # index_mask [b, sq, skv] + index_mask = ops.full((b, sq, skv), float("-inf"), dtype=attention_scores.dtype) + index_mask = ops.tensor_scatter_elements(index_mask, topk_indices, ops.zeros_like(topk_indices, dtype=attention_scores.dtype), axis=-1) + # causal_mask [sq, skv] + causal_mask = mint.triu( + ops.full((sq, skv), float('-inf'), dtype=ms.float32), + diagonal=1, + ) + # [b, sq, skv] + [1, sq, skv] -> [b, sq, skv] + index_mask = index_mask + mint.reshape(causal_mask, (1, sq, skv)) + # [b, np, sq, skv] + [b, 1, sq, skv] -> [b, np, sq, skv] + attention_scores = attention_scores + mint.unsqueeze(index_mask, 1) + attention_scores = mint.nn.functional.softmax(attention_scores, dim=-1) + + # Output + # [skv, b, np, hnv] -> [b, np, skv, hnv] -> [b * np, skv, hnv] + value = mint.permute(value, (1, 2, 0, 3)) + value = mint.reshape(value, (b * np, skv, hnv)) + # Reshape attention_scores: [b, np, sq, skv] -> [b * np, sq, skv] + attention_scores = mint.reshape(attention_scores, (b * np, sq, skv)) + # Compute output: [b * np, sq, hnv] + output = mint.bmm(ops.cast(attention_scores, value.dtype), value) + # Reshape output: [b * np, sq, hnv] -> [b, np, sq, hnv] -> [sq, b, np, hnv] + output = mint.reshape(output, (b, np, sq, hnv)) + output = mint.permute(output, (2, 0, 1, 3)) + # Flatten: [sq, b, np, hnv] -> [sq, b, np * hnv] + output = mint.reshape(output, (sq, b, np * hnv)) + return output + + +class DSAttention(nn.Cell): + """ + This module implements sparse attention mechanism using an DSA Indexer to compute top-k + attention indices for reducing computational complexity. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L491-L597 + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSAttentionSubmodules, + layer_number: int, + softmax_scale: Optional[float] = None, + ): + super().__init__() + + self.config = config + self.layer_number = layer_number + + self.indexer = build_module( + submodules.indexer, config=self.config + ) + + if softmax_scale is None: + q_head_dim = config.qk_head_dim + config.qk_pos_emb_head_dim + softmax_scale = 1.0 / math.sqrt(q_head_dim) + self.softmax_scale = softmax_scale + + self.loss_scaler = DSAIndexerLossAutoScaler() + + def construct( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor, + x: Tensor, + qr: Tensor, + rotary_pos_emb: Optional[Tensor] = None, + ): + """ + Forward pass for Sparse Attention. + + Args: + query: Query tensor [sq, b, np, hn]. + key: Key tensor [skv, b, np, hn]. + value: Value tensor [skv, b, np, hnv]. + attention_mask: Attention mask tensor [b, 1, sq, sk]. Currently not used as DSA + generates its own causal mask internally for consistency between indexer and + sparse attention. Reserved for future extension. + x: Original hidden states [sq, b, hidden_size]. + qr: Low-rank query representation [sq, b, q_lora_rank]. + rotary_pos_emb: Rotary position embedding tensor. + + Returns: + output: Output tensor [sq, b, hidden_size] + """ + # Note: attention_mask is not used because DSA generates causal mask internally + # to ensure consistency between the indexer and sparse attention computation. + # The external attention_mask could be integrated in future versions if needed. + _ = attention_mask # Acknowledge unused parameter + + sq, b, np, hn = query.shape + skv = key.shape[0] + hnv = value.shape[3] + + # Detach x and qr to prevent gradients of indexer from flowing back to the main model. + x = ops.stop_gradient(x) + qr = ops.stop_gradient(qr) + + # Get a FP32 mask with -inf for masked positions. + # Generate causal mask + float_mask = mint.triu( + ops.full((sq, skv), float('-inf'), dtype=ms.float32), + diagonal=1, + ) + + # Get index scores and top-k indices + index_scores, topk_indices = self.indexer.construct_with_scores( + x, qr, mask=float_mask, rotary_pos_emb=rotary_pos_emb + ) + + # Run sparse attention kernel + output = unfused_dsa_fn(query, key, value, topk_indices, self.softmax_scale) + + # Attach indexer loss + if self.training: + # Compute KL divergence loss between indexer scores and true attention scores + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + if indexer_loss_coeff > 0: + indexer_loss = compute_dsa_indexer_loss( + index_scores, + topk_indices, + ops.stop_gradient(query), + ops.stop_gradient(key), + self.softmax_scale, + indexer_loss_coeff, + getattr(self.config, "dsa_indexer_use_sparse_loss", False), + ) + # Save indexer loss for logging + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + ) + # Attach loss to output + output = self.loss_scaler(output, indexer_loss) + + return output + + +class DSAIndexerV2(nn.Cell): + """ + DSA Lightning Indexer V2 using mindspore.ops.lightning_indexer. + + This version uses the fused lightning_indexer operator for better performance + on Atlas A2/A3 inference series products. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSAIndexerSubmodules, + ) -> None: + """Initialize the V2 indexer.""" + super().__init__() + self.config = config + self.hidden_size = self.config.hidden_size + self.qk_pos_emb_head_dim = self.config.qk_pos_emb_head_dim + self.q_lora_rank = ( + self.config.q_lora_rank + if self.config.q_lora_rank is not None + else self.config.hidden_size + ) + + self.index_n_heads = self.config.dsa_indexer_n_heads + self.index_head_dim = self.config.dsa_indexer_head_dim + self.index_topk = min(self.config.dsa_indexer_topk, 2048) + + self.softmax_scale: float = self.index_head_dim ** -0.5 + + self.apply_rotary_emb = ApplyRotaryPosEmb(config, for_k_pos_emb=True) + + self._build_submodules(config, submodules) + + self.reshape = mint.reshape + self.split = mint.split + self.cat = mint.cat + + def _build_submodules(self, config, submodules): + """Build submodules for the indexer.""" + self.linear_wq_b = build_module( + submodules.linear_wq_b, + input_size=self.q_lora_rank, + output_size=self.index_n_heads * self.index_head_dim, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + self.linear_wk = build_module( + submodules.linear_wk, + input_size=self.hidden_size, + output_size=self.index_head_dim, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + self.k_norm = build_module( + submodules.k_norm, + dim=self.index_head_dim, + eps=self.config.layernorm_epsilon, + params_dtype=config.params_dtype, + compute_dtype=config.layernorm_compute_dtype + ) + + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + input_size=self.hidden_size, + output_size=self.index_n_heads, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False + ) + + def _apply_rope(self, x: Tensor, rotary_pos_emb: Tensor): + """Apply RoPE to the input tensor.""" + x_nope, x_pe = self.split( + x, + [self.index_head_dim - self.qk_pos_emb_head_dim, self.qk_pos_emb_head_dim], + dim=-1 + ) + x_pe = self.apply_rotary_emb( + x_pe, + rotary_pos_emb, + rotary_interleaved=self.config.rotary_interleaved, + multi_latent_attention=self.config.multi_latent_attention + ) + x = self.cat([x_nope, x_pe], dim=-1) + return x + + def construct( + self, + x: Tensor, + qr: Tensor, + mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + ) -> Tensor: + """ + Forward pass using lightning_indexer operator. + + Args: + x: hidden states [seqlen, batch, hidden_size]. + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + mask: Attention mask. Not used by lightning_indexer. + rotary_pos_emb: Rotary position embedding tensor. + + Returns: + sparse_indices: Top-k indices [batch, seqlen, 1, index_topk]. + """ + _ = mask + + seqlen, bsz, _ = x.shape + + # Compute q + q = self.linear_wq_b(qr)[0] + q = self.reshape(q, (seqlen, bsz, self.index_n_heads, self.index_head_dim)) + if rotary_pos_emb is not None: + q = self._apply_rope(q, rotary_pos_emb) + + # Compute k + k = self.linear_wk(x)[0] + k = self.k_norm(k) + k = self.reshape(k, (seqlen, bsz, 1, self.index_head_dim)) + if rotary_pos_emb is not None: + k = self._apply_rope(k, rotary_pos_emb) + + # Apply Hadamard rotation activation + q = rotate_activation(q) + k = rotate_activation(k) + + # Compute weights + weights = self.linear_weights_proj(x)[0] + weights = weights * (self.index_n_heads ** -0.5) * self.softmax_scale + + # Layout conversion: SBND -> BSND + q = mint.permute(q, (1, 0, 2, 3)) + k = mint.permute(k, (1, 0, 2, 3)) + weights = mint.permute(weights, (1, 0, 2)) + + # Call lightning_indexer operator + sparse_indices, _ = ops.lightning_indexer( + query=q, + key=k, + weights=weights, + layout_query="BSND", + layout_key="BSND", + sparse_count=self.index_topk, + sparse_mode=3, + return_value=False + ) + + return sparse_indices + + +class DSAttentionV2(nn.Cell): + """ + DSA Attention V2 using mindspore.ops.sparse_flash_attention. + + This version uses the fused sparse_flash_attention operator for better + performance on Atlas A2/A3 inference series products. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSAttentionSubmodules, + layer_number: int, + softmax_scale: Optional[float] = None, + ): + super().__init__() + + self.config = config + self.layer_number = layer_number + + # Dimension config + self.qk_head_dim = config.qk_head_dim + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.v_head_dim = config.v_head_dim + self.num_attention_heads = config.num_attention_heads + + # Build indexer using DSAIndexerV2 + self.indexer = DSAIndexerV2(config, submodules.indexer.submodules) + + if softmax_scale is None: + q_head_dim = config.qk_head_dim + config.qk_pos_emb_head_dim + softmax_scale = 1.0 / math.sqrt(q_head_dim) + self.softmax_scale = softmax_scale + + self.loss_scaler = DSAIndexerLossAutoScaler() + + def construct( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Tensor, + x: Tensor, + qr: Tensor, + rotary_pos_emb: Optional[Tensor] = None, + ): + """ + Forward pass using sparse_flash_attention operator. + + Args: + query: Query tensor [sq, b, np, hn]. + key: Key tensor [skv, b, np, hn]. + value: Value tensor [skv, b, np, hnv]. + attention_mask: Attention mask tensor. + x: Original hidden states [sq, b, hidden_size]. + qr: Low-rank query representation [sq, b, q_lora_rank]. + rotary_pos_emb: Rotary position embedding tensor. + + Returns: + output: Output tensor [sq, b, hidden_size] + """ + _ = attention_mask + + sq, b, np, _ = query.shape + hnv = value.shape[3] + + # Detach x and qr + x_detached = ops.stop_gradient(x) + qr_detached = ops.stop_gradient(qr) + + # Get sparse indices from indexer + sparse_indices = self.indexer( + x_detached, qr_detached, rotary_pos_emb=rotary_pos_emb + ) + + # Separate nope and rope parts + q_nope = query[..., :self.qk_head_dim] + q_pe = query[..., self.qk_head_dim:] + k_nope = key[..., :self.qk_head_dim] + k_pe = key[..., self.qk_head_dim:] + + # Layout conversion: SBND -> BSND + q_nope = mint.permute(q_nope, (1, 0, 2, 3)) + k_nope = mint.permute(k_nope, (1, 0, 2, 3)) + value_bsnd = mint.permute(value, (1, 0, 2, 3)) + q_pe = mint.permute(q_pe, (1, 0, 2, 3)) + k_pe = mint.permute(k_pe, (1, 0, 2, 3)) + + # Call sparse_flash_attention operator + attention_out, _, _ = ops.sparse_flash_attention( + query=q_nope, + key=k_nope, + value=value_bsnd, + sparse_indices=sparse_indices, + scale_value=self.softmax_scale, + query_rope=q_pe, + key_rope=k_pe, + layout_query='BSND', + layout_kv='BSND', + sparse_mode=3, + attention_mode=2, + return_softmax_lse=False + ) + + # Layout conversion: BSND -> SBND + output = mint.permute(attention_out, (1, 0, 2, 3)) + output = mint.reshape(output, (sq, b, np * hnv)) + + return output diff --git a/mindformers/pynative/transformers/multi_latent_attention.py b/mindformers/pynative/transformers/multi_latent_attention.py index 92c74e85348e92e5813e6d9a56b1b1bac8147dc0..94639c69b2caa3c96df6df998b47d5da23917269 100644 --- a/mindformers/pynative/transformers/multi_latent_attention.py +++ b/mindformers/pynative/transformers/multi_latent_attention.py @@ -129,14 +129,14 @@ class MultiLatentAttention(nn.Cell): def construct(self, x: Tensor, attention_mask=None, rotary_pos_emb=None, prefix_keys_values=None, pad_zeros=None, actual_seq_len=None): """ Forward pass of the Multi-head Latent Attention mechanism. - + Args: x: Input tensor with shape (seq_length, batch_size, hidden_size). attention_mask: Attention mask tensor (optional). rotary_pos_emb: Rotary position embedding tensor (optional). pad_zeros: Padding zeros tensor (not used). actual_seq_len: Actual sequence length for EOD mask compression (optional). - + Returns: Tensor: Output tensor with shape (seq_length, batch_size, hidden_size). """ @@ -153,7 +153,17 @@ class MultiLatentAttention(nn.Cell): query = self.cast(query, self.compute_dtype) key = self.cast(key, self.compute_dtype) value = self.cast(value, self.compute_dtype) - if self.use_flash_attention: + + # Check if using DSA (DeepSeek Sparse Attention) + use_dsa = getattr(self.config, 'experimental_attention_variant', None) == 'dsa' + + if use_dsa: + # DSA requires original hidden states and compressed query + attn_out = self.core_attention( + query, key, value, attention_mask, + x, self.q_compressed, rotary_pos_emb + ) + elif self.use_flash_attention: if self.use_eod_attn_mask_compression: context_layer = self.core_attention( query, key, value, attention_mask, @@ -307,14 +317,14 @@ class MLASelfAttention(MultiLatentAttention): ): """ Derive query, key, and value tensors from hidden states. - + This method generates query, key, and value tensors from the input hidden states using the configured projection layers and applies rotary position embeddings. - + Args: hidden_states: Input hidden states tensor with shape [seq_length, batch_size, hidden_size]. rotary_pos_emb: Rotary position embedding tensor (optional). - + Returns: tuple: A tuple containing query, key, and value tensors. - query: Query tensor with shape [seq_length, batch_size, num_heads, head_dim]. @@ -335,6 +345,9 @@ class MLASelfAttention(MultiLatentAttention): dim=-1, ) + # Store compressed query for DSA + self.q_compressed = q_a + if self.q_layernorm is not None: q_a = self.q_layernorm(q_a) q = self.linear_qb(q_a)[0]