diff --git a/.gitignore b/.gitignore index 9af6f991a6746dc0eb1c9f08cf5dd1bce45c6e5e..9cb41c969b7a004c8da7907cccf0a0bcf4e8abb8 100644 --- a/.gitignore +++ b/.gitignore @@ -140,3 +140,6 @@ dmypy.json # Cython debug symbols cython_debug/ + +DSA_vs_NSA_COMPARISON.md +NSA_LIMITATIONS_AND_IMPROVEMENTS.md \ No newline at end of file diff --git a/ds_pynative.yaml b/ds_pynative.yaml index 15461272e37878dd30e7265018c60fc5b0a42488..ed512c1ed2e5251ba83eb21ea7e131b8635bf1e7 100644 --- a/ds_pynative.yaml +++ b/ds_pynative.yaml @@ -26,7 +26,7 @@ parallelism: data_parallel: 1 training: - max_steps: 2000 + max_steps: 20000 global_batch_size: 1 local_batch_size: 1 save_steps: 2000 @@ -91,7 +91,7 @@ train_dataset: &train_dataset type: BlendedMegatronDatasetDataLoader datasets_type: "GPTDataset" sizes: - - 100 # Number of training set data samples. + - 20000 # Number of training set data samples. - 0 # Number of test set data samples. Currently, configuration is not supported. - 0 # Number of eval set data samples. Currently, configuration is not supported. config: # GPTDataset Configs @@ -108,7 +108,7 @@ train_dataset: &train_dataset pad: -1 # The token id of `pad` in the dataset. data_path: # Megatron dataset sampling ratio and path. - '1' - - "/home/w00932055/dsv4/deepseek-datasets/mmap_deepseekv3_datasets_text_document" + - "/home/w00932055/deepseek-datasets/fineweb_edu_10BT_text_document" input_columns: ["input_ids", "labels", "loss_mask", "position_ids"] construct_args_key: ["input_ids", "labels", "loss_mask", "position_ids"] num_parallel_workers: 8 @@ -123,7 +123,7 @@ train_dataset_task: # mindspore context init config context: - mode: 0 # 0--Graph Mode; 1--Pynative Mode + mode: 1 # 0--Graph Mode; 1--Pynative Mode device_target: "Ascend" max_call_depth: 10000 max_device_memory: "58GB" @@ -180,7 +180,7 @@ model: num_hidden_layers: 8 max_position_embeddings: 163840 hidden_act: 'silu' # 'fusedswiglu' - num_attention_heads: 4 + num_attention_heads: 8 rms_norm_eps: 1.e-6 add_bias_linear: False use_flash_attention: True @@ -203,6 +203,19 @@ model: # 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) + # NSA (Native Sparse Attention) Configuration + # experimental_attention_variant: 'nsa' + # nsa_local_window: 128 + # nsa_block_size: 32 + # nsa_stride: 32 + # nsa_topk_blocks: 4 + # nsa_compression: "grouped_mlp" # grouped_mlp | conv1d | avgpool + # nsa_gate_mode: "static" # static | q_cond + # nsa_gate_init: [2.0, -2.0, -2.0] + # nsa_dropout: 0.0 + # nsa_indexer_n_heads: 4 + # nsa_indexer_head_dim: 192 # qk_nope_head_dim + qk_rope_head_dim + # nsa_indexer_loss_coeff: 0.001 attention_dropout: 0.0 hidden_dropout: 0.0 params_dtype: "float32" @@ -233,12 +246,12 @@ model: moe_intermediate_size: 2048 routed_scaling_factor: 1.5 first_k_dense_replace: 1 - n_routed_experts: 16 + n_routed_experts: 16 # FFN + COPY num_experts_per_tok: 8 n_shared_experts: 1 num_copy_experts: 0 use_topk_router_with_load_balancing: False - moe_expected_ffn_experts: 2.0 # Default best value: top-k * FFN/(FFN + COPY) + moe_expected_ffn_experts: 8.0 # Default best value: top-k * FFN/(FFN + COPY) moe_router_bias_update_rate: 0.001 moe_shared_expert_intermediate_size: 2048 moe_grouped_gemm: True diff --git a/mindformers/parallel_core/transformer_config.py b/mindformers/parallel_core/transformer_config.py index 08f9b3e035dc16faa9884e14ca2b1ef9dc00a0f4..184726c882eb72f5d1d39f6f396a21d47b253714 100644 --- a/mindformers/parallel_core/transformer_config.py +++ b/mindformers/parallel_core/transformer_config.py @@ -903,11 +903,11 @@ class MLATransformerConfig(TransformerConfig): """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" #################### - # DSA (DeepSeek Sparse Attention) + # DSA / NSA (Sparse Attention Variants) #################### experimental_attention_variant: Optional[str] = None - """Experimental attention variant to use. Options: 'dsa' for DeepSeek Sparse Attention.""" + """Experimental attention variant to use. Options: 'dsa' or 'nsa'.""" dsa_indexer_n_heads: Optional[int] = None """Number of indexer heads for DSA. If None, defaults to num_attention_heads.""" @@ -927,6 +927,40 @@ class MLATransformerConfig(TransformerConfig): dsa_use_fused_ops: bool = False """Use fused DSA operators (lightning_indexer + sparse_flash_attention) for better performance.""" + # NSA (Native Sparse Attention) + nsa_local_window: int = 128 + """Sliding window size for local attention.""" + + nsa_block_size: int = 32 + """Block size for compression and selection.""" + + nsa_stride: int = 32 + """Stride for compression and selection blocks.""" + + nsa_topk_blocks: int = 4 + """Top-k blocks for selective attention.""" + + nsa_compression: str = "grouped_mlp" + """Compression type: grouped_mlp | conv1d | avgpool.""" + + nsa_gate_mode: str = "static" + """Gating mode: static | q_cond.""" + + nsa_gate_init: Tuple[float, float, float] = (2.0, -2.0, -2.0) + """Initial gate logits for (local, compressed, selected).""" + + nsa_dropout: Optional[float] = None + """NSA attention dropout. Defaults to attention_dropout when None.""" + + nsa_indexer_n_heads: Optional[int] = None + """Number of Lightning Indexer heads for NSA block selection. If None, defaults to 4.""" + + nsa_indexer_head_dim: Optional[int] = None + """Dimension per Lightning Indexer head for NSA. If None, defaults to qk_head_dim + qk_pos_emb_head_dim.""" + + nsa_indexer_loss_coeff: float = 0.0 + """Coefficient for KL divergence auxiliary loss in NSA Lightning Indexer. Set to 0 to disable.""" + def __post_init__(self): """Initialize DSA default values if not set.""" super().__post_init__() @@ -938,4 +972,19 @@ class MLATransformerConfig(TransformerConfig): if self.dsa_indexer_head_dim is None: self.dsa_indexer_head_dim = self.qk_head_dim + self.qk_pos_emb_head_dim + if self.experimental_attention_variant == 'nsa': + if self.nsa_dropout is None: + self.nsa_dropout = self.attention_dropout + if self.nsa_block_size % self.nsa_stride != 0: + raise ValueError("nsa_block_size must be divisible by nsa_stride") + if self.nsa_local_window % 2 != 0: + raise ValueError("nsa_local_window must be even") + if self.nsa_topk_blocks <= 0: + raise ValueError("nsa_topk_blocks must be positive") + # Set default NSA indexer parameters + if self.nsa_indexer_n_heads is None: + self.nsa_indexer_n_heads = 4 + if self.nsa_indexer_head_dim is None: + self.nsa_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 70209824293acd58ae608025577309f1b67631a1..80985b90f5e401cc73f2a95c431234231949be1f 100644 --- a/mindformers/parallel_core/transformer_config_utils.py +++ b/mindformers/parallel_core/transformer_config_utils.py @@ -420,6 +420,19 @@ COMMON_CONFIG_MAPPING = { "dsa_indexer_use_sparse_loss": "dsa_indexer_use_sparse_loss", "dsa_use_fused_ops": "dsa_use_fused_ops", + # NSA (Native Sparse Attention) + "nsa_local_window": "nsa_local_window", + "nsa_block_size": "nsa_block_size", + "nsa_stride": "nsa_stride", + "nsa_topk_blocks": "nsa_topk_blocks", + "nsa_compression": "nsa_compression", + "nsa_gate_mode": "nsa_gate_mode", + "nsa_gate_init": "nsa_gate_init", + "nsa_dropout": "nsa_dropout", + "nsa_indexer_n_heads": "nsa_indexer_n_heads", + "nsa_indexer_head_dim": "nsa_indexer_head_dim", + "nsa_indexer_loss_coeff": "nsa_indexer_loss_coeff", + # 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 ca8d24e460de7b141775209c0654d16c790f31fe..b39502ff62a453f0dd4a91aced4003cd0e71f253 100644 --- a/mindformers/pynative/base_models/gpt/gpt_layer_specs.py +++ b/mindformers/pynative/base_models/gpt/gpt_layer_specs.py @@ -36,6 +36,7 @@ from mindformers.pynative.transformers.multi_latent_attention import MLASelfAtte MLASelfAttentionSubmodules from mindformers.pynative.transformers.dsa import DSAttention, DSAttentionSubmodules, \ DSAIndexer, DSAIndexerSubmodules, DSAIndexerV2, DSAttentionV2 +from mindformers.pynative.transformers.nsa import NSAAttention def get_mlp_module_spec( num_experts: Optional[int] = None, @@ -127,6 +128,8 @@ def get_gpt_layer_local_spec( ), ), ) + elif experimental_attention_variant == 'nsa': + core_attention = NSAAttention else: # Use standard FlashAttention core_attention = FlashAttention diff --git a/mindformers/pynative/transformers/multi_latent_attention.py b/mindformers/pynative/transformers/multi_latent_attention.py index 94639c69b2caa3c96df6df998b47d5da23917269..da2cdbd1cc607c5c651e77d6c326945cc99f7abd 100644 --- a/mindformers/pynative/transformers/multi_latent_attention.py +++ b/mindformers/pynative/transformers/multi_latent_attention.py @@ -154,8 +154,9 @@ class MultiLatentAttention(nn.Cell): key = self.cast(key, self.compute_dtype) value = self.cast(value, self.compute_dtype) - # Check if using DSA (DeepSeek Sparse Attention) + # Check if using DSA / NSA use_dsa = getattr(self.config, 'experimental_attention_variant', None) == 'dsa' + use_nsa = getattr(self.config, 'experimental_attention_variant', None) == 'nsa' if use_dsa: # DSA requires original hidden states and compressed query @@ -163,6 +164,13 @@ class MultiLatentAttention(nn.Cell): query, key, value, attention_mask, x, self.q_compressed, rotary_pos_emb ) + elif use_nsa: + # NSA requires original hidden states for compression and + # compressed query (qr) for the Lightning Indexer block selection. + 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( diff --git a/mindformers/pynative/transformers/nsa.py b/mindformers/pynative/transformers/nsa.py new file mode 100644 index 0000000000000000000000000000000000000000..851c1d858ab2f54669f3e26005a02b98acb540cc --- /dev/null +++ b/mindformers/pynative/transformers/nsa.py @@ -0,0 +1,745 @@ +""" +Native Sparse Attention (NSA) for MindSpore Pynative. + +Implements three-branch sparse attention: + 1) local sliding window attention + 2) compressed token attention + 3) selective block attention (with Lightning Indexer for block selection) + +The block selection branch supports two scoring modes: + - Legacy: standard Softmax(Q @ K_compressed^T) scoring + - Lightning Indexer: weighted ReLU attention scoring (from DSA/NSA paper) + Reference: https://arxiv.org/pdf/2502.11089 +""" +import math +from dataclasses import dataclass +from typing import Optional + +import mindspore as ms +from mindspore import nn, Tensor, Parameter, mint, ops + +from mindformers.parallel_core.transformer_config import MLATransformerConfig +from mindformers.pynative.layers.linear import Linear +from mindformers.pynative.layers.layer_norm import get_norm_cls +from mindformers.pynative.layers.dropout import Dropout + + +@dataclass +class NSACompressionConfig: + block_size: int + compression: str + normalization: str + fused_norm: bool + params_dtype: str + compute_dtype: str + init_method: Optional[callable] + norm_eps: float + + +def _pad_to_multiple(x: Tensor, block: int) -> Tensor: + """Pad sequence length to a multiple of block size by repeating the last token.""" + seq_len = x.shape[1] + pad = (block - seq_len % block) % block + if pad == 0: + return x + last = x[:, -1:, :] + pad_tensor = mint.tile(last, (1, pad, 1)) + return mint.cat((x, pad_tensor), dim=1) + + +class GroupedMLPCompression(nn.Cell): + """Project each block of tokens into a single vector using a grouped MLP.""" + + def __init__(self, dim: int, cfg: NSACompressionConfig): + super().__init__() + self.block = cfg.block_size + self.linear = Linear( + input_size=dim * cfg.block_size, + output_size=dim, + params_dtype=cfg.params_dtype, + compute_dtype=cfg.compute_dtype, + init_method=cfg.init_method, + bias=False, + skip_bias_add=False, + ) + norm_cls = get_norm_cls(cfg.normalization, cfg.fused_norm) + self.norm = norm_cls(dim=dim, eps=cfg.norm_eps) + # Remove bias as it can cause training instability + # The linear layer and norm should be sufficient for learning + self.reshape = mint.reshape + + def construct(self, x: Tensor) -> Tensor: + # x: (b, n, d) + x = _pad_to_multiple(x, self.block) + b, n, d = x.shape + x = self.reshape(x, (b, n // self.block, self.block, d)) + # Removed bias addition for better stability + x = self.reshape(x, (b, n // self.block, self.block * d)) + x = self.linear(x)[0] + return self.norm(x) + + +class Conv1dCompression(nn.Cell): + """Depthwise 1D convolution compression.""" + + def __init__(self, dim: int, cfg: NSACompressionConfig): + super().__init__() + self.block = cfg.block_size + self.conv = nn.Conv1d( + in_channels=dim, + out_channels=dim, + kernel_size=cfg.block_size, + stride=cfg.block_size, + group=dim, + has_bias=False, + pad_mode="valid", + ) + norm_cls = get_norm_cls(cfg.normalization, cfg.fused_norm) + self.norm = norm_cls(dim=dim, eps=cfg.norm_eps) + self.transpose = mint.transpose + + def construct(self, x: Tensor) -> Tensor: + # x: (b, n, d) + x = _pad_to_multiple(x, self.block) + x = self.transpose(x, (0, 2, 1)) # (b, d, n) + x = self.conv(x) + x = self.transpose(x, (0, 2, 1)) # (b, n', d) + return self.norm(x) + + +class AvgPoolCompression(nn.Cell): + """Average pooling compression.""" + + def __init__(self, dim: int, cfg: NSACompressionConfig): + super().__init__() + self.block = cfg.block_size + self.pool = nn.AvgPool1d(kernel_size=cfg.block_size, stride=cfg.block_size) + norm_cls = get_norm_cls(cfg.normalization, cfg.fused_norm) + self.norm = norm_cls(dim=dim, eps=cfg.norm_eps) + + def construct(self, x: Tensor) -> Tensor: + # x: (b, n, d) + x = _pad_to_multiple(x, self.block) + x = mint.transpose(x, (0, 2, 1)) # (b, d, n) + x = self.pool(x) + x = mint.transpose(x, (0, 2, 1)) # (b, n', d) + return self.norm(x) + + +class TokenCompressor(nn.Cell): + """Compression wrapper to select grouped MLP, conv1d, or avgpool.""" + + def __init__(self, dim: int, cfg: NSACompressionConfig): + super().__init__() + if cfg.compression == "grouped_mlp": + self.op = GroupedMLPCompression(dim, cfg) + elif cfg.compression == "conv1d": + self.op = Conv1dCompression(dim, cfg) + elif cfg.compression == "avgpool": + self.op = AvgPoolCompression(dim, cfg) + else: + raise ValueError(f"Unsupported NSA compression: {cfg.compression}") + + def construct(self, x: Tensor) -> Tensor: + return self.op(x) + + +def _rotate_activation(x: Tensor) -> Tensor: + """Apply Hadamard rotation activation (simplified). + + Scales the input by 1/sqrt(hidden_dim) as an approximation of the + Hadamard transform used in DSA's Lightning Indexer. + + Reference: + https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/inference/model.py#L424-L428 + + Args: + x: Input tensor. + + Returns: + Scaled tensor. + """ + hidden_size = x.shape[-1] + scale = hidden_size ** -0.5 + return x * scale + + +class NSABlockIndexer(nn.Cell): + """Lightning Indexer for NSA block-level top-k selection. + + Replaces the standard Softmax(Q @ K_compressed^T) block scoring with a + weighted ReLU attention mechanism, inspired by DSA's Lightning Indexer. + + The scoring process: + 1. Project low-rank query (qr) into multi-head indexer queries. + 2. Project compressed block representations into indexer keys. + 3. Compute ``ReLU(q_idx @ k_idx^T)`` instead of ``Softmax(Q @ K^T)``. + 4. Weight each head via a learned projection from original hidden states. + 5. Sum across indexer heads to produce per-block importance scores. + + This naturally enforces that all query heads within a GQA group select the + same key/value blocks, since the scores are head-agnostic. + + Reference: + NSA paper Section 2.1: https://arxiv.org/pdf/2502.11089 + """ + + def __init__(self, config: MLATransformerConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.q_lora_rank = ( + config.q_lora_rank + if config.q_lora_rank is not None + else config.hidden_size + ) + self.index_n_heads = config.nsa_indexer_n_heads + self.index_head_dim = config.nsa_indexer_head_dim + self.softmax_scale: float = self.index_head_dim ** -0.5 + + # Query projection: q_lora_rank -> index_n_heads * index_head_dim + self.linear_wq_b = Linear( + 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=config.init_method, + bias=False, + skip_bias_add=False, + ) + + # Key projection: hidden_size -> index_head_dim (from compressed blocks) + self.linear_wk = Linear( + input_size=self.hidden_size, + output_size=self.index_head_dim, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + ) + + # Key normalization + norm_cls = get_norm_cls(config.normalization, config.fused_norm) + self.k_norm = norm_cls(dim=self.index_head_dim, eps=config.layernorm_epsilon) + + # Weight projection: hidden_size -> index_n_heads (per-query position) + self.linear_weights_proj = Linear( + input_size=self.hidden_size, + output_size=self.index_n_heads, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + ) + + def construct( + self, + qr: Tensor, + xc_blocks: Tensor, + x: Tensor, + ) -> Tensor: + """Compute block importance scores using weighted ReLU attention. + + Args: + qr: Low-rank query tensor [seqlen, batch, q_lora_rank]. + xc_blocks: Compressed block representations [batch, n_blocks, hidden_size]. + x: Original hidden states [seqlen, batch, hidden_size] (for per-query weights). + + Returns: + block_scores: Block importance scores [batch, seqlen, n_blocks]. + """ + sq, b, _ = qr.shape + n_blk = xc_blocks.shape[1] + + # ---- Query path ---- + # [sq, b, q_lora_rank] -> [sq, b, n_heads * head_dim] + q = self.linear_wq_b(qr)[0] + q = mint.reshape(q, (sq, b, self.index_n_heads, self.index_head_dim)) + + # ---- Key path ---- + # [b, n_blk, hidden] -> [n_blk, b, hidden] + xc_t = mint.permute(xc_blocks, (1, 0, 2)) + # [n_blk, b, hidden] -> [n_blk, b, head_dim] + k = self.linear_wk(xc_t)[0] + k = self.k_norm(k) + + # Hadamard rotation activation (approximation) + q = _rotate_activation(q) + k = _rotate_activation(k) + + # ---- Weights path ---- + # [sq, b, hidden] -> [sq, b, n_heads] + weights = self.linear_weights_proj(x)[0] + weights = weights * (self.index_n_heads ** -0.5) * self.softmax_scale + + # ---- Weighted ReLU attention ---- + # q: [sq, b, n_heads, head_dim], k: [n_blk, b, head_dim] + # -> scores: [sq, b, n_heads, n_blk] + scores = mint.einsum( + 'sbhd,tbd->sbht', + ops.cast(q, ms.float32), + ops.cast(k, ms.float32), + ) + # ReLU instead of Softmax + scores = ops.relu(scores) + + # Weight each head: [sq, b, n_heads, n_blk] * [sq, b, n_heads, 1] + scores = scores * mint.unsqueeze(ops.cast(weights, ms.float32), -1) + + # Sum across indexer heads: [sq, b, n_blk] + scores = scores.sum(axis=2) + + # Transpose to [b, sq, n_blk] + block_scores = mint.transpose(scores, 0, 1) + + return block_scores + + +def _bool_to_score_mask(mask: Tensor, dtype: ms.dtype) -> Tensor: + """Convert boolean mask to additive score mask.""" + # Use a large negative value to avoid NaNs when all positions are masked. + neg_inf = mint.full(mask.shape, float(-1e9), dtype=dtype) + zero = mint.zeros_like(neg_inf) + return mint.where(mask, zero, neg_inf) + + +def _normalize_attention_mask(attention_mask: Tensor, dtype: ms.dtype) -> Tensor: + """Normalize attention_mask to additive mask with shape (b, 1, s, s).""" + if attention_mask.dtype in (ms.bool_, ms.uint8): + attn = _bool_to_score_mask(attention_mask.astype(ms.bool_), dtype) + else: + attn = attention_mask.astype(dtype) + + if attn.ndim == 2: + attn = attn.reshape((attn.shape[0], 1, 1, attn.shape[1])) + elif attn.ndim == 3: + attn = attn.reshape((attn.shape[0], 1, attn.shape[1], attn.shape[2])) + return attn + + +class NSAAttention(nn.Cell): + """Native Sparse Attention core attention for MLA.""" + + def __init__( + self, + config: MLATransformerConfig, + layer_number: int, + softmax_scale: Optional[float] = None, + ): + super().__init__() + self.config = config + self.layer_number = layer_number + + self.num_heads = config.num_attention_heads + self.q_head_dim = config.qk_head_dim + config.qk_pos_emb_head_dim + self.v_head_dim = config.v_head_dim + self.hidden_size = config.hidden_size + + self.local_window = config.nsa_local_window + self.block_size = config.nsa_block_size + self.stride = config.nsa_stride + self.topk_blocks = config.nsa_topk_blocks + self.compression = config.nsa_compression + self.gate_mode = config.nsa_gate_mode + self.gate_init = config.nsa_gate_init + self.nsa_dropout = config.nsa_dropout if config.nsa_dropout is not None else config.attention_dropout + + if self.block_size % self.stride != 0: + raise ValueError("nsa_block_size must be divisible by nsa_stride") + if self.local_window % 2 != 0: + raise ValueError("nsa_local_window must be even") + if self.topk_blocks <= 0: + raise ValueError("nsa_topk_blocks must be positive") + + self.softmax_scale = softmax_scale if softmax_scale is not None else (self.q_head_dim ** -0.5) + + # Lightning Indexer for block selection (replaces softmax-based scoring) + self.nsa_indexer_loss_coeff = getattr(config, 'nsa_indexer_loss_coeff', 0.0) + nsa_indexer_n_heads = getattr(config, 'nsa_indexer_n_heads', None) + if nsa_indexer_n_heads is not None and nsa_indexer_n_heads > 0: + self.block_indexer = NSABlockIndexer(config) + else: + self.block_indexer = None + + comp_cfg = NSACompressionConfig( + block_size=self.block_size, + compression=self.compression, + normalization=config.normalization, + fused_norm=config.fused_norm, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=config.init_method, + norm_eps=config.layernorm_epsilon, + ) + self.compressor = TokenCompressor(self.hidden_size, comp_cfg) + self.kvc_proj = Linear( + input_size=self.hidden_size, + output_size=self.num_heads * (self.q_head_dim + self.v_head_dim), + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + ) + + # Normalize compressed keys to prevent attention score explosion + # (analogous to k_layernorm in the standard MLA path) + norm_cls = get_norm_cls(config.normalization, config.fused_norm) + self.kc_norm = norm_cls(dim=self.q_head_dim, eps=config.layernorm_epsilon) + + if self.gate_mode == "static": + init = Tensor(self.gate_init, ms.float32) + self.gate = Parameter(mint.tile(init, (self.num_heads, 1)), name="nsa_gate") + self.gate_proj = None + elif self.gate_mode == "q_cond": + self.gate = Parameter(mint.zeros((self.num_heads, 3), dtype=ms.float32), name="nsa_gate") + self.gate_proj = Linear( + input_size=self.q_head_dim, + output_size=3, + params_dtype=config.params_dtype, + compute_dtype=config.compute_dtype, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + ) + else: + raise ValueError(f"Unsupported nsa_gate_mode: {self.gate_mode}") + + self.attn_dropout = Dropout(drop_prob=float(self.nsa_dropout)) + self.softmax = mint.nn.functional.softmax + + def construct( + self, + query: Tensor, + key: Tensor, + value: Tensor, + attention_mask: Optional[Tensor], + x: Tensor, + qr: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + ): + _ = rotary_pos_emb + + # SBHD -> B H S D + q = mint.permute(query, (1, 2, 0, 3)) + k = mint.permute(key, (1, 2, 0, 3)) + v = mint.permute(value, (1, 2, 0, 3)) + b, h, n, _ = q.shape + + # Compressed tokens from hidden states + # Detach x to prevent gradient shortcut through the compression path + # back to earlier transformer layers. The compressor parameters still + # receive gradients from the attention loss via kc/vc. + # (analogous to ops.stop_gradient(x) in DSA's indexer) + x_b = mint.permute(x, (1, 0, 2)) # (b, n, hidden) + x_b = ops.stop_gradient(x_b) + xc = self.compressor(x_b) + kvc = self.kvc_proj(xc)[0] + kvc = mint.reshape(kvc, (b, xc.shape[1], h, self.q_head_dim + self.v_head_dim)) + kc, vc = mint.split(kvc, [self.q_head_dim, self.v_head_dim], dim=-1) + kc = mint.permute(kc, (0, 2, 1, 3)) # (b, h, n_c, q_head_dim) + vc = mint.permute(vc, (0, 2, 1, 3)) # (b, h, n_c, v_head_dim) + + # Normalize compressed keys to constrain norms and prevent attention + # score explosion as kvc_proj weights evolve during training + kc = self.kc_norm(kc) + + # --- Block selection scoring --- + # Use Lightning Indexer (weighted ReLU) if available, else fall back + # to legacy Softmax(Q @ K_compressed^T) scoring. + block_scores = None + if self.block_indexer is not None and qr is not None: + n_c = xc.shape[1] + tokens_per_block = self.block_size // self.stride + blk_total = math.ceil(n_c / tokens_per_block) + + # Aggregate compressed tokens into block representations + if tokens_per_block == 1: + xc_blocks = xc # (b, blk_total, hidden) + else: + pad = blk_total * tokens_per_block - n_c + if pad > 0: + pad_tensor = mint.zeros((b, pad, self.hidden_size), dtype=xc.dtype) + xc_padded = mint.cat((xc, pad_tensor), dim=1) + else: + xc_padded = xc + xc_padded = mint.reshape(xc_padded, (b, blk_total, tokens_per_block, self.hidden_size)) + xc_blocks = xc_padded.mean(dim=2) # (b, blk_total, hidden) + + # Detach inputs to indexer — the indexer learns via auxiliary loss, + # not through the main attention gradient path. + qr_det = ops.stop_gradient(qr) + x_det = ops.stop_gradient(x) + xc_blocks_det = ops.stop_gradient(xc_blocks) + + block_scores = self.block_indexer(qr_det, xc_blocks_det, x_det) + + # Branch masks + local_mask = self._build_local_mask(n, self.local_window, q.dtype) + comp_mask = self._build_compressed_mask(n, kc.shape[2], self.stride, q.dtype) + sel_mask, sel_valid = self._build_selected_mask(q, kc, q.dtype, block_scores=block_scores) + + if attention_mask is not None: + attn_mask = _normalize_attention_mask(attention_mask, q.dtype) + if attn_mask.shape[-1] == n: + local_mask = local_mask + attn_mask + sel_mask = sel_mask + attn_mask + # Check if any position in each row is valid (not masked) + attn_row_valid = (mint.sum((attn_mask > -1e8).astype(ms.int32), dim=-1, keepdim=True) > 0).astype(ms.float32) + sel_valid = sel_valid * attn_row_valid + if attn_mask.shape[-2] == n: + # For compressed branch, we need to handle the sequence length mismatch + # Take the max mask value per block of tokens that gets compressed + n_c = kc.shape[2] + if n_c > 0: + # Simple approach: check if query position can attend based on causal constraint + # Query at position i can attend to compressed token j if i >= j * stride + q_pos = mint.arange(n, dtype=ms.int32).reshape((n, 1)) + c_pos = mint.arange(n_c, dtype=ms.int32).reshape((1, n_c)) + comp_causal = (q_pos >= c_pos * self.stride).reshape((1, 1, n, n_c)) + comp_causal_mask = _bool_to_score_mask(comp_causal, q.dtype) + comp_mask = comp_mask + comp_causal_mask + + # Branch outputs + local_out = self._attend(q, k, v, local_mask) + comp_out = self._attend(q, kc, vc, comp_mask) + sel_out = self._attend(q, k, v, sel_mask) + + # Instead of multiplying by sel_valid (which can kill gradients), + # we incorporate validity into the gating weights + # This provides more stable gradients during training + sel_valid_mask = sel_valid # (1, 1, n, 1) + + # Gating + if self.gate_mode == "static": + g = self.softmax(self.gate, dim=-1) # (h, 3) + w = g.reshape((1, h, 1, 3)) + else: + q_mean = q.mean(dim=-2) # (b, h, d) + q_mean_flat = mint.reshape(q_mean, (b * h, self.q_head_dim)) + gate_proj = self.gate_proj(q_mean_flat)[0] + gate_proj = mint.reshape(gate_proj, (b, h, 3)) + g = gate_proj + self.gate.reshape((1, h, 3)) + w = mint.unsqueeze(self.softmax(g, dim=-1), -2) # (b, h, 1, 3) + + # Apply validity mask to selective branch gate weight instead of output + # This prevents gradient flow issues when sel_valid is 0 + # Redistribute the masked weight to other branches for stability + w_sel = w[..., 2, None] * sel_valid_mask + w_redistrib = w[..., 2, None] * (1.0 - sel_valid_mask) + w_local = w[..., 0, None] + w_redistrib * 0.5 + w_comp = w[..., 1, None] + w_redistrib * 0.5 + + out = w_local * local_out + w_comp * comp_out + w_sel * sel_out + + # --- Auxiliary indexer loss --- + # Train the Lightning Indexer to predict block importance by matching + # the compressed attention scores (analogous to DSA's KL divergence loss). + if self.training and self.block_indexer is not None and block_scores is not None: + if self.nsa_indexer_loss_coeff > 0: + indexer_loss = self._compute_indexer_loss(q, kc, block_scores) + # Attach loss without affecting forward values (same as DSA pattern) + out = out + indexer_loss * 0.0 + + out = mint.permute(out, (2, 0, 1, 3)) + out = mint.reshape(out, (n, b, h * self.v_head_dim)) + return out + + def _attend(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor]) -> Tensor: + scores = mint.einsum("bhid,bhjd->bhij", q.astype(ms.float32), k.astype(ms.float32)) + scores = scores * self.softmax_scale + if mask is not None: + scores = scores + mask + # Clip scores to prevent overflow/underflow in softmax + # This is critical for training stability + scores = mint.clamp(scores, min=-50.0, max=50.0) + attn = self.softmax(scores, dim=-1) + # Replace NaN values with uniform distribution if they occur + # This prevents cascading NaN issues during training + attn = mint.where(mint.isnan(attn), mint.full_like(attn, 1.0 / attn.shape[-1]), attn) + attn = self.attn_dropout(attn) + out = mint.einsum("bhij,bhjd->bhid", attn.astype(v.dtype), v) + return out + + def _build_local_mask(self, seq_len: int, window: int, dtype: ms.dtype) -> Tensor: + idx = mint.arange(seq_len, dtype=ms.int32) + diff = idx.reshape((seq_len, 1)) - idx.reshape((1, seq_len)) + mask = (diff >= 0) & (diff < window) + mask = mask.reshape((1, 1, seq_len, seq_len)) + return _bool_to_score_mask(mask, dtype) + + def _build_causal_mask(self, seq_len: int) -> Tensor: + idx = mint.arange(seq_len, dtype=ms.int32) + mask = idx.reshape((seq_len, 1)) >= idx.reshape((1, seq_len)) + return mask.reshape((1, 1, seq_len, seq_len)) + + def _build_compressed_mask(self, seq_len: int, comp_len: int, stride: int, dtype: ms.dtype) -> Tensor: + t = mint.arange(seq_len, dtype=ms.int32).reshape((seq_len, 1)) + c = mint.arange(comp_len, dtype=ms.int32).reshape((1, comp_len)) + mask = (t // stride) > c + mask = mask.reshape((1, 1, seq_len, comp_len)) + return _bool_to_score_mask(mask, dtype) + + def _build_selected_mask( + self, + q: Tensor, + kc: Tensor, + dtype: ms.dtype, + block_scores: Optional[Tensor] = None, + ) -> tuple[Tensor, Tensor]: + """Build the selection mask for the selective attention branch. + + Supports two scoring modes: + - ``block_scores is None``: Legacy mode — uses Softmax(Q @ K_compressed_mean^T) + to score blocks, per attention head. + - ``block_scores is not None``: Lightning Indexer mode — uses pre-computed + weighted-ReLU block scores (shared across all heads). + + Args: + q: Query tensor (b, h, n, d). + kc: Compressed keys (b, h, n_c, d). + dtype: Output mask dtype. + block_scores: Optional block importance scores from Lightning Indexer, + shape (b, n, blk_total). When provided, all heads share the same + block selection (GQA-compatible). + + Returns: + sel_mask: Additive score mask (b, h, n, n). + sel_valid: Per-position validity float (1, 1, n, 1). + """ + b, h, n, d = q.shape + n_c = kc.shape[2] + tokens_per_block = self.block_size // self.stride + blk_total = math.ceil(n_c / tokens_per_block) + + # --- Common causal + diagonal validity mask --- + tok_blk = mint.arange(n, dtype=ms.int32) // self.block_size + blk_id = mint.arange(blk_total, dtype=ms.int32) + causal = tok_blk.reshape((n, 1)) >= blk_id.reshape((1, blk_total)) + diag = tok_blk.reshape((n, 1)) != blk_id.reshape((1, blk_total)) + valid = (causal & diag).reshape((1, 1, n, blk_total)) + valid_any = (mint.sum(valid.astype(ms.int32), dim=-1, keepdim=True) > 0) + + if block_scores is not None: + # ---- Lightning Indexer path ---- + # block_scores: (b, n, blk_total) — shared across all heads. + valid_2d = valid.reshape((1, n, blk_total)) + masked_scores = block_scores + _bool_to_score_mask(valid_2d, block_scores.dtype) + + topk_k = min(self.topk_blocks, blk_total) + topk_values, topk_indices = mint.topk(masked_scores, topk_k, dim=-1) + # Valid if score > large-negative (i.e., not masked out) + updates = (topk_values > -1e8).astype(ms.float32) + + blk_mask_2d = mint.zeros((b, n, blk_total), dtype=ms.float32) + blk_mask_2d = ops.tensor_scatter_elements( + blk_mask_2d, topk_indices, updates, axis=-1 + ) + blk_mask_2d = blk_mask_2d > 0 + valid_any_2d = valid_any.reshape((1, 1, n, 1)) + + # Expand to all heads: (b, n, blk) -> (b, 1, n, blk) -> (b, h, n, blk) + blk_mask = mint.unsqueeze(blk_mask_2d, 1) + blk_mask = mint.tile(blk_mask, (1, h, 1, 1)) + blk_mask = blk_mask & valid_any_2d + else: + # ---- Legacy softmax attention path ---- + q_det = ops.stop_gradient(q) + kc_det = ops.stop_gradient(kc) + + pad = blk_total * tokens_per_block - n_c + if pad > 0: + pad_tensor = mint.zeros((b, h, pad, d), dtype=kc_det.dtype) + kc_det = mint.cat((kc_det, pad_tensor), dim=2) + + kc_det = mint.reshape(kc_det, (b, h, blk_total, tokens_per_block, d)) + kc_mean = kc_det.mean(dim=3) # (b, h, blk, d) + logits = mint.einsum("bhid,bhjd->bhij", q_det, kc_mean) * self.softmax_scale + logits = logits + _bool_to_score_mask(valid, logits.dtype) + + attn = self.softmax(logits, dim=-1) + topk_k = min(self.topk_blocks, blk_total) + topk_values, topk_indices = mint.topk(attn, topk_k, dim=-1) + updates = (topk_values > 1e-3).astype(ms.float32) + + blk_mask = mint.zeros((b, h, n, blk_total), dtype=ms.float32) + blk_mask = ops.tensor_scatter_elements(blk_mask, topk_indices, updates, axis=-1) + blk_mask = blk_mask > 0 + blk_mask = blk_mask & valid_any + + # --- Expand block mask to token-level mask --- + blk_mask = mint.unsqueeze(blk_mask, -1) + blk_mask = mint.tile(blk_mask, (1, 1, 1, 1, self.block_size)) + blk_mask = mint.reshape(blk_mask, (b, h, n, blk_total * self.block_size)) + blk_mask = blk_mask[..., :n] + blk_mask = blk_mask & self._build_causal_mask(n) + sel_valid = valid_any.astype(ms.float32).reshape((1, 1, n, 1)) + return _bool_to_score_mask(blk_mask, dtype), sel_valid + + def _compute_indexer_loss( + self, + q: Tensor, + kc: Tensor, + block_scores: Tensor, + ) -> Tensor: + """Compute KL divergence auxiliary loss to train the Lightning Indexer. + + The target distribution comes from the compressed attention scores + (Softmax over Q @ K_compressed_mean^T, aggregated across heads). + The indexer distribution comes from the Softmax of the Lightning Indexer's + block scores. The loss minimises KL(target || indexer). + + Args: + q: Query tensor (b, h, n, d). + kc: Compressed keys (b, h, n_c, d). + block_scores: Indexer block scores (b, n, blk_total). + + Returns: + Scalar loss tensor. + """ + q_det = ops.stop_gradient(q) + kc_det = ops.stop_gradient(kc) + + b, h, n, d = q_det.shape + n_c = kc_det.shape[2] + tokens_per_block = self.block_size // self.stride + blk_total = math.ceil(n_c / tokens_per_block) + + # Aggregate compressed keys into block means + pad = blk_total * tokens_per_block - n_c + if pad > 0: + pad_tensor = mint.zeros((b, h, pad, d), dtype=kc_det.dtype) + kc_det = mint.cat((kc_det, pad_tensor), dim=2) + kc_det = mint.reshape(kc_det, (b, h, blk_total, tokens_per_block, d)) + kc_mean = kc_det.mean(dim=3) + + # True attention scores: (b, h, n, blk_total) + true_logits = mint.einsum( + "bhid,bhjd->bhij", + ops.cast(q_det, ms.float32), + ops.cast(kc_mean, ms.float32), + ) * self.softmax_scale + + # Causal mask + tok_blk = mint.arange(n, dtype=ms.int32) // self.block_size + blk_id = mint.arange(blk_total, dtype=ms.int32) + causal = tok_blk.reshape((n, 1)) >= blk_id.reshape((1, blk_total)) + diag = tok_blk.reshape((n, 1)) != blk_id.reshape((1, blk_total)) + valid = (causal & diag).reshape((1, 1, n, blk_total)) + true_logits = true_logits + _bool_to_score_mask(valid, true_logits.dtype) + + # Target distribution: softmax, sum across heads, L1-normalise + target = self.softmax(true_logits, dim=-1) # (b, h, n, blk) + target = target.sum(axis=1) # (b, n, blk) + target = target / (target.sum(axis=-1, keepdims=True) + 1e-10) + + # Indexer distribution: softmax of block_scores (with same mask) + valid_2d = valid.reshape((1, n, blk_total)) + idx_logits = block_scores + _bool_to_score_mask(valid_2d, block_scores.dtype) + idx_dist = self.softmax(ops.cast(idx_logits, ms.float32), dim=-1) + + # KL(target || idx_dist) + kl = target * (mint.log(target + 1e-10) - mint.log(idx_dist + 1e-10)) + loss = kl.sum(axis=-1).mean() * self.nsa_indexer_loss_coeff + return loss