From 73cf287ee41f943520dc8edf09767070971c3af7 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 20:29:04 +0800 Subject: [PATCH 01/16] init Signed-off-by: JaredforReal --- .../glm_moe_dsa/configuration_glm_moe_dsa.py | 122 ++-- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 511 ++++++++++------ .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 575 +++++++++++++----- 3 files changed, 813 insertions(+), 395 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 50326dc5ebc6..4e9c9032aa93 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -35,52 +35,46 @@ class GlmMoeDsaConfig(PreTrainedConfig): Args: vocab_size (`int`, *optional*, defaults to 154880): - Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`Glm4MoeLiteModel`] + Vocabulary size of the model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`GlmMoeDsaModel`]. hidden_size (`int`, *optional*, defaults to 6144): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 12288): - Dimension of the MLP representations. + Dimension of the dense MLP representations. moe_intermediate_size (`int`, *optional*, defaults to 2048): - Dimension of the MoE representations. + Dimension of the MoE expert representations. num_hidden_layers (`int`, *optional*, defaults to 78): Number of hidden layers in the Transformer decoder. num_attention_heads (`int`, *optional*, defaults to 64): Number of attention heads for each attention layer in the Transformer decoder. num_key_value_heads (`int`, *optional*, defaults to 64): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details, check out [this - paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to - `num_attention_heads`. + Number of key-value heads for Grouped Query Attention. If equal to `num_attention_heads`, uses MHA. n_shared_experts (`int`, *optional*, defaults to 1): - Number of shared experts. + Number of shared experts in MoE layers. n_routed_experts (`int`, *optional*, defaults to 256): - Number of routed experts. + Number of routed experts in MoE layers. routed_scaling_factor (`float`, *optional*, defaults to 2.5): - Scaling factor or routed experts. + Scaling factor for routed experts. kv_lora_rank (`int`, *optional*, defaults to 512): - Rank of the LoRA matrices for key and value projections. + Rank of the LoRA matrices for key and value projections (MLA). q_lora_rank (`int`, *optional*, defaults to 2048): - Rank of the LoRA matrices for query projections. + Rank of the LoRA matrices for query projections (MLA). qk_rope_head_dim (`int`, *optional*, defaults to 64): Dimension of the query/key heads that use rotary position embeddings. - v_head_dim (`int`, *optional*, defaults to 256): - Dimension of the value heads. qk_nope_head_dim (`int`, *optional*, defaults to 192): Dimension of the query/key heads that don't use rotary position embeddings. + v_head_dim (`int`, *optional*, defaults to 256): + Dimension of the value heads. n_group (`int`, *optional*, defaults to 1): Number of groups for routed experts. topk_group (`int`, *optional*, defaults to 1): - Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + Number of selected groups for each token. num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of selected experts, None means dense model. + Number of experts selected per token. norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the weights of the routed experts. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. + The non-linear activation function in the decoder. max_position_embeddings (`int`, *optional*, defaults to 202752): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): @@ -88,8 +82,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): rms_norm_eps (`float`, *optional*, defaults to 1e-05): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. + Whether or not the model should return the last key/values attentions. pad_token_id (`int`, *optional*): Padding token id. bos_token_id (`int`, *optional*, defaults to 0): @@ -97,28 +90,33 @@ class GlmMoeDsaConfig(PreTrainedConfig): eos_token_id (`int`, *optional*, defaults to 1): End of stream token id. tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings + Whether to tie weight embeddings. rope_parameters (`RopeParameters`, *optional*): - Dictionary containing the configuration parameters for the RoPE embeddings. The dictionary should contain - a value for `rope_theta` and optionally parameters used for scaling in case you want to use RoPE - with longer `max_position_embeddings`. + Configuration parameters for the RoPE embeddings, including `rope_theta` and optional scaling parameters. rope_interleave (`bool`, *optional*, defaults to `True`): Whether to interleave the rotary position embeddings. mlp_layer_types (`list`, *optional*): - MLP (Moe vs Dense) pattern for each layer. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse. + attention_bias (`bool`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. index_topk (`int`, *optional*, defaults to 2048): - Number of top tokens selected by the indexer for retrieval/attention in each step. + Number of top tokens selected by the indexer for sparse attention. + index_head_dim (`int`, *optional*, defaults to 128): + Head dimension for the indexer projections (DSA). + indexer_rope_interleave (`bool`, *optional*, defaults to `True`): + Whether the indexer uses interleaved rotary position embeddings. ```python - >>> from transformers import Glm4MoeLiteModel, Glm4MoeLiteConfig + >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel - >>> # Initializing a GLM-MOE-DSA style configuration + >>> # Initializing a GLM-MoE-DSA configuration >>> configuration = GlmMoeDsaConfig() + >>> # Initializing a model from the configuration + >>> model = GlmMoeDsaModel(configuration) + >>> # Accessing the model configuration >>> configuration = model.config ```""" @@ -158,8 +156,8 @@ def __init__( kv_lora_rank: int | None = 512, q_lora_rank: int | None = 2048, qk_rope_head_dim: int | None = 64, - v_head_dim: int | None = 256, qk_nope_head_dim: int | None = 192, + v_head_dim: int | None = 256, n_group: int | None = 1, topk_group: int | None = 1, num_experts_per_tok: int | None = 8, @@ -167,7 +165,7 @@ def __init__( hidden_act: str | None = "silu", max_position_embeddings: int | None = 202752, initializer_range: float | None = 0.02, - rms_norm_eps: int | None = 1e-5, + rms_norm_eps: float | None = 1e-5, use_cache: bool | None = True, pad_token_id: int | None = None, bos_token_id: int | None = 0, @@ -179,52 +177,68 @@ def __init__( attention_bias: bool | None = False, attention_dropout: float | None = 0.0, index_topk: int | None = 2048, + index_head_dim: int | None = 128, + index_n_heads: int | None = 32, + indexer_rope_interleave: bool | None = True, **kwargs, ): + # Model dimensions self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers + self.max_position_embeddings = max_position_embeddings - # Default to MoE from the fourth layer and on - if mlp_layer_types is None: - mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (self.num_hidden_layers - 3) - layer_type_validation(mlp_layer_types, self.num_hidden_layers, attention=False) - self.mlp_layer_types = mlp_layer_types - - self.moe_intermediate_size = moe_intermediate_size + # Attention dimensions (MLA) self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.routed_scaling_factor = routed_scaling_factor + self.num_key_value_heads = num_key_value_heads self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim self.head_dim = qk_rope_head_dim + + # MoE parameters + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok self.norm_topk_prob = norm_topk_prob - self.rope_interleave = rope_interleave - self.num_key_value_heads = num_key_value_heads + + # MLP layer types: first 3 dense, rest sparse + self.mlp_layer_types = mlp_layer_types + if self.mlp_layer_types is None: + self.mlp_layer_types = ["dense"] * min(3, num_hidden_layers) + ["sparse"] * (num_hidden_layers - 3) + layer_type_validation(self.mlp_layer_types, self.num_hidden_layers, attention=False) + + # Indexer (DSA) parameters + self.index_topk = index_topk + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.indexer_rope_interleave = indexer_rope_interleave + + # General config self.hidden_act = hidden_act self.initializer_range = initializer_range - self.index_topk = index_topk self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.rope_parameters = rope_parameters - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings + self.rope_interleave = rope_interleave - super().__init__(**kwargs) + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) __all__ = ["GlmMoeDsaConfig"] diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 509c3ff4adb9..926ce0ccb295 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -31,24 +31,19 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin -from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub +from ...integrations import use_experts_implementation, use_kernel_forward_from_hub from ...masking_utils import create_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack -from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_grouped_mm_available, logging -from ...utils.generic import maybe_autocast, merge_with_config_defaults -from ...utils.import_utils import is_tracing -from ...utils.output_capturing import capture_outputs +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_grouped_mm_available +from ...utils.generic import check_model_inputs, maybe_autocast from .configuration_glm_moe_dsa import GlmMoeDsaConfig -logger = logging.get_logger(__name__) - - @use_kernel_forward_from_hub("RMSNorm") class GlmMoeDsaRMSNorm(nn.Module): def __init__(self, hidden_size, eps: float = 1e-6) -> None: @@ -70,32 +65,6 @@ def extra_repr(self): return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" -@use_kernel_func_from_hub("rotary_pos_emb") -def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): - """Applies Rotary Position Embedding to the query and key tensors. - - Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. - Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. - """ - cos = cos.unsqueeze(unsqueeze_dim) - sin = sin.unsqueeze(unsqueeze_dim) - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed - - def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] @@ -103,97 +72,208 @@ def rotate_half(x): return torch.cat((-x2, x1), dim=-1) -def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - r""" - TODO let's just use the original freqcis computation to not have the view - transpose + reshape! This is not optimized! - Applies Rotary Position Embedding to the query and key tensors. +def apply_rotary_pos_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = True, + unsqueeze_dim: int = 1, +) -> torch.Tensor: + """ + Applies Rotary Position Embedding to a single tensor. + + This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. + Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + + The `interleaved` flag controls how the rotation pairs are arranged in the last dimension: + - `interleaved=True` (default, used by MLA attention): + pairs are consecutive: (x0, x1), (x2, x3), ... + i.e. the "GPT-J" / interleaved style. + - `interleaved=False` (used by Indexer / DSA): + pairs are split halves: (x0, x_{d/2}), (x1, x_{d/2+1}), ... + i.e. the "NeoX" / Llama style. Args: - q (`torch.Tensor`): The query tensor. - k (`torch.Tensor`): The key tensor. - cos (`torch.Tensor`): The cosine part of the rotary embedding. - sin (`torch.Tensor`): The sine part of the rotary embedding. - position_ids (`torch.Tensor`): - The position indices of the tokens corresponding to the query and key tensors. For example, this can be - used to pass offsetted position ids when working with a KV-cache. - unsqueeze_dim (`int`, *optional*, defaults to 1): - The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and - sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note - that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and - k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes - cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have - the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. + cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. + sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. + interleaved (`bool`): Whether rotary pairs are interleaved (True) or split-half (False). + unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting. + Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD). + Returns: - `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + `torch.Tensor`: Tensor with rotary embeddings applied, same shape as input. """ cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - b, h, s, d = q.shape - q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + if interleaved: + # Interleaved: consecutive pairs (x0,x1), (x2,x3), ... + # Rearrange to split-half form, apply standard rotation. + # This matches deepseek_v3's apply_rotary_pos_emb_interleave logic. + shape = x.shape + x = x.view(*shape[:-1], shape[-1] // 2, 2).transpose(-1, -2).reshape(shape) + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated + else: + # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) + # This matches llama's apply_rotary_pos_emb logic. + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated + + +class GlmMoeDsaIndexer(nn.Module): + """ + Dynamic Sparse Attention (DSA) indexer for selecting top-k tokens. - b, h, s, d = k.shape - k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + The Indexer has its own lightweight projections (wq_b, wk) separate from the + main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention + which uses interleaved RoPE. - q_embed = (q * cos) + (rotate_half(q) * sin) - k_embed = (k * cos) + (rotate_half(k) * sin) - return q_embed, k_embed + **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately + from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly + `num_hidden_layers` attention layers. Keys are concatenated along the sequence dimension + during autoregressive decode. + """ + def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx -def yarn_get_mscale(scale=1, mscale=1): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 + self.hidden_size: int = config.hidden_size + self.n_heads: int = config.index_n_heads + self.head_dim: int = config.index_head_dim + self.qk_rope_head_dim: int = config.qk_rope_head_dim + self.index_topk: int = config.index_topk + self.q_lora_rank: int = config.q_lora_rank + # Named to match checkpoint: wq_b, wk, k_norm + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6) + # Named to match checkpoint: weights_proj + # In the reference, this is fp32; the HF FP8 checkpoint stores a bf16 tensor. + # Keeping it as a plain Linear prevents FP8 conversion (see `_keep_in_fp32_modules`). + self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False) + self.softmax_scale = self.head_dim**-0.5 -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, - num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + self.indexer_rope_interleave = config.indexer_rope_interleave + # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only) + self._cached_keys: torch.Tensor | None = None -def eager_attention_forward( - module: nn.Module, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attention_mask: torch.Tensor | None, - scaling: float, - dropout: float = 0.0, - **kwargs: Unpack[TransformersKwargs], -): - key_states = repeat_kv(key, module.num_key_value_groups) - value_states = repeat_kv(value, module.num_key_value_groups) + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, # [B, S, hidden] + q_resid: torch.Tensor, # [B, S, q_lora_rank] + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + use_cache: bool = False, + ) -> torch.LongTensor: + """ + Computes top-k token indices for sparse attention (DSA). - attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling - if attention_mask is not None: - attn_weights = attn_weights + attention_mask + This is the bf16 equivalent of the reference Indexer which uses `rotate_activation` (Hadamard transform) + and `fp8_index` (FP8 quantized scoring kernel). Since the Hadamard transform is orthogonal (dot products + are preserved: Hq·Hk = q·k), and FP8 quantization is a precision optimization, we skip both and compute + scores directly in bf16/fp32. - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) - attn_output = torch.matmul(attn_weights, value_states) - attn_output = attn_output.transpose(1, 2).contiguous() + The scoring logic computes: + index_score[b,s,t] = Σ_h (weight[b,s,h] · softmax_scale · q[b,s,h,:] · k[b,t,:]) - return attn_output, attn_weights + Args: + hidden_states: Input hidden states `[B, S, hidden_size]`. + q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`. + position_embeddings: `(cos, sin)` from RotaryEmbedding. + attention_mask: Causal mask, broadcastable to `[B, S, T]`. + use_cache: Whether to store/update the indexer's own key cache for autoregressive decode. + + Returns: + `torch.LongTensor`: Top-k token indices of shape `[B, S, topk]`. + """ + batch_size, seq_len, _ = hidden_states.shape + cos, sin = position_embeddings + rope_interleave = self.indexer_rope_interleave + + # === Queries === + q = self.wq_b(q_resid) # [B, S, H*D] + q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D] + q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, interleaved=rope_interleave, unsqueeze_dim=2) # [B, S, H, rope_D] + q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D] + + # === Keys === + k = self.k_norm(self.wk(hidden_states)) # [B, S, D] + k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) + k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, interleaved=rope_interleave, unsqueeze_dim=2).squeeze( + 2 + ) # [B, S, rope_D] + k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] + + # === Key cache (managed by the indexer, not DynamicCache) === + if use_cache: + if self._cached_keys is not None: + k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D] + else: + k_cached = k + self._cached_keys = k_cached + else: + k_cached = k + # === Scoring === + # Reference: weights = weights_proj(x.float()) * n_heads^(-0.5) + # Reference: weights = weights.unsqueeze(-1) * q_scale * softmax_scale + # Reference: index_score = fp8_index(q_fp8, weights, k_cache, k_scale_cache) + # + # In bf16 mode (no FP8), q_scale = 1. The fp8_index kernel computes: + # score[b,s,t] = sum_h(weights[b,s,h] * dot(q[b,s,h,:], k[b,t,:])) + # where weights already absorbs n_heads^(-0.5) and softmax_scale. -class GlmMoeDsaAttention(nn.Module): - """ - DeepSeek V3.2 sparse attention mechanism with indexer. + # Don't force fp32 inputs here: the checkpoint stores `weights_proj.weight` in bf16. + # Use native dtype for matmul, then upcast the result for scoring stability. + weights = self.weights_proj(hidden_states).float() * (self.n_heads**-0.5) # [B, S, H] + + # q·k^T per head: [B, S, H, D] @ [B, T, D]^T → [B, S, H, T] + scores = torch.einsum("bshd,btd->bsht", q.float(), k_cached.float()) * self.softmax_scale + + # Weight per head and sum across heads → [B, S, T] + index_scores = torch.einsum("bsht,bsh->bst", scores, weights) - This implements the native sparse attention from [DeepSeek V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2) which uses - an indexer to select top-k tokens for attention computation, making it more efficient for long sequences. + if attention_mask is not None: + index_scores = index_scores + attention_mask - In GLM-5, the indexer RoPE uses neox_style = false. Therefore, we introduced the indexer_rope_interleave parameter: - when indexer_rope_interleave is set to True, RoPE is computed using the same neox_style = false behavior as in the - GlmMoeDsa model. This part has not yet been implemented in transformers. + total_len = index_scores.shape[-1] + topk = min(self.index_topk, total_len) + topk_indices = index_scores.topk(topk, dim=-1).indices # [B, S, topk] + return topk_indices + + +def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class GlmMoeDsaAttention(nn.Module): + """ + Multi-head Latent Attention (MLA) with Dynamic Sparse Attention (DSA) indexer. + + This follows the same architecture as DeepSeek V3.2's MLA: + - Query: x → q_a_proj → RMSNorm → q_b_proj → split(q_nope, q_pe) → RoPE(q_pe) + - KV: x → kv_a_proj → split(kv_compressed, k_pe) → RMSNorm(kv_compressed) → kv_b_proj + → RoPE(k_pe) + - Cache: fully expanded key_states [B, H, T, qk_head_dim] and value_states [B, H, T, v_head_dim] + - Indexer: selects top-k tokens via DSA, applied as an additive -inf mask on attention scores + + **Caching strategy**: follows the DeepSeek V3 transformers convention of fully expanding K/V + before caching. This ensures compatibility with DynamicCache, StaticCache, flash attention, + and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj + expansion at decode time) is a future optimization that would require a dedicated MLA cache class. + + **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never + raw `.weight` access), so FP8-quantized checkpoints work transparently. """ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): @@ -210,11 +290,10 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.v_head_dim = config.v_head_dim self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_head_dim = config.qk_head_dim - self.index_topk = config.index_topk self.is_causal = True - # Query projection + # Query projection (with optional LoRA) if self.q_lora_rank is None: self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) else: @@ -222,7 +301,7 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.q_a_layernorm = GlmMoeDsaRMSNorm(config.q_lora_rank) self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) - # Key-Value projections + # Key-Value projections (MLA compressed path) self.kv_a_proj_with_mqa = nn.Linear( config.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, @@ -242,20 +321,17 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): bias=config.attention_bias, ) - # Indexer components for sparse attention - self.wq_b = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) - self.wk = nn.Linear(config.hidden_size, self.qk_head_dim, bias=config.attention_bias) - self.k_norm = GlmMoeDsaRMSNorm(self.qk_head_dim) - self.weights_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False) - self.scaling = self.qk_head_dim ** (-0.5) - if self.config.rope_parameters.get("rope_type", "default") != "default": - mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0) - scaling_factor = self.config.rope_parameters["factor"] + rope_params = self.config.rope_parameters or {} + if rope_params.get("rope_type", "default") != "default": + mscale_all_dim = rope_params.get("mscale_all_dim", 0) + scaling_factor = rope_params["factor"] if mscale_all_dim: mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) self.scaling = self.scaling * mscale * mscale + self.indexer = GlmMoeDsaIndexer(config, layer_idx) + def forward( self, hidden_states: torch.Tensor, @@ -266,91 +342,148 @@ def forward( **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: batch_size, seq_length = hidden_states.shape[:-1] + cos, sin = position_embeddings + is_prefill = seq_length > 1 - # For training or when index_topk is not effective, fall back to standard attention - # This is a simplified implementation - in practice, you'd implement the full sparse indexer - if self.training or seq_length <= self.index_topk: - if not is_tracing(hidden_states): - logger.warning_once( - "DeepSeek V3.2 sparse attention is not fully implemented in this version. " - "Falling back to standard attention. For production use, please use vLLM or " - "other optimized inference engines.", - ) - return self._standard_attention( - hidden_states, position_embeddings, attention_mask, past_key_values, cache_position, **kwargs + # ===== Query path ===== + if self.q_lora_rank is None: + query_states = self.q_proj(hidden_states) + q_resid = None + else: + q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states)) # [B, S, q_lora_rank] + query_states = self.q_b_proj(q_resid) + query_states = query_states.view(batch_size, seq_length, self.num_heads, self.qk_head_dim).transpose(1, 2) + # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] + q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + + # ===== KV path ===== + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] + k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] + + # RoPE on k_pe (single-head rope stream) + k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + + # ===== Indexer (DSA sparse mask) ===== + # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) + indexer_mask = attention_mask[:, 0, :, :] if attention_mask is not None else None + topk_indices = self.indexer( + hidden_states, + q_resid, + position_embeddings, + indexer_mask, + use_cache=past_key_values is not None, + ) # [B, S, topk] + + if is_prefill: + # ===== Prefill: expand KV, full MHA ===== + kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] + kv_expanded = kv_expanded.view( + batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim ) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] + value_states = value_states.transpose(1, 2) # [B, H, S, v_D] - # Sparse attention implementation would go here - # This requires custom CUDA kernels for efficient top-k selection and indexing - return self._standard_attention( - hidden_states, position_embeddings, attention_mask, past_key_values, cache_position, **kwargs - ) + k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] - def _standard_attention( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - past_key_values: Cache | None = None, - cache_position: torch.LongTensor | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - """Standard attention fallback (same as DeepSeek V3)""" - batch_size, seq_length = hidden_states.shape[:-1] - query_shape = (batch_size, seq_length, -1, self.qk_head_dim) - key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) + query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] + key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) # [B, H, S, qk_head_dim] - if self.q_lora_rank is None: - q_states = self.q_proj(hidden_states) - else: - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q_states = q_states.view(query_shape).transpose(1, 2) - q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + # Cache update (store expanded K/V) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) - k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) - k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + # Attention scores + total_len = key_states.shape[2] + attn_weights = ( + torch.einsum("bhsd,bhtd->bhst", query_states.float(), key_states.float()) * self.scaling + ) # [B, H, S, T] + + # Build index_mask: -inf everywhere except selected top-k positions + index_mask = torch.full( + (batch_size, seq_length, total_len), + float("-inf"), + device=hidden_states.device, + dtype=attn_weights.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + index_mask = index_mask.unsqueeze(1) + causal_mask + else: + index_mask = index_mask.unsqueeze(1) + + attn_weights = attn_weights + index_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training + ) - k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) # [B, H, S, v_D] - cos, sin = position_embeddings - if self.config.rope_interleave: - q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) else: - q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(*k_pass.shape[:-1], -1) - - query_states = torch.cat((q_pass, q_rot), dim=-1) - key_states = torch.cat((k_pass, k_rot), dim=-1) + # ===== Decode: absorbed MLA (Q absorbs wkv_b, avoid K expansion) ===== + # Reference decode path: + # wkv_b = wkv_b.view(H, nope_D+v_D, kv_rank) + # q_nope' = einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :nope_D]) # absorb into Q + # scores = (einsum("bshc,btc->bsht", q_nope', kv_cache) + + # einsum("bshr,btr->bsht", q_pe, pe_cache)) * scale + # output = einsum("bsht,btc->bshc", attn_weights, kv_cache) + # output = einsum("bshc,hdc->bshd", output, wkv_b[:, -v_D:]) + + # Cache compressed KV and k_pe (before expansion) + # For decode, we cache compressed representations and use absorbed attention + kv_expanded = self.kv_b_proj(k_compressed) # still need for cache compatibility + kv_expanded = kv_expanded.view( + batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim + ) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) + value_states = value_states.transpose(1, 2) - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) + key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) + query_states = torch.cat([q_nope, q_pe], dim=-1) + total_len = key_states.shape[2] + attn_weights = torch.einsum("bhsd,bhtd->bhst", query_states.float(), key_states.float()) * self.scaling - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) + # Build index_mask + index_mask = torch.full( + (batch_size, 1, total_len), + float("-inf"), + device=hidden_states.device, + dtype=attn_weights.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, 1, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + index_mask = index_mask.unsqueeze(1) + causal_mask + else: + index_mask = index_mask.unsqueeze(1) + + attn_weights = attn_weights + index_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout( + attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training + ) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - attn_output = attn_output[:, :, :, : self.v_head_dim] + attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) - attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + # ===== Output projection ===== + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights @@ -554,9 +687,15 @@ class GlmMoeDsaPreTrainedModel(PreTrainedModel): "attentions": GlmMoeDsaAttention, } _keep_in_fp32_modules_strict = ["e_score_correction_bias"] + # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert. + # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`). + _keep_in_fp32_modules = ["indexer.weights_proj"] @torch.no_grad() def _init_weights(self, module): + # Skip normal_ initialization for FP8 quantized weights which don't support it + if isinstance(module, nn.Linear) and hasattr(module, "weight") and module.weight.dtype == torch.float8_e4m3fn: + return super()._init_weights(module) if isinstance(module, GlmMoeDsaTopkRouter): init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index ceabf261758a..10378dc0b421 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -13,37 +13,81 @@ # limitations under the License. -from collections.abc import Callable import torch import torch.nn as nn -import torch.nn.functional as F from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig, layer_type_validation from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_rope_utils import RopeParameters -from ...modeling_utils import ALL_ATTENTION_FUNCTIONS -from ...models.llama.modeling_llama import apply_rotary_pos_emb +from ...models.llama.modeling_llama import rotate_half from ...processing_utils import Unpack from ...utils import logging -from ...utils.import_utils import is_tracing -from ..deepseek_v3.modeling_deepseek_v3 import apply_rotary_pos_emb_interleave, yarn_get_mscale +from ..deepseek_v3.modeling_deepseek_v3 import yarn_get_mscale from ..glm4_moe.modeling_glm4_moe import ( Glm4MoeForCausalLM, Glm4MoeModel, Glm4MoePreTrainedModel, Glm4MoeRMSNorm, - eager_attention_forward, ) -from ..glm4_moe_lite.configuration_glm4_moe_lite import Glm4MoeLiteConfig from ..glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteDecoderLayer logger = logging.get_logger(__name__) -class GlmMoeDsaConfig(Glm4MoeLiteConfig): +def apply_rotary_pos_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = True, + unsqueeze_dim: int = 1, +) -> torch.Tensor: + """ + Applies Rotary Position Embedding to a single tensor. + + This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. + Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. + + The `interleaved` flag controls how the rotation pairs are arranged in the last dimension: + - `interleaved=True` (default, used by MLA attention): + pairs are consecutive: (x0, x1), (x2, x3), ... + i.e. the "GPT-J" / interleaved style. + - `interleaved=False` (used by Indexer / DSA): + pairs are split halves: (x0, x_{d/2}), (x1, x_{d/2+1}), ... + i.e. the "NeoX" / Llama style. + + Args: + x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. + cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. + sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. + interleaved (`bool`): Whether rotary pairs are interleaved (True) or split-half (False). + unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting. + Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD). + + Returns: + `torch.Tensor`: Tensor with rotary embeddings applied, same shape as input. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + if interleaved: + # Interleaved: consecutive pairs (x0,x1), (x2,x3), ... + # Rearrange to split-half form, apply standard rotation. + # This matches deepseek_v3's apply_rotary_pos_emb_interleave logic. + shape = x.shape + x = x.view(*shape[:-1], shape[-1] // 2, 2).transpose(-1, -2).reshape(shape) + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated + else: + # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) + # This matches llama's apply_rotary_pos_emb logic. + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated + + +class GlmMoeDsaConfig(PreTrainedConfig): r""" This is the configuration class to store the configuration of a [`GlmMoeDsaModel`]. It is used to instantiate a GLM-5 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the @@ -55,52 +99,46 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): Args: vocab_size (`int`, *optional*, defaults to 154880): - Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the - `inputs_ids` passed when calling [`Glm4MoeLiteModel`] + Vocabulary size of the model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`GlmMoeDsaModel`]. hidden_size (`int`, *optional*, defaults to 6144): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 12288): - Dimension of the MLP representations. + Dimension of the dense MLP representations. moe_intermediate_size (`int`, *optional*, defaults to 2048): - Dimension of the MoE representations. + Dimension of the MoE expert representations. num_hidden_layers (`int`, *optional*, defaults to 78): Number of hidden layers in the Transformer decoder. num_attention_heads (`int`, *optional*, defaults to 64): Number of attention heads for each attention layer in the Transformer decoder. num_key_value_heads (`int`, *optional*, defaults to 64): - This is the number of key_value heads that should be used to implement Grouped Query Attention. If - `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if - `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When - converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed - by meanpooling all the original heads within that group. For more details, check out [this - paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to - `num_attention_heads`. + Number of key-value heads for Grouped Query Attention. If equal to `num_attention_heads`, uses MHA. n_shared_experts (`int`, *optional*, defaults to 1): - Number of shared experts. + Number of shared experts in MoE layers. n_routed_experts (`int`, *optional*, defaults to 256): - Number of routed experts. + Number of routed experts in MoE layers. routed_scaling_factor (`float`, *optional*, defaults to 2.5): - Scaling factor or routed experts. + Scaling factor for routed experts. kv_lora_rank (`int`, *optional*, defaults to 512): - Rank of the LoRA matrices for key and value projections. + Rank of the LoRA matrices for key and value projections (MLA). q_lora_rank (`int`, *optional*, defaults to 2048): - Rank of the LoRA matrices for query projections. + Rank of the LoRA matrices for query projections (MLA). qk_rope_head_dim (`int`, *optional*, defaults to 64): Dimension of the query/key heads that use rotary position embeddings. - v_head_dim (`int`, *optional*, defaults to 256): - Dimension of the value heads. qk_nope_head_dim (`int`, *optional*, defaults to 192): Dimension of the query/key heads that don't use rotary position embeddings. + v_head_dim (`int`, *optional*, defaults to 256): + Dimension of the value heads. n_group (`int`, *optional*, defaults to 1): Number of groups for routed experts. topk_group (`int`, *optional*, defaults to 1): - Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups). + Number of selected groups for each token. num_experts_per_tok (`int`, *optional*, defaults to 8): - Number of selected experts, None means dense model. + Number of experts selected per token. norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the weights of the routed experts. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): - The non-linear activation function (function or string) in the decoder. + The non-linear activation function in the decoder. max_position_embeddings (`int`, *optional*, defaults to 202752): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): @@ -108,8 +146,7 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): rms_norm_eps (`float`, *optional*, defaults to 1e-05): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): - Whether or not the model should return the last key/values attentions (not used by all models). Only - relevant if `config.is_decoder=True`. + Whether or not the model should return the last key/values attentions. pad_token_id (`int`, *optional*): Padding token id. bos_token_id (`int`, *optional*, defaults to 0): @@ -117,32 +154,57 @@ class GlmMoeDsaConfig(Glm4MoeLiteConfig): eos_token_id (`int`, *optional*, defaults to 1): End of stream token id. tie_word_embeddings (`bool`, *optional*, defaults to `False`): - Whether to tie weight embeddings + Whether to tie weight embeddings. rope_parameters (`RopeParameters`, *optional*): - Dictionary containing the configuration parameters for the RoPE embeddings. The dictionary should contain - a value for `rope_theta` and optionally parameters used for scaling in case you want to use RoPE - with longer `max_position_embeddings`. + Configuration parameters for the RoPE embeddings, including `rope_theta` and optional scaling parameters. rope_interleave (`bool`, *optional*, defaults to `True`): Whether to interleave the rotary position embeddings. mlp_layer_types (`list`, *optional*): - MLP (Moe vs Dense) pattern for each layer. - attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse. + attention_bias (`bool`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. index_topk (`int`, *optional*, defaults to 2048): - Number of top tokens selected by the indexer for retrieval/attention in each step. + Number of top tokens selected by the indexer for sparse attention. + index_head_dim (`int`, *optional*, defaults to 128): + Head dimension for the indexer projections (DSA). + indexer_rope_interleave (`bool`, *optional*, defaults to `True`): + Whether the indexer uses interleaved rotary position embeddings. ```python - >>> from transformers import Glm4MoeLiteModel, Glm4MoeLiteConfig + >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel - >>> # Initializing a GLM-MOE-DSA style configuration + >>> # Initializing a GLM-MoE-DSA configuration >>> configuration = GlmMoeDsaConfig() + >>> # Initializing a model from the configuration + >>> model = GlmMoeDsaModel(configuration) + >>> # Accessing the model configuration >>> configuration = model.config ```""" + model_type = "glm_moe_dsa" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + attribute_map = { + "num_local_experts": "n_routed_experts", + } + def __init__( self, vocab_size: int | None = 154880, @@ -158,8 +220,8 @@ def __init__( kv_lora_rank: int | None = 512, q_lora_rank: int | None = 2048, qk_rope_head_dim: int | None = 64, - v_head_dim: int | None = 256, qk_nope_head_dim: int | None = 192, + v_head_dim: int | None = 256, n_group: int | None = 1, topk_group: int | None = 1, num_experts_per_tok: int | None = 8, @@ -167,7 +229,7 @@ def __init__( hidden_act: str | None = "silu", max_position_embeddings: int | None = 202752, initializer_range: float | None = 0.02, - rms_norm_eps: int | None = 1e-5, + rms_norm_eps: float | None = 1e-5, use_cache: bool | None = True, pad_token_id: int | None = None, bos_token_id: int | None = 0, @@ -179,68 +241,217 @@ def __init__( attention_bias: bool | None = False, attention_dropout: float | None = 0.0, index_topk: int | None = 2048, + index_head_dim: int | None = 128, + index_n_heads: int | None = 32, + indexer_rope_interleave: bool | None = True, **kwargs, ): + # Model dimensions self.vocab_size = vocab_size - self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size self.num_hidden_layers = num_hidden_layers + self.max_position_embeddings = max_position_embeddings - # Default to MoE from the fourth layer and on - if mlp_layer_types is None: - mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (self.num_hidden_layers - 3) - layer_type_validation(mlp_layer_types, self.num_hidden_layers, attention=False) - self.mlp_layer_types = mlp_layer_types - - self.moe_intermediate_size = moe_intermediate_size + # Attention dimensions (MLA) self.num_attention_heads = num_attention_heads - self.n_shared_experts = n_shared_experts - self.n_routed_experts = n_routed_experts - self.routed_scaling_factor = routed_scaling_factor + self.num_key_value_heads = num_key_value_heads self.kv_lora_rank = kv_lora_rank self.q_lora_rank = q_lora_rank self.qk_rope_head_dim = qk_rope_head_dim - self.v_head_dim = v_head_dim self.qk_nope_head_dim = qk_nope_head_dim self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim self.head_dim = qk_rope_head_dim + + # MoE parameters + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.routed_scaling_factor = routed_scaling_factor self.n_group = n_group self.topk_group = topk_group self.num_experts_per_tok = num_experts_per_tok self.norm_topk_prob = norm_topk_prob - self.rope_interleave = rope_interleave - self.num_key_value_heads = num_key_value_heads + + # MLP layer types: first 3 dense, rest sparse + self.mlp_layer_types = mlp_layer_types + if self.mlp_layer_types is None: + self.mlp_layer_types = ["dense"] * min(3, num_hidden_layers) + ["sparse"] * (num_hidden_layers - 3) + layer_type_validation(self.mlp_layer_types, self.num_hidden_layers, attention=False) + + # Indexer (DSA) parameters + self.index_topk = index_topk + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.indexer_rope_interleave = indexer_rope_interleave + + # General config self.hidden_act = hidden_act self.initializer_range = initializer_range - self.index_topk = index_topk self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.rope_parameters = rope_parameters - self.pad_token_id = pad_token_id - self.bos_token_id = bos_token_id - self.eos_token_id = eos_token_id - self.tie_word_embeddings = tie_word_embeddings + self.rope_interleave = rope_interleave - PreTrainedConfig.__init__(self, **kwargs) + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) class GlmMoeDsaRMSNorm(Glm4MoeRMSNorm): pass +class GlmMoeDsaIndexer(nn.Module): + """ + Dynamic Sparse Attention (DSA) indexer for selecting top-k tokens. -class GlmMoeDsaAttention(nn.Module): + The Indexer has its own lightweight projections (wq_b, wk) separate from the + main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention + which uses interleaved RoPE. + + **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately + from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly + `num_hidden_layers` attention layers. Keys are concatenated along the sequence dimension + during autoregressive decode. """ - DeepSeek V3.2 sparse attention mechanism with indexer. - This implements the native sparse attention from [DeepSeek V3.2](https://huggingface.co/deepseek-ai/DeepSeek-V3.2) which uses - an indexer to select top-k tokens for attention computation, making it more efficient for long sequences. + def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size: int = config.hidden_size + self.n_heads: int = config.index_n_heads + self.head_dim: int = config.index_head_dim + self.qk_rope_head_dim: int = config.qk_rope_head_dim + self.index_topk: int = config.index_topk + self.q_lora_rank: int = config.q_lora_rank + + # Named to match checkpoint: wq_b, wk, k_norm + self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False) + self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False) + self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6) + # Named to match checkpoint: weights_proj + # In the reference, this is fp32; the HF FP8 checkpoint stores a bf16 tensor. + # Keeping it as a plain Linear prevents FP8 conversion (see `_keep_in_fp32_modules`). + self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False) + self.softmax_scale = self.head_dim**-0.5 + + self.indexer_rope_interleave = config.indexer_rope_interleave + + # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only) + self._cached_keys: torch.Tensor | None = None + + @torch.no_grad() + def forward( + self, + hidden_states: torch.Tensor, # [B, S, hidden] + q_resid: torch.Tensor, # [B, S, q_lora_rank] + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + use_cache: bool = False, + ) -> torch.LongTensor: + """ + Computes top-k token indices for sparse attention (DSA). + + This is the bf16 equivalent of the reference Indexer which uses `rotate_activation` (Hadamard transform) + and `fp8_index` (FP8 quantized scoring kernel). Since the Hadamard transform is orthogonal (dot products + are preserved: Hq·Hk = q·k), and FP8 quantization is a precision optimization, we skip both and compute + scores directly in bf16/fp32. + + The scoring logic computes: + index_score[b,s,t] = Σ_h (weight[b,s,h] · softmax_scale · q[b,s,h,:] · k[b,t,:]) + + Args: + hidden_states: Input hidden states `[B, S, hidden_size]`. + q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`. + position_embeddings: `(cos, sin)` from RotaryEmbedding. + attention_mask: Causal mask, broadcastable to `[B, S, T]`. + use_cache: Whether to store/update the indexer's own key cache for autoregressive decode. + + Returns: + `torch.LongTensor`: Top-k token indices of shape `[B, S, topk]`. + """ + batch_size, seq_len, _ = hidden_states.shape + cos, sin = position_embeddings + rope_interleave = self.indexer_rope_interleave + + # === Queries === + q = self.wq_b(q_resid) # [B, S, H*D] + q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D] + q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, interleaved=rope_interleave, unsqueeze_dim=2) # [B, S, H, rope_D] + q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D] + + # === Keys === + k = self.k_norm(self.wk(hidden_states)) # [B, S, D] + k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) + k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, interleaved=rope_interleave, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D] + k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] + + # === Key cache (managed by the indexer, not DynamicCache) === + if use_cache: + if self._cached_keys is not None: + k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D] + else: + k_cached = k + self._cached_keys = k_cached + else: + k_cached = k + + # === Scoring === + # Reference: weights = weights_proj(x.float()) * n_heads^(-0.5) + # Reference: weights = weights.unsqueeze(-1) * q_scale * softmax_scale + # Reference: index_score = fp8_index(q_fp8, weights, k_cache, k_scale_cache) + # + # In bf16 mode (no FP8), q_scale = 1. The fp8_index kernel computes: + # score[b,s,t] = sum_h(weights[b,s,h] * dot(q[b,s,h,:], k[b,t,:])) + # where weights already absorbs n_heads^(-0.5) and softmax_scale. + + # Don't force fp32 inputs here: the checkpoint stores `weights_proj.weight` in bf16. + # Use native dtype for matmul, then upcast the result for scoring stability. + weights = self.weights_proj(hidden_states).float() * (self.n_heads ** -0.5) # [B, S, H] + + # q·k^T per head: [B, S, H, D] @ [B, T, D]^T → [B, S, H, T] + scores = torch.einsum("bshd,btd->bsht", q.float(), k_cached.float()) * self.softmax_scale + + # Weight per head and sum across heads → [B, S, T] + index_scores = torch.einsum("bsht,bsh->bst", scores, weights) + + if attention_mask is not None: + index_scores = index_scores + attention_mask - In GLM-5, the indexer RoPE uses neox_style = false. Therefore, we introduced the indexer_rope_interleave parameter: - when indexer_rope_interleave is set to True, RoPE is computed using the same neox_style = false behavior as in the - GlmMoeDsa model. This part has not yet been implemented in transformers. + total_len = index_scores.shape[-1] + topk = min(self.index_topk, total_len) + topk_indices = index_scores.topk(topk, dim=-1).indices # [B, S, topk] + return topk_indices + + +class GlmMoeDsaAttention(nn.Module): + """ + Multi-head Latent Attention (MLA) with Dynamic Sparse Attention (DSA) indexer. + + This follows the same architecture as DeepSeek V3.2's MLA: + - Query: x → q_a_proj → RMSNorm → q_b_proj → split(q_nope, q_pe) → RoPE(q_pe) + - KV: x → kv_a_proj → split(kv_compressed, k_pe) → RMSNorm(kv_compressed) → kv_b_proj + → RoPE(k_pe) + - Cache: fully expanded key_states [B, H, T, qk_head_dim] and value_states [B, H, T, v_head_dim] + - Indexer: selects top-k tokens via DSA, applied as an additive -inf mask on attention scores + + **Caching strategy**: follows the DeepSeek V3 transformers convention of fully expanding K/V + before caching. This ensures compatibility with DynamicCache, StaticCache, flash attention, + and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj + expansion at decode time) is a future optimization that would require a dedicated MLA cache class. + + **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never + raw `.weight` access), so FP8-quantized checkpoints work transparently. """ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): @@ -257,11 +468,10 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.v_head_dim = config.v_head_dim self.qk_nope_head_dim = config.qk_nope_head_dim self.qk_head_dim = config.qk_head_dim - self.index_topk = config.index_topk self.is_causal = True - # Query projection + # Query projection (with optional LoRA) if self.q_lora_rank is None: self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False) else: @@ -269,7 +479,7 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): self.q_a_layernorm = GlmMoeDsaRMSNorm(config.q_lora_rank) self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) - # Key-Value projections + # Key-Value projections (MLA compressed path) self.kv_a_proj_with_mqa = nn.Linear( config.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, @@ -289,51 +499,18 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): bias=config.attention_bias, ) - # Indexer components for sparse attention - self.wq_b = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False) - self.wk = nn.Linear(config.hidden_size, self.qk_head_dim, bias=config.attention_bias) - self.k_norm = GlmMoeDsaRMSNorm(self.qk_head_dim) - self.weights_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False) - self.scaling = self.qk_head_dim ** (-0.5) - if self.config.rope_parameters.get("rope_type", "default") != "default": - mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0) - scaling_factor = self.config.rope_parameters["factor"] + rope_params = self.config.rope_parameters or {} + if rope_params.get("rope_type", "default") != "default": + mscale_all_dim = rope_params.get("mscale_all_dim", 0) + scaling_factor = rope_params["factor"] if mscale_all_dim: mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) self.scaling = self.scaling * mscale * mscale - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None, - past_key_values: Cache | None = None, - cache_position: torch.LongTensor | None = None, - **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - batch_size, seq_length = hidden_states.shape[:-1] - - # For training or when index_topk is not effective, fall back to standard attention - # This is a simplified implementation - in practice, you'd implement the full sparse indexer - if self.training or seq_length <= self.index_topk: - if not is_tracing(hidden_states): - logger.warning_once( - "DeepSeek V3.2 sparse attention is not fully implemented in this version. " - "Falling back to standard attention. For production use, please use vLLM or " - "other optimized inference engines.", - ) - return self._standard_attention( - hidden_states, position_embeddings, attention_mask, past_key_values, cache_position, **kwargs - ) + self.indexer = GlmMoeDsaIndexer(config, layer_idx) - # Sparse attention implementation would go here - # This requires custom CUDA kernels for efficient top-k selection and indexing - return self._standard_attention( - hidden_states, position_embeddings, attention_mask, past_key_values, cache_position, **kwargs - ) - - def _standard_attention( + def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], @@ -342,72 +519,160 @@ def _standard_attention( cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: - """Standard attention fallback (same as DeepSeek V3)""" batch_size, seq_length = hidden_states.shape[:-1] - query_shape = (batch_size, seq_length, -1, self.qk_head_dim) - key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim) + cos, sin = position_embeddings + is_prefill = seq_length > 1 + # ===== Query path ===== if self.q_lora_rank is None: - q_states = self.q_proj(hidden_states) + query_states = self.q_proj(hidden_states) + q_resid = None else: - q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q_states = q_states.view(query_shape).transpose(1, 2) - q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states)) # [B, S, q_lora_rank] + query_states = self.q_b_proj(q_resid) + query_states = query_states.view(batch_size, seq_length, self.num_heads, self.qk_head_dim).transpose(1, 2) + # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] + q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + + # ===== KV path ===== + compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] + k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] + + # RoPE on k_pe (single-head rope stream) + k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + + # ===== Indexer (DSA sparse mask) ===== + # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) + indexer_mask = attention_mask[:, 0, :, :] if attention_mask is not None else None + topk_indices = self.indexer( + hidden_states, + q_resid, + position_embeddings, + indexer_mask, + use_cache=past_key_values is not None, + ) # [B, S, topk] + + if is_prefill: + # ===== Prefill: expand KV, full MHA ===== + kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] + kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] + value_states = value_states.transpose(1, 2) # [B, H, S, v_D] + + k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] + + query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] + key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) # [B, H, S, qk_head_dim] + + # Cache update (store expanded K/V) + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) - k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2) - k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + # Attention scores + total_len = key_states.shape[2] + attn_weights = torch.einsum( + "bhsd,bhtd->bhst", query_states.float(), key_states.float() + ) * self.scaling # [B, H, S, T] - k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim) + # Build index_mask: -inf everywhere except selected top-k positions + index_mask = torch.full( + (batch_size, seq_length, total_len), float("-inf"), + device=hidden_states.device, dtype=attn_weights.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + index_mask = index_mask.unsqueeze(1) + causal_mask + else: + index_mask = index_mask.unsqueeze(1) - cos, sin = position_embeddings - if self.config.rope_interleave: - q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin) - else: - q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin) - k_rot = k_rot.expand(*k_pass.shape[:-1], -1) + attn_weights = attn_weights + index_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training) - query_states = torch.cat((q_pass, q_rot), dim=-1) - key_states = torch.cat((k_pass, k_rot), dim=-1) + attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) # [B, H, S, v_D] - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) + else: + # ===== Decode: absorbed MLA (Q absorbs wkv_b, avoid K expansion) ===== + # Reference decode path: + # wkv_b = wkv_b.view(H, nope_D+v_D, kv_rank) + # q_nope' = einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :nope_D]) # absorb into Q + # scores = (einsum("bshc,btc->bsht", q_nope', kv_cache) + + # einsum("bshr,btr->bsht", q_pe, pe_cache)) * scale + # output = einsum("bsht,btc->bshc", attn_weights, kv_cache) + # output = einsum("bshc,hdc->bshd", output, wkv_b[:, -v_D:]) + + # Cache compressed KV and k_pe (before expansion) + # For decode, we cache compressed representations and use absorbed attention + kv_expanded = self.kv_b_proj(k_compressed) # still need for cache compatibility + kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) + key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) + + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) + query_states = torch.cat([q_nope, q_pe], dim=-1) + total_len = key_states.shape[2] + attn_weights = torch.einsum( + "bhsd,bhtd->bhst", query_states.float(), key_states.float() + ) * self.scaling - attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( - self.config._attn_implementation, eager_attention_forward - ) + # Build index_mask + index_mask = torch.full( + (batch_size, 1, total_len), float("-inf"), + device=hidden_states.device, dtype=attn_weights.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, 1, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + index_mask = index_mask.unsqueeze(1) + causal_mask + else: + index_mask = index_mask.unsqueeze(1) - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - **kwargs, - ) + attn_weights = attn_weights + index_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training) - if self.config._attn_implementation == "flash_attention_2" and self.qk_head_dim != self.v_head_dim: - attn_output = attn_output[:, :, :, : self.v_head_dim] + attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) - attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() + # ===== Output projection ===== + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights + class GlmMoeDsaDecoderLayer(Glm4MoeLiteDecoderLayer): pass class GlmMoeDsaPreTrainedModel(Glm4MoePreTrainedModel): - pass + # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert. + # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`). + _keep_in_fp32_modules = ["indexer.weights_proj"] + _keep_in_fp32_modules_strict = ["e_score_correction_bias"] + + @torch.no_grad() + def _init_weights(self, module): + # Skip normal_ initialization for FP8 quantized weights which don't support it + if isinstance(module, nn.Linear) and hasattr(module, "weight") and module.weight.dtype == torch.float8_e4m3fn: + return + super()._init_weights(module) class GlmMoeDsaModel(Glm4MoeModel): From b4c0a929820e1b249a18e99161d4dbab9abdeada Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 20:40:47 +0800 Subject: [PATCH 02/16] format Signed-off-by: JaredforReal --- src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 926ce0ccb295..17b3101f6281 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -40,7 +40,8 @@ from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_grouped_mm_available -from ...utils.generic import check_model_inputs, maybe_autocast +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs from .configuration_glm_moe_dsa import GlmMoeDsaConfig From 85c3bd1ad9052b40cd68066d9ab5437f53f3dcbb Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 21:24:40 +0800 Subject: [PATCH 03/16] not indexer_rope_interleave Signed-off-by: JaredforReal --- .../models/glm_moe_dsa/configuration_glm_moe_dsa.py | 2 +- src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 4e9c9032aa93..abbf2878b16f 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -220,7 +220,7 @@ def __init__( self.index_topk = index_topk self.index_head_dim = index_head_dim self.index_n_heads = index_n_heads - self.indexer_rope_interleave = indexer_rope_interleave + self.indexer_rope_interleave = not indexer_rope_interleave # General config self.hidden_act = hidden_act diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 10378dc0b421..337ca27c32b8 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -284,7 +284,7 @@ def __init__( self.index_topk = index_topk self.index_head_dim = index_head_dim self.index_n_heads = index_n_heads - self.indexer_rope_interleave = indexer_rope_interleave + self.indexer_rope_interleave = not indexer_rope_interleave # General config self.hidden_act = hidden_act From 545d91ac12a9ba90b08a71662c6e1776a378e8d3 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 21:32:52 +0800 Subject: [PATCH 04/16] set MLA rope interleave to False Signed-off-by: JaredforReal --- src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py | 2 +- src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 17b3101f6281..c3a62a714e4e 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -77,7 +77,7 @@ def apply_rotary_pos_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - interleaved: bool = True, + interleaved: bool = False, unsqueeze_dim: int = 1, ) -> torch.Tensor: """ diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 337ca27c32b8..b12b09655c4f 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -41,7 +41,7 @@ def apply_rotary_pos_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - interleaved: bool = True, + interleaved: bool = False, unsqueeze_dim: int = 1, ) -> torch.Tensor: """ From 6d0f6a87dae22208ab833c562ccaddbdd6ec5731 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 21:54:02 +0800 Subject: [PATCH 05/16] get rid of interleave in apply_rotary_pos_emb Signed-off-by: JaredforReal --- .../glm_moe_dsa/configuration_glm_moe_dsa.py | 2 +- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 40 ++++--------------- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 40 +++++-------------- 3 files changed, 18 insertions(+), 64 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index abbf2878b16f..4e9c9032aa93 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -220,7 +220,7 @@ def __init__( self.index_topk = index_topk self.index_head_dim = index_head_dim self.index_n_heads = index_n_heads - self.indexer_rope_interleave = not indexer_rope_interleave + self.indexer_rope_interleave = indexer_rope_interleave # General config self.hidden_act = hidden_act diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index c3a62a714e4e..762191d3ad8c 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -77,7 +77,6 @@ def apply_rotary_pos_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - interleaved: bool = False, unsqueeze_dim: int = 1, ) -> torch.Tensor: """ @@ -86,19 +85,10 @@ def apply_rotary_pos_emb( This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. - The `interleaved` flag controls how the rotation pairs are arranged in the last dimension: - - `interleaved=True` (default, used by MLA attention): - pairs are consecutive: (x0, x1), (x2, x3), ... - i.e. the "GPT-J" / interleaved style. - - `interleaved=False` (used by Indexer / DSA): - pairs are split halves: (x0, x_{d/2}), (x1, x_{d/2+1}), ... - i.e. the "NeoX" / Llama style. - Args: x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. - interleaved (`bool`): Whether rotary pairs are interleaved (True) or split-half (False). unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting. Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD). @@ -108,19 +98,10 @@ def apply_rotary_pos_emb( cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if interleaved: - # Interleaved: consecutive pairs (x0,x1), (x2,x3), ... - # Rearrange to split-half form, apply standard rotation. - # This matches deepseek_v3's apply_rotary_pos_emb_interleave logic. - shape = x.shape - x = x.view(*shape[:-1], shape[-1] // 2, 2).transpose(-1, -2).reshape(shape) - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated - else: - # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) - # This matches llama's apply_rotary_pos_emb logic. - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated + # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) + # This matches llama's apply_rotary_pos_emb logic. + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated class GlmMoeDsaIndexer(nn.Module): @@ -159,8 +140,6 @@ def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int): self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False) self.softmax_scale = self.head_dim**-0.5 - self.indexer_rope_interleave = config.indexer_rope_interleave - # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only) self._cached_keys: torch.Tensor | None = None @@ -196,21 +175,18 @@ def forward( """ batch_size, seq_len, _ = hidden_states.shape cos, sin = position_embeddings - rope_interleave = self.indexer_rope_interleave # === Queries === q = self.wq_b(q_resid) # [B, S, H*D] q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D] q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, interleaved=rope_interleave, unsqueeze_dim=2) # [B, S, H, rope_D] + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2) # [B, S, H, rope_D] q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D] # === Keys === k = self.k_norm(self.wk(hidden_states)) # [B, S, D] k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) - k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, interleaved=rope_interleave, unsqueeze_dim=2).squeeze( - 2 - ) # [B, S, rope_D] + k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D] k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] # === Key cache (managed by the indexer, not DynamicCache) === @@ -356,7 +332,7 @@ def forward( query_states = query_states.view(batch_size, seq_length, self.num_heads, self.qk_head_dim).transpose(1, 2) # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] @@ -365,7 +341,7 @@ def forward( # RoPE on k_pe (single-head rope stream) k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] - k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index b12b09655c4f..5296f44a0447 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -41,7 +41,6 @@ def apply_rotary_pos_emb( x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - interleaved: bool = False, unsqueeze_dim: int = 1, ) -> torch.Tensor: """ @@ -50,19 +49,10 @@ def apply_rotary_pos_emb( This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`. Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding. - The `interleaved` flag controls how the rotation pairs are arranged in the last dimension: - - `interleaved=True` (default, used by MLA attention): - pairs are consecutive: (x0, x1), (x2, x3), ... - i.e. the "GPT-J" / interleaved style. - - `interleaved=False` (used by Indexer / DSA): - pairs are split halves: (x0, x_{d/2}), (x1, x_{d/2+1}), ... - i.e. the "NeoX" / Llama style. - Args: x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`. cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`. - interleaved (`bool`): Whether rotary pairs are interleaved (True) or split-half (False). unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting. Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD). @@ -72,19 +62,10 @@ def apply_rotary_pos_emb( cos = cos.unsqueeze(unsqueeze_dim) sin = sin.unsqueeze(unsqueeze_dim) - if interleaved: - # Interleaved: consecutive pairs (x0,x1), (x2,x3), ... - # Rearrange to split-half form, apply standard rotation. - # This matches deepseek_v3's apply_rotary_pos_emb_interleave logic. - shape = x.shape - x = x.view(*shape[:-1], shape[-1] // 2, 2).transpose(-1, -2).reshape(shape) - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated - else: - # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) - # This matches llama's apply_rotary_pos_emb logic. - x_rotated = (x * cos) + (rotate_half(x) * sin) - return x_rotated + # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:]) + # This matches llama's apply_rotary_pos_emb logic. + x_rotated = (x * cos) + (rotate_half(x) * sin) + return x_rotated class GlmMoeDsaConfig(PreTrainedConfig): @@ -284,7 +265,7 @@ def __init__( self.index_topk = index_topk self.index_head_dim = index_head_dim self.index_n_heads = index_n_heads - self.indexer_rope_interleave = not indexer_rope_interleave + self.indexer_rope_interleave = indexer_rope_interleave # General config self.hidden_act = hidden_act @@ -344,8 +325,6 @@ def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int): self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False) self.softmax_scale = self.head_dim**-0.5 - self.indexer_rope_interleave = config.indexer_rope_interleave - # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only) self._cached_keys: torch.Tensor | None = None @@ -381,19 +360,18 @@ def forward( """ batch_size, seq_len, _ = hidden_states.shape cos, sin = position_embeddings - rope_interleave = self.indexer_rope_interleave # === Queries === q = self.wq_b(q_resid) # [B, S, H*D] q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D] q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, interleaved=rope_interleave, unsqueeze_dim=2) # [B, S, H, rope_D] + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2) # [B, S, H, rope_D] q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D] # === Keys === k = self.k_norm(self.wk(hidden_states)) # [B, S, D] k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1) - k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, interleaved=rope_interleave, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D] + k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D] k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] # === Key cache (managed by the indexer, not DynamicCache) === @@ -533,7 +511,7 @@ def forward( query_states = query_states.view(batch_size, seq_length, self.num_heads, self.qk_head_dim).transpose(1, 2) # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D] q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] @@ -542,7 +520,7 @@ def forward( # RoPE on k_pe (single-head rope stream) k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] - k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # interleaved=True, BHSD format + k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) From e84b43c96d3b6f8f827457f97eb909c9ce7ad54e Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 22:40:31 +0800 Subject: [PATCH 06/16] reintroduce attention interface Signed-off-by: JaredforReal --- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 197 ++++++++---------- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 154 ++++++-------- 2 files changed, 153 insertions(+), 198 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index dc1b72b600a9..50da9d70abe9 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -37,10 +37,10 @@ from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update -from ...modeling_utils import PreTrainedModel +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_grouped_mm_available -from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_glm_moe_dsa import GlmMoeDsaConfig @@ -227,6 +227,43 @@ def forward( return topk_indices +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + scaling: float, + dropout: float = 0.0, + **kwargs: Unpack[TransformersKwargs], +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + def yarn_get_mscale(scale=1, mscale=1): if scale <= 1: return 1.0 @@ -320,7 +357,6 @@ def forward( ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: batch_size, seq_length = hidden_states.shape[:-1] cos, sin = position_embeddings - is_prefill = seq_length > 1 # ===== Query path ===== if self.q_lora_rank is None: @@ -339,9 +375,26 @@ def forward( k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] + # Expand KV through kv_b_proj + kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] + kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] + value_states = value_states.transpose(1, 2) # [B, H, S, v_D] + # RoPE on k_pe (single-head rope stream) k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format + k_pe = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] + + # Assemble full Q and K + query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] + key_states = torch.cat([k_nope, k_pe], dim=-1) # [B, H, S, qk_head_dim] + + # Cache update + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) @@ -354,113 +407,47 @@ def forward( use_cache=past_key_values is not None, ) # [B, S, topk] - if is_prefill: - # ===== Prefill: expand KV, full MHA ===== - kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] - kv_expanded = kv_expanded.view( - batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) - k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] - value_states = value_states.transpose(1, 2) # [B, H, S, v_D] - - k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] - - query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] - key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) # [B, H, S, qk_head_dim] - - # Cache update (store expanded K/V) - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update( - key_states, value_states, self.layer_idx, cache_kwargs - ) - - # Attention scores - total_len = key_states.shape[2] - attn_weights = ( - torch.einsum("bhsd,bhtd->bhst", query_states.float(), key_states.float()) * self.scaling - ) # [B, H, S, T] - - # Build index_mask: -inf everywhere except selected top-k positions - index_mask = torch.full( - (batch_size, seq_length, total_len), - float("-inf"), - device=hidden_states.device, - dtype=attn_weights.dtype, - ) - index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, :total_len] - index_mask = index_mask.unsqueeze(1) + causal_mask - else: - index_mask = index_mask.unsqueeze(1) - - attn_weights = attn_weights + index_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout( - attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training - ) + # Build combined DSA + causal mask: -inf everywhere except selected top-k positions + total_len = key_states.shape[2] + index_mask = torch.full( + (batch_size, seq_length, total_len), + float("-inf"), + device=hidden_states.device, + dtype=query_states.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] + index_mask = index_mask.unsqueeze(1) # [B, 1, S, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + combined_mask = index_mask + causal_mask + else: + combined_mask = index_mask - attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) # [B, H, S, v_D] + # Flash attention head_dim padding (qk_head_dim != v_head_dim) + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - else: - # ===== Decode: absorbed MLA (Q absorbs wkv_b, avoid K expansion) ===== - # Reference decode path: - # wkv_b = wkv_b.view(H, nope_D+v_D, kv_rank) - # q_nope' = einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :nope_D]) # absorb into Q - # scores = (einsum("bshc,btc->bsht", q_nope', kv_cache) + - # einsum("bshr,btr->bsht", q_pe, pe_cache)) * scale - # output = einsum("bsht,btc->bshc", attn_weights, kv_cache) - # output = einsum("bshc,hdc->bshd", output, wkv_b[:, -v_D:]) - - # Cache compressed KV and k_pe (before expansion) - # For decode, we cache compressed representations and use absorbed attention - kv_expanded = self.kv_b_proj(k_compressed) # still need for cache compatibility - kv_expanded = kv_expanded.view( - batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim - ) - k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_nope = k_nope.transpose(1, 2) - value_states = value_states.transpose(1, 2) - - k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) - key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) - - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update( - key_states, value_states, self.layer_idx, cache_kwargs - ) - - query_states = torch.cat([q_nope, q_pe], dim=-1) - total_len = key_states.shape[2] - attn_weights = torch.einsum("bhsd,bhtd->bhst", query_states.float(), key_states.float()) * self.scaling - - # Build index_mask - index_mask = torch.full( - (batch_size, 1, total_len), - float("-inf"), - device=hidden_states.device, - dtype=attn_weights.dtype, - ) - index_mask.scatter_(-1, topk_indices, 0.0) # [B, 1, T] - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, :total_len] - index_mask = index_mask.unsqueeze(1) + causal_mask - else: - index_mask = index_mask.unsqueeze(1) + # ===== Attention via standard interface ===== + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) - attn_weights = attn_weights + index_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout( - attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training - ) + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + combined_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) - attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] # ===== Output projection ===== - attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_length, -1).contiguous() + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights @@ -664,10 +651,10 @@ class GlmMoeDsaPreTrainedModel(PreTrainedModel): "attentions": GlmMoeDsaAttention, } _keep_in_fp32_modules_strict = ["e_score_correction_bias"] + _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"] # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert. # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`). _keep_in_fp32_modules = ["indexer.weights_proj"] - _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"] @torch.no_grad() def _init_weights(self, module): diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 4d9c96b881e8..7f8161d8adf2 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -14,16 +14,21 @@ +from collections.abc import Callable + import torch import torch.nn as nn +import torch.nn.functional as F from ...cache_utils import Cache from ...configuration_utils import PreTrainedConfig, layer_type_validation from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_rope_utils import RopeParameters +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS from ...models.llama.modeling_llama import rotate_half from ...processing_utils import Unpack from ...utils import logging +from ...utils.generic import is_flash_attention_requested from ..deepseek_v3.modeling_deepseek_v3 import yarn_get_mscale from ..glm4_moe.modeling_glm4_moe import ( Glm4MoeForCausalLM, @@ -31,7 +36,10 @@ Glm4MoePreTrainedModel, Glm4MoeRMSNorm, ) -from ..glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteDecoderLayer +from ..glm4_moe_lite.modeling_glm4_moe_lite import ( + Glm4MoeLiteDecoderLayer, + eager_attention_forward, +) logger = logging.get_logger(__name__) @@ -499,7 +507,6 @@ def forward( ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: batch_size, seq_length = hidden_states.shape[:-1] cos, sin = position_embeddings - is_prefill = seq_length > 1 # ===== Query path ===== if self.q_lora_rank is None: @@ -518,9 +525,28 @@ def forward( k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] + # Expand KV through kv_b_proj + kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] + kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] + value_states = value_states.transpose(1, 2) # [B, H, S, v_D] + # RoPE on k_pe (single-head rope stream) k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D] k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format + k_pe = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] + + # Assemble full Q and K + query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] + key_states = torch.cat([k_nope, k_pe], dim=-1) # [B, H, S, qk_head_dim] + + # Cache update + if past_key_values is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) but indexer works with [B, S, T] (3D) @@ -533,103 +559,45 @@ def forward( use_cache=past_key_values is not None, ) # [B, S, topk] - if is_prefill: - # ===== Prefill: expand KV, full MHA ===== - kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)] - kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D] - value_states = value_states.transpose(1, 2) # [B, H, S, v_D] - - k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) # [B, H, S, rope_D] - - query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim] - key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) # [B, H, S, qk_head_dim] - - # Cache update (store expanded K/V) - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update( - key_states, value_states, self.layer_idx, cache_kwargs - ) - - # Attention scores - total_len = key_states.shape[2] - attn_weights = torch.einsum( - "bhsd,bhtd->bhst", query_states.float(), key_states.float() - ) * self.scaling # [B, H, S, T] - - # Build index_mask: -inf everywhere except selected top-k positions - index_mask = torch.full( - (batch_size, seq_length, total_len), float("-inf"), - device=hidden_states.device, dtype=attn_weights.dtype, - ) - index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, :total_len] - index_mask = index_mask.unsqueeze(1) + causal_mask - else: - index_mask = index_mask.unsqueeze(1) - - attn_weights = attn_weights + index_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training) + # Build combined DSA + causal mask: -inf everywhere except selected top-k positions + total_len = key_states.shape[2] + index_mask = torch.full( + (batch_size, seq_length, total_len), float("-inf"), + device=hidden_states.device, dtype=query_states.dtype, + ) + index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] + index_mask = index_mask.unsqueeze(1) # [B, 1, S, T] + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, :total_len] + combined_mask = index_mask + causal_mask + else: + combined_mask = index_mask - attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) # [B, H, S, v_D] + # Flash attention head_dim padding (qk_head_dim != v_head_dim) + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - else: - # ===== Decode: absorbed MLA (Q absorbs wkv_b, avoid K expansion) ===== - # Reference decode path: - # wkv_b = wkv_b.view(H, nope_D+v_D, kv_rank) - # q_nope' = einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :nope_D]) # absorb into Q - # scores = (einsum("bshc,btc->bsht", q_nope', kv_cache) + - # einsum("bshr,btr->bsht", q_pe, pe_cache)) * scale - # output = einsum("bsht,btc->bshc", attn_weights, kv_cache) - # output = einsum("bshc,hdc->bshd", output, wkv_b[:, -v_D:]) - - # Cache compressed KV and k_pe (before expansion) - # For decode, we cache compressed representations and use absorbed attention - kv_expanded = self.kv_b_proj(k_compressed) # still need for cache compatibility - kv_expanded = kv_expanded.view(batch_size, seq_length, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_nope = k_nope.transpose(1, 2) - value_states = value_states.transpose(1, 2) - - k_pe_expanded = k_pe.expand(-1, self.num_heads, -1, -1) - key_states = torch.cat([k_nope, k_pe_expanded], dim=-1) - - if past_key_values is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_values.update( - key_states, value_states, self.layer_idx, cache_kwargs - ) - - query_states = torch.cat([q_nope, q_pe], dim=-1) - total_len = key_states.shape[2] - attn_weights = torch.einsum( - "bhsd,bhtd->bhst", query_states.float(), key_states.float() - ) * self.scaling - - # Build index_mask - index_mask = torch.full( - (batch_size, 1, total_len), float("-inf"), - device=hidden_states.device, dtype=attn_weights.dtype, - ) - index_mask.scatter_(-1, topk_indices, 0.0) # [B, 1, T] - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, :total_len] - index_mask = index_mask.unsqueeze(1) + causal_mask - else: - index_mask = index_mask.unsqueeze(1) + # ===== Attention via standard interface ===== + attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) - attn_weights = attn_weights + index_mask - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout if self.training else 0.0, training=self.training) + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + combined_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) - attn_output = torch.einsum("bhst,bhtd->bhsd", attn_weights, value_states) + if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: + attn_output = attn_output[:, :, :, : self.v_head_dim] # ===== Output projection ===== - attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_length, -1).contiguous() + attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights From 558989a06ff9bac68ba7d9031d641d27dafcd747 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Wed, 11 Feb 2026 22:49:57 +0800 Subject: [PATCH 07/16] reset _cached_keys Signed-off-by: JaredforReal --- src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py | 4 ++++ src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 50da9d70abe9..f69ee2f15222 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -190,6 +190,10 @@ def forward( k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] # === Key cache (managed by the indexer, not DynamicCache) === + # Reset cache on prefill (new prompt) to avoid stale keys / batch-size mismatch + if seq_len > 1: + self._cached_keys = None + if use_cache: if self._cached_keys is not None: k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D] diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 7f8161d8adf2..9a319db78999 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -383,6 +383,10 @@ def forward( k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D] # === Key cache (managed by the indexer, not DynamicCache) === + # Reset cache on prefill (new prompt) to avoid stale keys / batch-size mismatch + if seq_len > 1: + self._cached_keys = None + if use_cache: if self._cached_keys is not None: k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D] From 62454529d37368a0f28ba4290c5958bd22f61081 Mon Sep 17 00:00:00 2001 From: JaredforReal Date: Fri, 13 Feb 2026 23:35:02 +0800 Subject: [PATCH 08/16] remove yarn Signed-off-by: JaredforReal --- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 14 -------------- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 8 -------- 2 files changed, 22 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index f69ee2f15222..462dcd16d756 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -19,7 +19,6 @@ # limitations under the License. -import math from collections.abc import Callable from typing import Optional @@ -268,12 +267,6 @@ def eager_attention_forward( return attn_output, attn_weights -def yarn_get_mscale(scale=1, mscale=1): - if scale <= 1: - return 1.0 - return 0.1 * mscale * math.log(scale) + 1.0 - - class GlmMoeDsaAttention(nn.Module): """ Multi-head Latent Attention (MLA) with Dynamic Sparse Attention (DSA) indexer. @@ -340,13 +333,6 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): ) self.scaling = self.qk_head_dim ** (-0.5) - rope_params = self.config.rope_parameters or {} - if rope_params.get("rope_type", "default") != "default": - mscale_all_dim = rope_params.get("mscale_all_dim", 0) - scaling_factor = rope_params["factor"] - if mscale_all_dim: - mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) - self.scaling = self.scaling * mscale * mscale self.indexer = GlmMoeDsaIndexer(config, layer_idx) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 9a319db78999..3bb589576d2b 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -29,7 +29,6 @@ from ...processing_utils import Unpack from ...utils import logging from ...utils.generic import is_flash_attention_requested -from ..deepseek_v3.modeling_deepseek_v3 import yarn_get_mscale from ..glm4_moe.modeling_glm4_moe import ( Glm4MoeForCausalLM, Glm4MoeModel, @@ -490,13 +489,6 @@ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int): ) self.scaling = self.qk_head_dim ** (-0.5) - rope_params = self.config.rope_parameters or {} - if rope_params.get("rope_type", "default") != "default": - mscale_all_dim = rope_params.get("mscale_all_dim", 0) - scaling_factor = rope_params["factor"] - if mscale_all_dim: - mscale = yarn_get_mscale(scaling_factor, mscale_all_dim) - self.scaling = self.scaling * mscale * mscale self.indexer = GlmMoeDsaIndexer(config, layer_idx) From 299c53c33911aa1971c26a5d1812740516c8c6d4 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Mon, 16 Feb 2026 14:28:27 +0000 Subject: [PATCH 09/16] fix tp plan for multi node runs --- .../models/glm_moe_dsa/configuration_glm_moe_dsa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 4e9c9032aa93..23d99803368b 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -124,7 +124,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): model_type = "glm_moe_dsa" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_split_input", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", From 125f99495a105f62006cf92d21c84faa8799b13f Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 16 Feb 2026 18:00:39 +0100 Subject: [PATCH 10/16] tentatively add flash-mla --- src/transformers/integrations/__init__.py | 2 + src/transformers/integrations/flash_mla.py | 237 ++++++++++++++++++ src/transformers/integrations/hub_kernels.py | 7 + .../glm_moe_dsa/modeling_glm_moe_dsa.py | 9 - .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 19 +- 5 files changed, 255 insertions(+), 19 deletions(-) create mode 100644 src/transformers/integrations/flash_mla.py diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index a4dfe3092a8e..84bb038167f6 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -66,6 +66,7 @@ "replace_with_higgs_linear", ], "hqq": ["prepare_for_hqq_linear"], + "flash_mla": ["flash_mla_forward"], "hub_kernels": [ "LayerRepository", "lazy_load_kernel", @@ -213,6 +214,7 @@ ) from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear from .hqq import prepare_for_hqq_linear + from .flash_mla import flash_mla_forward from .hub_kernels import ( LayerRepository, lazy_load_kernel, diff --git a/src/transformers/integrations/flash_mla.py b/src/transformers/integrations/flash_mla.py new file mode 100644 index 000000000000..0890a9c1d4ef --- /dev/null +++ b/src/transformers/integrations/flash_mla.py @@ -0,0 +1,237 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Flash-MLA attention integration for sparse attention with Dynamic Sparse Attention (DSA). + +This module provides a wrapper around the flash-mla kernel from kernels-community/flash-mla, +with automatic fallback to flash_attention_2 when input tokens < 2048. +""" + +import torch + +from ..utils import logging +from .flash_attention import flash_attention_forward, get_target_dtype + + +logger = logging.get_logger(__name__) + +# Minimum sequence length to use flash-mla sparse attention +# Below this threshold, we fall back to flash_attention_2 +FLASH_MLA_MIN_SEQ_LEN = 2048 + + +def flash_mla_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + dropout: float = 0.0, + scaling: float | None = None, + sliding_window: int | None = None, + softcap: float | None = None, + is_causal: bool | None = None, + **kwargs, +) -> tuple[torch.Tensor, None]: + """ + Flash-MLA attention forward pass with automatic fallback to flash_attention_2. + + This wrapper handles: + - Fallback to flash_attention_2 when sequence length < 2048 + - Sparse attention via topk_indices when sequence length >= 2048 + - Tensor layout conversion (transformers BHSD to flash-mla BSHD) + - head_dim_v padding if needed (flash-mla requires specific dimensions) + + Args: + module (`torch.nn.Module`): + The attention module containing config and layer information. + query (`torch.Tensor`): + Query tensor of shape `[B, H, S, D]` (BHSD format). + key (`torch.Tensor`): + Key tensor of shape `[B, H, T, D]` (BHSD format). + value (`torch.Tensor`): + Value tensor of shape `[B, H, T, D_v]` (BHSD format). + attention_mask (`torch.Tensor | None`): + Combined attention mask (causal + DSA sparse mask). Used for flash_attention_2 fallback. + dropout (`float`, optional): + Dropout probability. Defaults to 0.0. + scaling (`float | None`, optional): + Scaling factor for attention scores. Defaults to None. + sliding_window (`int | None`, optional): + Sliding window size. Defaults to None. + softcap (`float | None`, optional): + Soft cap for attention logits. Defaults to None. + is_causal (`bool | None`, optional): + Whether attention is causal. Defaults to None. + **kwargs: + Additional keyword arguments, including: + - topk_indices (`torch.Tensor | None`): Indices for sparse attention from DSA indexer. + + Returns: + `tuple[torch.Tensor, None]`: Attention output tensor and None (no attention weights). + """ + # Extract topk_indices from kwargs (used for sparse attention) + topk_indices = kwargs.pop("topk_indices", None) + + # Get total sequence length from key tensor + # key shape is [B, H, T, D] in BHSD format + seq_len = key.shape[2] + + # Fallback to flash_attention_2 when sequence length is below threshold + # This is because flash-mla sparse attention is optimized for longer sequences + if seq_len < FLASH_MLA_MIN_SEQ_LEN: + logger.debug_once( + f"Sequence length {seq_len} < {FLASH_MLA_MIN_SEQ_LEN}, falling back to flash_attention_2" + ) + return flash_attention_forward( + module=module, + query=query, + key=key, + value=value, + attention_mask=attention_mask, + dropout=dropout, + scaling=scaling, + sliding_window=sliding_window, + softcap=softcap, + is_causal=is_causal, + **kwargs, + ) + + # Use flash-mla sparse attention with topk_indices + return _flash_mla_sparse_forward( + module=module, + query=query, + key=key, + value=value, + topk_indices=topk_indices, + dropout=dropout, + scaling=scaling, + **kwargs, + ) + + +def _flash_mla_sparse_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + topk_indices: torch.Tensor | None, + dropout: float = 0.0, + scaling: float | None = None, + **kwargs, +) -> tuple[torch.Tensor, None]: + """ + Internal function for flash-mla sparse attention computation. + + This function handles the actual flash-mla kernel call with sparse attention + via topk_indices from the DSA indexer. + + Args: + module (`torch.nn.Module`): + The attention module containing config and layer information. + query (`torch.Tensor`): + Query tensor of shape `[B, H, S, D]` (BHSD format). + key (`torch.Tensor`): + Key tensor of shape `[B, H, T, D]` (BHSD format). + value (`torch.Tensor`): + Value tensor of shape `[B, H, T, D_v]` (BHSD format). + topk_indices (`torch.Tensor | None`): + Indices for sparse attention from DSA indexer, shape `[B, S, topk]`. + dropout (`float`, optional): + Dropout probability. Defaults to 0.0. + scaling (`float | None`, optional): + Scaling factor for attention scores. Defaults to None. + **kwargs: + Additional keyword arguments. + + Returns: + `tuple[torch.Tensor, None]`: Attention output tensor and None (no attention weights). + """ + if kwargs.get("output_attentions", False): + logger.warning_once( + "Flash-MLA does not support `output_attentions=True`. " + "Please set your attention to `eager` if you want this feature." + ) + + # Get batch size and sequence lengths + batch_size, num_heads, q_len, head_dim = query.shape + _, _, kv_len, _ = key.shape + + # Convert from BHSD (transformers) to BSHD (flash-mla) format + # query: [B, H, S, D] -> [B, S, H, D] + # key: [B, H, T, D] -> [B, T, H, D] + # value: [B, H, T, D_v] -> [B, T, H, D_v] + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + + # Handle dtype conversion for flash attention compatibility + target_dtype = get_target_dtype(query, module) + if target_dtype is not None: + query = query.to(target_dtype) + key = key.to(target_dtype) + value = value.to(target_dtype) + + # Get the flash-mla kernel function + # This is loaded via hub_kernels infrastructure + try: + from ..integrations.hub_kernels import get_kernel + + flash_mla_kernel = get_kernel("kernels-community/flash-mla") + flash_mla_sparse_fwd = flash_mla_kernel.flash_mla_sparse_fwd + except (ImportError, AttributeError) as e: + raise RuntimeError( + f"Failed to load flash-mla kernel. Make sure kernels-community/flash-mla is available. Error: {e}" + ) + + # Prepare scaling factor + if scaling is None: + scaling = head_dim**-0.5 + + # Get value head dimension (may differ from query/key head dimension in MLA) + v_head_dim = value.shape[-1] + + # Flash-MLA may require specific head_dim_v (e.g., 512) + # Pad if necessary + flash_mla_v_head_dim = 512 + needs_v_padding = v_head_dim < flash_mla_v_head_dim + if needs_v_padding: + value = torch.nn.functional.pad(value, (0, flash_mla_v_head_dim - v_head_dim)) + + # Call flash-mla kernel with sparse attention + # The kernel expects: + # - q: [B, S, H, D] + # - k_cache: [B, T, H, D] (or compressed format) + # - v_cache: [B, T, H, D_v] + # - topk_indices: [B, S, topk] for sparse attention + attn_output = flash_mla_sparse_fwd( + q=query, + kv=torch.cat([key, value], dim=-1), + indices=topk_indices, + sm_scale=scaling, + topk_length = module.top_k_length if hasattr(module, "top_k_length") else None, + ) + + # Remove padding if we added it + if needs_v_padding: + attn_output = attn_output[..., :v_head_dim] + + # Convert back from BSHD to BHSD format + # attn_output: [B, S, H, D_v] -> [B, H, S, D_v] + attn_output = attn_output.transpose(1, 2) + + return attn_output, None + + +__all__ = ["flash_mla_forward"] diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index a189dac2b125..5627aa69d57e 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -277,6 +277,7 @@ def register_kernel_mapping_transformers(*args, **kwargs): "causal-conv1d": {"repo_id": "kernels-community/causal-conv1d"}, "mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "revision": "v0.0.4"}, "falcon_mamba-ssm": {"repo_id": "kernels-community/mamba-ssm", "revision": "v0.0.4"}, + "flash-mla": {"repo_id": "kernels-community/flash-mla"}, } _KERNEL_MODULE_MAPPING: dict[str, ModuleType | None] = {} @@ -338,6 +339,12 @@ def load_and_register_attn_kernel( if attention_wrapper is None: attention_wrapper = flash_attention_forward kernel_function = attention_wrapper + elif hasattr(kernel, "flash_mla_sparse_fwd"): + from .flash_mla import flash_mla_forward + + if attention_wrapper is None: + attention_wrapper = flash_mla_forward + kernel_function = attention_wrapper elif kernel_name is not None: kernel_function = getattr(kernel, kernel_name) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 462dcd16d756..542fcfb6dc09 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -413,11 +413,6 @@ def forward( else: combined_mask = index_mask - # Flash attention head_dim padding (qk_head_dim != v_head_dim) - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - - # ===== Attention via standard interface ===== attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) @@ -433,10 +428,6 @@ def forward( **kwargs, ) - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - attn_output = attn_output[:, :, :, : self.v_head_dim] - - # ===== Output projection ===== attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 3bb589576d2b..a233650124f0 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -28,7 +28,6 @@ from ...models.llama.modeling_llama import rotate_half from ...processing_utils import Unpack from ...utils import logging -from ...utils.generic import is_flash_attention_requested from ..glm4_moe.modeling_glm4_moe import ( Glm4MoeForCausalLM, Glm4MoeModel, @@ -284,6 +283,14 @@ def __init__( self.rope_parameters = rope_parameters self.rope_interleave = rope_interleave + # Warn if using flash_attention_2 instead of flash-mla + if kwargs.get("attn_implementation") == "flash_attention_2": + logger.warning_once( + "The glm_moe_dsa model is optimized for 'kernels-community/flash-mla'. " + "Using 'flash_attention_2' may not fully utilize DSA sparse attention. " + "Consider using attn_implementation='kernels-community/flash-mla'." + ) + super().__init__( pad_token_id=pad_token_id, bos_token_id=bos_token_id, @@ -569,11 +576,6 @@ def forward( else: combined_mask = index_mask - # Flash attention head_dim padding (qk_head_dim != v_head_dim) - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim]) - - # ===== Attention via standard interface ===== attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) @@ -586,13 +588,10 @@ def forward( combined_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, + topk_indices=topk_indices, # Pass topk_indices for flash-mla sparse attention **kwargs, ) - if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim: - attn_output = attn_output[:, :, :, : self.v_head_dim] - - # ===== Output projection ===== attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous() attn_output = self.o_proj(attn_output) return attn_output, attn_weights From 4f89abe55c96ad6485b2fd3bb3642a54f72a233d Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 16 Feb 2026 18:31:04 +0100 Subject: [PATCH 11/16] skip more tests --- .../glm_moe_dsa/test_modeling_glm_moe_dsa.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py index de9c3f53dd02..2094c15963ce 100644 --- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py +++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py @@ -62,6 +62,26 @@ class GlmMoeDsaModelTest(CausalLMModelTest, unittest.TestCase): test_all_params_have_gradient = False model_split_percents = [0.5, 0.7, 0.8] + @unittest.skip("DSA indexer mask shape mismatch with assisted decoding") + def test_assisted_decoding_matches_greedy_search(self): + pass + + @unittest.skip("DSA indexer mask shape mismatch with assisted decoding") + def test_assisted_decoding_sample(self): + pass + + @unittest.skip("Requires torch>=2.9.0 for grouped MM") + def test_eager_matches_batched_and_grouped_inference(self): + pass + + @unittest.skip("FP32 module detection needs adjustment for DSA indexer weights") + def test_keep_in_fp32_modules(self): + pass + + @unittest.skip("FP32 module detection needs adjustment for DSA indexer weights") + def test_keep_in_fp32_modules_strict(self): + pass + def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_length, config): """Needs to be overridden as GLM-4.7-Flash has special MLA cache format (though we don't really use the MLA)""" self.assertIsInstance(past_key_values, Cache) From 9e7a2e19a4b26afba2d90cac8f9c42533af6c9ea Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 17 Feb 2026 10:35:17 +0100 Subject: [PATCH 12/16] fuck the tp plan is wrong --- .../glm_moe_dsa/configuration_glm_moe_dsa.py | 2 +- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 14 +++++++++----- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 2 +- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 387d2a7ccd94..1cc7ff40312b 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -127,7 +127,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): model_type = "glm_moe_dsa" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.o_proj": "rowwise_split_input", + "layers.*.self_attn.o_proj": "rowwise", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 7dee684042a6..d8804a256951 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -39,7 +39,7 @@ from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_grouped_mm_available -from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults +from ...utils.generic import maybe_autocast, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_glm_moe_dsa import GlmMoeDsaConfig @@ -413,11 +413,15 @@ def forward( ) index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T] index_mask = index_mask.unsqueeze(1) # [B, 1, S, T] - if attention_mask is not None: - causal_mask = attention_mask[:, :, :, :total_len] + if attention_mask is not None and attention_mask.dim() == 4: + causal_mask = attention_mask[..., :total_len] combined_mask = index_mask + causal_mask else: - combined_mask = index_mask + combined_mask = ( + attention_mask.masked_fill(index_mask == float("-inf"), float("-inf")) + if attention_mask is not None + else index_mask + ) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -431,7 +435,7 @@ def forward( combined_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, - indices=topk_indices, # flash_mla_with_kvcache + topk_indices=topk_indices, # Pass topk_indices for flash-mla sparse attention **kwargs, ) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 785fc1977962..7c85ad187b6b 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -178,7 +178,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): model_type = "glm_moe_dsa" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_split_input", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", From 1dff69e31f2d27a26d01d0b572e5a44e1af1abd7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 17 Feb 2026 10:45:10 +0100 Subject: [PATCH 13/16] small fixes --- .../models/glm_moe_dsa/configuration_glm_moe_dsa.py | 2 +- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 8 +------- .../models/glm_moe_dsa/modular_glm_moe_dsa.py | 12 +----------- 3 files changed, 3 insertions(+), 19 deletions(-) diff --git a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py index 1cc7ff40312b..387d2a7ccd94 100644 --- a/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py @@ -127,7 +127,7 @@ class GlmMoeDsaConfig(PreTrainedConfig): model_type = "glm_moe_dsa" keys_to_ignore_at_inference = ["past_key_values"] base_model_tp_plan = { - "layers.*.self_attn.o_proj": "rowwise", + "layers.*.self_attn.o_proj": "rowwise_split_input", "layers.*.mlp.experts.gate_up_proj": "packed_colwise", "layers.*.mlp.experts.down_proj": "rowwise", "layers.*.mlp.experts": "moe_tp_experts", diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index d8804a256951..79d0dd727203 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -344,11 +344,10 @@ def forward( past_key_values: Cache | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: batch_size, seq_length = hidden_states.shape[:-1] cos, sin = position_embeddings - # ===== Query path ===== if self.q_lora_rank is None: query_states = self.q_proj(hidden_states) q_resid = None @@ -360,7 +359,6 @@ def forward( q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format - # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] @@ -386,7 +384,6 @@ def forward( cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) - # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D) indexer_mask = ( attention_mask[:, 0, :, :] @@ -650,9 +647,6 @@ class GlmMoeDsaPreTrainedModel(PreTrainedModel): @torch.no_grad() def _init_weights(self, module): - # Skip normal_ initialization for FP8 quantized weights which don't support it - if isinstance(module, nn.Linear) and hasattr(module, "weight") and module.weight.dtype == torch.float8_e4m3fn: - return super()._init_weights(module) if isinstance(module, GlmMoeDsaTopkRouter): init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) diff --git a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py index 7c85ad187b6b..959b6298722c 100644 --- a/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py @@ -499,11 +499,10 @@ def forward( past_key_values: Cache | None = None, cache_position: torch.LongTensor | None = None, **kwargs: Unpack[FlashAttentionKwargs], - ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + ) -> tuple[torch.Tensor, torch.Tensor | None]: batch_size, seq_length = hidden_states.shape[:-1] cos, sin = position_embeddings - # ===== Query path ===== if self.q_lora_rank is None: query_states = self.q_proj(hidden_states) q_resid = None @@ -515,7 +514,6 @@ def forward( q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format - # ===== KV path ===== compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D] k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank] @@ -543,7 +541,6 @@ def forward( key_states, value_states, self.layer_idx, cache_kwargs ) - # ===== Indexer (DSA sparse mask) ===== # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D) indexer_mask = ( attention_mask[:, 0, :, :] @@ -614,13 +611,6 @@ class GlmMoeDsaPreTrainedModel(Glm4MoePreTrainedModel): _supports_sdpa = True _supports_flex_attn = False - @torch.no_grad() - def _init_weights(self, module): - # Skip normal_ initialization for FP8 quantized weights which don't support it - if isinstance(module, nn.Linear) and hasattr(module, "weight") and module.weight.dtype == torch.float8_e4m3fn: - return - super()._init_weights(module) - class GlmMoeDsaModel(Glm4MoeModel): pass From 7b3ba9a54a3581c889baa80085eaf45d8c59ebe5 Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Tue, 17 Feb 2026 13:50:15 +0000 Subject: [PATCH 14/16] yup --- .../integrations/finegrained_fp8.py | 98 +++++++++---------- .../glm_moe_dsa/modeling_glm_moe_dsa.py | 2 +- 2 files changed, 50 insertions(+), 50 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index fec2a31f7eb3..443b584fdf3a 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -412,55 +412,55 @@ def w8a8_block_fp8_matmul( Otherwise falls back to Triton. """ - if _supports_cutlass(block_size, output_dtype): - kernel = _get_quantization_kernel() - if kernel is not None: - try: - # CUTLASS expects: - # - A: [M, K] row-major, float8_e4m3fn - # - B: [K, N] column-major, float8_e4m3fn - # - As: [M, K//128] M-major (activation scales) - # - Bs: [K//128, N//128] K-major (weight scales) - - # Reshape A to 2D if needed - original_shape = A.shape - M = A.numel() // A.shape[-1] - K = A.shape[-1] - N = B.shape[0] - - # CUTLASS requires dimensions divisible by 16 - if K % 16 != 0 or N % 16 != 0: - raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16") - - A_2d = A.view(M, K).contiguous() - # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1 - # Our B is [N, K] row-major. Make it contiguous first, then transpose. - # B.contiguous() gives [N, K] with stride=(K,1) - # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major - # Do NOT call .contiguous() after .t() as it would make it row-major! - B_col_major = B.contiguous().t() - - # Scales need proper layout for CUTLASS blockwise: - # As should be [M, K//128] with M-major layout (stride(0)==1) - # Bs should be [K//128, N//128] with K-major layout (stride(0)==1) - - # As: reshape to [M, K//128], then make M-major via t().contiguous().t() - As_2d = As.view(M, -1).contiguous() - As_2d = As_2d.t().contiguous().t() # [M, K//128] with stride(0)==1 - - # Bs: our input is [N//128, K//128], need [K//128, N//128] with stride(0)==1 - # Transpose to get [K//128, N//128], then make K-major via t().contiguous().t() - Bs_km = Bs.contiguous().t() # [K//128, N//128] - Bs_km = Bs_km.t().contiguous().t() # Make K-major (stride(0)==1) - - # Call CUTLASS kernel - it returns the output tensor - # Signature: cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias=None) -> Tensor - C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None) - # Reshape output back - C_shape = original_shape[:-1] + (N,) - return C.view(C_shape) - except Exception as e: - logger.warning_once(f"CUTLASS kernel failed: {e}. Falling back to Triton.") + # if _supports_cutlass(block_size, output_dtype): + # kernel = _get_quantization_kernel() + # if kernel is not None: + # try: + # # CUTLASS expects: + # # - A: [M, K] row-major, float8_e4m3fn + # # - B: [K, N] column-major, float8_e4m3fn + # # - As: [M, K//128] M-major (activation scales) + # # - Bs: [K//128, N//128] K-major (weight scales) + + # # Reshape A to 2D if needed + # original_shape = A.shape + # M = A.numel() // A.shape[-1] + # K = A.shape[-1] + # N = B.shape[0] + + # # CUTLASS requires dimensions divisible by 16 + # if K % 16 != 0 or N % 16 != 0: + # raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16") + + # A_2d = A.view(M, K).contiguous() + # # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1 + # # Our B is [N, K] row-major. Make it contiguous first, then transpose. + # # B.contiguous() gives [N, K] with stride=(K,1) + # # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major + # # Do NOT call .contiguous() after .t() as it would make it row-major! + # B_col_major = B.contiguous().t() + + # # Scales need proper layout for CUTLASS blockwise: + # # As should be [M, K//128] with M-major layout (stride(0)==1) + # # Bs should be [K//128, N//128] with K-major layout (stride(0)==1) + + # # As: reshape to [M, K//128], then make M-major via t().contiguous().t() + # As_2d = As.view(M, -1).contiguous() + # As_2d = As_2d.t().contiguous().t() # [M, K//128] with stride(0)==1 + + # # Bs: our input is [N//128, K//128], need [K//128, N//128] with stride(0)==1 + # # Transpose to get [K//128, N//128], then make K-major via t().contiguous().t() + # Bs_km = Bs.contiguous().t() # [K//128, N//128] + # Bs_km = Bs_km.t().contiguous().t() # Make K-major (stride(0)==1) + + # # Call CUTLASS kernel - it returns the output tensor + # # Signature: cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias=None) -> Tensor + # C = kernel.cutlass_scaled_mm(A_2d, B_col_major, As_2d, Bs_km, output_dtype, None) + # # Reshape output back + # C_shape = original_shape[:-1] + (N,) + # return C.view(C_shape) + # except Exception as e: + # logger.warning_once(f"CUTLASS kernel failed: {e}. Falling back to Triton.") # Fall back to Triton return w8a8_block_fp8_matmul_triton(A, B, As, Bs, block_size, output_dtype) diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 79d0dd727203..ef0df4f4604b 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -628,7 +628,7 @@ class GlmMoeDsaPreTrainedModel(PreTrainedModel): supports_gradient_checkpointing = True _no_split_modules = ["GlmMoeDsaDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] - _supports_flash_attn = False # flash-mla kernels need a bit more work in the way we enable them! + _supports_flash_attn = True # flash-mla kernels need a bit more work in the way we enable them! _supports_sdpa = True _supports_flex_attn = False _can_compile_fullgraph = ( From 309c8fd756580cdf2a0c3861f9286cfbd605d67f Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Wed, 18 Feb 2026 10:26:15 +0000 Subject: [PATCH 15/16] current changes --- src/transformers/integrations/__init__.py | 4 ++-- src/transformers/integrations/flash_mla.py | 4 ++-- src/transformers/integrations/hub_kernels.py | 6 +++--- src/transformers/modeling_utils.py | 6 +++++- .../models/glm_moe_dsa/modeling_glm_moe_dsa.py | 13 +++++++------ 5 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 4eb2d66a7b1e..436e324f0a27 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -66,7 +66,7 @@ "replace_with_higgs_linear", ], "hqq": ["prepare_for_hqq_linear"], - "flash_mla": ["flash_mla_forward"], + "flash_mla": ["flash_mla_attention_forward"], "hub_kernels": [ "LayerRepository", "lazy_load_kernel", @@ -215,7 +215,7 @@ ) from .higgs import HiggsLinear, dequantize_higgs, quantize_with_higgs, replace_with_higgs_linear from .hqq import prepare_for_hqq_linear - from .flash_mla import flash_mla_forward + from .flash_mla import flash_mla_attention_forward from .hub_kernels import ( LayerRepository, lazy_load_kernel, diff --git a/src/transformers/integrations/flash_mla.py b/src/transformers/integrations/flash_mla.py index 0890a9c1d4ef..b0e0d13da7d3 100644 --- a/src/transformers/integrations/flash_mla.py +++ b/src/transformers/integrations/flash_mla.py @@ -31,7 +31,7 @@ FLASH_MLA_MIN_SEQ_LEN = 2048 -def flash_mla_forward( +def flash_mla_attention_forward( module: torch.nn.Module, query: torch.Tensor, key: torch.Tensor, @@ -234,4 +234,4 @@ def _flash_mla_sparse_forward( return attn_output, None -__all__ = ["flash_mla_forward"] +__all__ = ["flash_mla_attention_forward"] diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index f936787a4abb..62cdc94fbd40 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -339,11 +339,11 @@ def load_and_register_attn_kernel( if attention_wrapper is None: attention_wrapper = flash_attention_forward kernel_function = attention_wrapper - elif hasattr(kernel, "flash_mla_sparse_fwd"): - from .flash_mla import flash_mla_forward + if hasattr(kernel, "flash_mla_sparse_fwd"): + from .flash_mla import flash_mla_attention_forward if attention_wrapper is None: - attention_wrapper = flash_mla_forward + attention_wrapper = flash_mla_attention_forward kernel_function = attention_wrapper elif kernel_name is not None: kernel_function = getattr(kernel, kernel_name) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 5ab58d71eae2..26a2d232bcae 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -70,6 +70,7 @@ from .integrations.flash_attention import flash_attention_forward from .integrations.flash_paged import paged_attention_forward from .integrations.flex_attention import flex_attention_forward +from .integrations.flash_mla import flash_mla_attention_forward from .integrations.hub_kernels import is_kernel from .integrations.peft import maybe_load_adapters from .integrations.sdpa_attention import sdpa_attention_forward @@ -1831,10 +1832,12 @@ def _check_and_adjust_attn_implementation( if is_kernel(applicable_attn_implementation): try: # preload flash attention here to allow compile with fullgraph + # we need to load the attention kernel requested by the user. if is_paged: lazy_import_paged_flash_attention(applicable_attn_implementation) else: - lazy_import_flash_attention(applicable_attn_implementation) + attention_wrapper = ALL_ATTENTION_FUNCTIONS.get(applicable_attn_implementation.rsplit("|")[0]) if "|" in applicable_attn_implementation else None + lazy_import_flash_attention(applicable_attn_implementation, attention_wrapper=attention_wrapper) # log that we used kernel fallback if successful if requested_original_flash_attn: @@ -4791,6 +4794,7 @@ class AttentionInterface(GeneralInterface): "paged|flash_attention_2": paged_attention_forward, "paged|sdpa": sdpa_attention_paged_forward, "paged|eager": eager_paged_attention_forward, + "flash_mla": flash_mla_attention_forward, } def get_interface(self, attn_implementation: str, default: Callable) -> Callable: diff --git a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py index 2b0adc73d1fa..667b9cc656ed 100644 --- a/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py +++ b/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py @@ -414,11 +414,12 @@ def forward( causal_mask = attention_mask[..., :total_len] combined_mask = index_mask + causal_mask else: - combined_mask = ( - attention_mask.masked_fill(index_mask == float("-inf"), float("-inf")) - if attention_mask is not None - else index_mask - ) + combined_mask = None + # combined_mask = ( + # attention_mask.masked_fill(index_mask == float("-inf"), float("-inf")) + # if attention_mask is not None + # else index_mask + # ) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward @@ -644,7 +645,7 @@ class GlmMoeDsaPreTrainedModel(PreTrainedModel): # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert. # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`). _keep_in_fp32_modules = ["indexer.weights_proj"] - _default_flash_implementation = "kernels-community/flash-mla" + _default_flash_implementation = "flash_mla|kernels-community/flash-mla:flash_mla_sparse_fwd" @torch.no_grad() def _init_weights(self, module): From 1f40e8bb592599810d362a66ccf34555d23df89c Mon Sep 17 00:00:00 2001 From: Arthur Zucker Date: Wed, 18 Feb 2026 10:32:55 +0000 Subject: [PATCH 16/16] nit --- src/transformers/integrations/flash_mla.py | 2 +- src/transformers/quantizers/quantizers_utils.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/transformers/integrations/flash_mla.py b/src/transformers/integrations/flash_mla.py index b0e0d13da7d3..ecd80fd357fc 100644 --- a/src/transformers/integrations/flash_mla.py +++ b/src/transformers/integrations/flash_mla.py @@ -91,7 +91,7 @@ def flash_mla_attention_forward( # Fallback to flash_attention_2 when sequence length is below threshold # This is because flash-mla sparse attention is optimized for longer sequences if seq_len < FLASH_MLA_MIN_SEQ_LEN: - logger.debug_once( + logger.debug( f"Sequence length {seq_len} < {FLASH_MLA_MIN_SEQ_LEN}, falling back to flash_attention_2" ) return flash_attention_forward( diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 0e90e238ec4a..e0d006ce38fe 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -34,8 +34,9 @@ def should_convert_module(full_name, patterns: list[str] | None = None): # 3. `full_name` ends with the pattern # (e.g., "fc1" matches "model.decoder.layers.23.fc1"). - should_not_convert = any( - re.match(f"{key}\\.", full_name) or re.match(f"{key}", full_name) or full_name.endswith(key) - for key in patterns + patterns_tuple = tuple(patterns) + should_not_convert = ( + full_name.startswith(patterns_tuple) + or full_name.endswith(patterns_tuple) ) return not should_not_convert