import math import torch import torch.nn.functional as F from typing import Optional, Tuple from .model_mapper import ModelMapper from .custom_op import FusedAttention, FusedRoPE, MoE, FusedLinearAttention class Embedding(torch.nn.Module): def __init__(self, embed, config): super().__init__() self.hidden_size = config.hidden_size self.embed = embed self.embed_scale = 1.0 config_embed_scale = getattr(config, 'scale_emb', None) if config_embed_scale is not None: self.embed_scale = config_embed_scale elif config.model_type == 'gemma' or config.model_type == 'gemma2': self.embed_scale = self.hidden_size**0.5 if hasattr(embed, 'embed_scale'): self.embed_scale = embed.embed_scale # Replace ScaledWordEmbedding with plain Embedding to avoid double # scaling (scale is applied separately via model.scale_emb) if hasattr(embed, 'scalar_embed_scale'): plain_embed = torch.nn.Embedding(embed.num_embeddings, embed.embedding_dim, embed.padding_idx) plain_embed.weight = embed.weight self.embed = plain_embed def forward(self, input_ids): inputs_embeds = self.embed(input_ids).view(-1, 1, self.hidden_size) return inputs_embeds class RMSNorm(torch.nn.Module): def __init__(self, hidden_size, eps=1e-6): """ LlamaRMSNorm is equivalent to T5LayerNorm """ super().__init__() self.weight = torch.nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps def forward(self, hidden_states, gate=None): input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) hidden_states = self.weight * hidden_states.to(input_dtype) if gate is not None: hidden_states = hidden_states * F.silu(gate.to(torch.float32)) return hidden_states.to(input_dtype) def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 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) class Attention(torch.nn.Module): def __init__(self, attn, layer_id, config, rotary, mapper): super().__init__() self.export_fused_attn = False if config is None: return self.config = config self.kv_cache = True self.layer_id = layer_id self.rotary = rotary export_args = getattr(config, 'export_args', None) self.export_fused_rope = getattr(export_args, 'transformer_c4', False) self.hidden_size = config.hidden_size self.head_dim = config.head_dim if isinstance(config.num_attention_heads, list): self.num_heads = config.num_attention_heads[layer_id] self.num_key_value_heads = config.num_key_value_heads[layer_id] else: self.head_dim = config.head_dim self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads ModelMapper.do_map(self, attn, mapper['attention']) self.qk_norm_after_rope = getattr(config, 'qk_norm_after_rope', False) if not self.qk_norm_after_rope: self.qk_norm_after_rope = ( hasattr(attn, 'query_layernorm') and hasattr(attn, 'key_layernorm') ) # Read attention scaling from the original HF attention module if hasattr(attn, 'scaling'): self.attn_scaling = attn.scaling # k_eq_v / KV sharing detection (gemma4 and similar models) # Mapper key 'k_eq_v' acts as sentinel: its presence means per-layer detection is needed. # Detection is structural (works across HF versions): # - k_proj exists, v_proj missing → k_eq_v (K serves as both K and V) # - both missing + is_kv_shared_layer → pure KV sharing (no local K/V computation) if getattr(self, 'k_eq_v', None) is not None: has_k_proj = hasattr(self, 'k_proj') and self.k_proj is not None has_v_proj = hasattr(self, 'v_proj') and self.v_proj is not None self.k_eq_v = has_k_proj and not has_v_proj # per-layer head_dim auto-detection (gemma4 has varying head_dim) if hasattr(self, 'q_proj') and self.q_proj is not None: actual_head_dim = self.q_proj.out_features // self.num_heads if actual_head_dim != self.head_dim: self.head_dim = actual_head_dim if has_k_proj: actual_kv_heads = self.k_proj.out_features // self.head_dim if actual_kv_heads != self.num_key_value_heads: self.num_key_value_heads = actual_kv_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads else: self.k_eq_v = False # KV sharing (gemma4): track which layers share KV self.is_kv_shared_layer = getattr(attn, 'is_kv_shared_layer', False) self.kv_shared_layer_index = getattr(attn, 'kv_shared_layer_index', None) self.store_full_length_kv = getattr(attn, 'store_full_length_kv', False) # Create FusedAttention with KV sharing info kv_shared_idx = self.kv_shared_layer_index if self.is_kv_shared_layer else -1 self.fused_attn = FusedAttention(self.num_heads * self.head_dim, self.kv_cache, f'/layers.{layer_id}/self_attn/FusedAttention', layer_id, kv_shared_idx) self.fused_rope = FusedRoPE(self.head_dim, f'/layers.{layer_id}/self_attn/FusedRoPE') if hasattr(self, 'qkv_proj') and self.qkv_proj is not None: # split qkv linear to q, k, v split_sizes = [self.hidden_size] * 3 if self.qkv_proj.weight.shape[0] != self.hidden_size * 3: # M/GQA split_sizes = [ self.num_heads * self.head_dim, # q_size self.num_key_value_heads * self.head_dim, # k_size self.num_key_value_heads * self.head_dim # v_size ] self.q_proj = torch.nn.Linear(self.hidden_size, split_sizes[0]) self.k_proj = torch.nn.Linear(self.hidden_size, split_sizes[1]) self.v_proj = torch.nn.Linear(self.hidden_size, split_sizes[2]) if config.model_type == 'chatglm': # chatglm-6b qkv_weight = self.qkv_proj.weight.data.view(self.num_heads, 3, self.head_dim, self.hidden_size) self.q_proj.weight.data = qkv_weight[:, 0, :, :].reshape(self.hidden_size, self.hidden_size) self.k_proj.weight.data = qkv_weight[:, 1, :, :].reshape(self.hidden_size, self.hidden_size) self.v_proj.weight.data = qkv_weight[:, 2, :, :].reshape(self.hidden_size, self.hidden_size) qkv_bias = self.qkv_proj.bias.data.view(self.num_heads, 3, self.head_dim) self.q_proj.bias.data = qkv_bias[:, 0, :].reshape(self.hidden_size) self.k_proj.bias.data = qkv_bias[:, 1, :].reshape(self.hidden_size) self.v_proj.bias.data = qkv_bias[:, 2, :].reshape(self.hidden_size) else: # other qw, kw, vw = torch.split(self.qkv_proj.weight, split_sizes) self.q_proj.weight.data = qw self.k_proj.weight.data = kw self.v_proj.weight.data = vw if self.qkv_proj.bias is not None: qb, kb, vb = torch.split(self.qkv_proj.bias, split_sizes) self.q_proj.bias.data = qb self.k_proj.bias.data = kb self.v_proj.bias.data = vb else: data_type = self.q_proj.weight.dtype self.q_proj.bias.data = torch.zeros(split_sizes[0], dtype=data_type) self.k_proj.bias.data = torch.zeros(split_sizes[1], dtype=data_type) self.v_proj.bias.data = torch.zeros(split_sizes[2], dtype=data_type) self.q_proj.weight.requires_grad = False self.k_proj.weight.requires_grad = False self.v_proj.weight.requires_grad = False self.q_proj.bias.requires_grad = False self.k_proj.bias.requires_grad = False self.v_proj.bias.requires_grad = False self.past_key_value = None def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = None value_states = None if self.q_proj.out_features == 2 * self.num_heads * self.head_dim: reshaped = query_states.view(bsz, q_len, self.num_heads, self.head_dim * 2) query_states, gate = torch.split(reshaped, self.head_dim, dim=-1) gate = gate.reshape(bsz, q_len, -1) else: gate = None qk_norm_after_rope = getattr(self, 'qk_norm_after_rope', getattr(self.config, 'qk_norm_after_rope', False)) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim) q_norm_before_rope = not qk_norm_after_rope and hasattr(self, 'q_norm') and self.q_norm is not None # KV sharing: for shared layers, reuse KV from source layer (test mode only) shared_kv_cache = getattr(self, '_shared_kv_cache', None) use_shared_kv = (self.is_kv_shared_layer and shared_kv_cache is not None and self.kv_shared_layer_index in shared_kv_cache and not torch.onnx.is_in_onnx_export()) k_norm_before_rope = False if use_shared_kv: key_states, value_states = shared_kv_cache[self.kv_shared_layer_index] elif self.k_proj is not None: key_states = self.k_proj(hidden_states) if self.k_eq_v: value_states = key_states.clone() else: value_states = self.v_proj(hidden_states) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim) k_norm_before_rope = not qk_norm_after_rope and hasattr(self, 'k_norm') and self.k_norm is not None # gemma4 has v_norm (RMSNorm without scale) if hasattr(self, 'v_norm') and self.v_norm is not None: value_states = self.v_norm(value_states) else: # Pure KV sharing layer: no local K/V projections (e.g. gemma4 in HF>=5.5.4) # Dummy K/V for ONNX tracing; FusedAttention handles sharing via kv_shared_layer_index key_states = query_states.new_zeros(bsz, q_len, self.num_key_value_heads, self.head_dim) value_states = key_states k_norm_before_rope = False kv_seq_len = key_states.shape[1] if self.past_key_value is not None: kv_seq_len += self.past_key_value[0].shape[1] # rope if self.rotary is not None: cos, sin = rotary_pos_emb[0], rotary_pos_emb[1] use_fused_rope = ( self.export_fused_attn and torch.onnx.is_in_onnx_export() and self.export_fused_rope and not qk_norm_after_rope and not use_shared_kv and self.k_proj is not None and self.rotary.model_type not in ['chatglm', 'chatglm2', 'ernie4_5', 'glm_ocr'] and cos.shape[-1] == self.head_dim and sin.shape[-1] == self.head_dim ) fuse_qk_norm = use_fused_rope and q_norm_before_rope and k_norm_before_rope if use_fused_rope: if not fuse_qk_norm: if q_norm_before_rope: query_states = self.q_norm(query_states) if k_norm_before_rope: key_states = self.k_norm(key_states) query_states, key_states = self.fused_rope( query_states, key_states, cos, sin, self.q_norm if fuse_qk_norm else None, self.k_norm if fuse_qk_norm else None, ) else: # Most models apply q/k norm before rotary, but HunYuan applies it after rotary. if q_norm_before_rope: query_states = self.q_norm(query_states) if k_norm_before_rope: key_states = self.k_norm(key_states) query_states = self.rotary.apply_rotary_pos(query_states, cos, sin) if not use_shared_kv and self.k_proj is not None: key_states = self.rotary.apply_rotary_pos(key_states, cos, sin) elif q_norm_before_rope or k_norm_before_rope: if q_norm_before_rope: query_states = self.q_norm(query_states) if k_norm_before_rope: key_states = self.k_norm(key_states) if qk_norm_after_rope: if hasattr(self, 'q_norm') and self.q_norm is not None: query_states = self.q_norm(query_states) if not use_shared_kv and self.k_proj is not None and hasattr(self, 'k_norm') and self.k_norm is not None: key_states = self.k_norm(key_states) # MobileLLM model llama4_text has qk_norm after rotary if hasattr(self, 'qk_norm') and self.qk_norm is not None : query_states = self.qk_norm(query_states) key_states = self.qk_norm(key_states) if self.export_fused_attn and torch.onnx.is_in_onnx_export(): attn_output = self.fused_attn(query_states, key_states, value_states, attention_mask) if gate is not None: attn_output = attn_output * torch.sigmoid(gate) attn_output = self.o_proj(attn_output) return attn_output # kv cache if self.past_key_value is not None: past_key, past_value = self.past_key_value[0], self.past_key_value[1] key_states = torch.cat((past_key, key_states), dim=1) value_states = torch.cat((past_value, value_states), dim=1) if not use_shared_kv: self.past_key_value = torch.stack((key_states, value_states)) query_states = query_states.transpose(1, 2) if use_shared_kv: # Shared KV is already in transposed format [B, heads, head_dim, seq] / [B, heads, seq, head_dim] pass else: key_states = key_states.permute([0, 2, 3, 1]) value_states = value_states.transpose(1, 2) # Store KV for sharing (source layers that other layers will read from) if self.store_full_length_kv and shared_kv_cache is not None: shared_kv_cache[self.layer_id] = (key_states.clone(), value_states.clone()) # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) #------- attention ---------- # query_states @ key_states attn_scaling = getattr(self, 'attn_scaling', 1.0 / math.sqrt(self.head_dim)) attn_weights = torch.matmul(query_states, key_states) * attn_scaling # attention_mask if attention_mask.dtype in (torch.bool, torch.int32): # chatglm attn_weights.masked_fill_(attention_mask, -10000.0) else: attn_weights = attn_weights + attention_mask if hasattr(self, 'sinks'): sinks = self.sinks.reshape(1, -1, 1, 1).to(torch.float32).expand(query_states.shape[0], -1, query_states.shape[-2], -1) combined_logits = torch.cat([attn_weights, sinks], dim=-1) combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values probs = F.softmax(combined_logits, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_weights = probs[..., :-1] # we drop the sink here else: # upcast softmax to fp32 attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) # attn_weights @ value_states attn_output = torch.matmul(attn_weights, value_states) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, -1) if gate is not None: attn_output = attn_output * torch.sigmoid(gate) attn_output = self.o_proj(attn_output) return attn_output def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6): """This function is intended to align with the l2norm implementation in the FLA library.""" inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) return x * inv_norm def torch_chunk_gated_delta_rule( query, key, value, g, beta, chunk_size=64, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) query, key, value, beta, g = [ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) ] batch_size, num_heads, sequence_length, k_head_dim = key.shape v_head_dim = value.shape[-1] pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size query = F.pad(query, (0, 0, 0, pad_size)) key = F.pad(key, (0, 0, 0, pad_size)) value = F.pad(value, (0, 0, 0, pad_size)) beta = F.pad(beta, (0, pad_size)) g = F.pad(g, (0, pad_size)) total_sequence_length = sequence_length + pad_size scale = 1 / (query.shape[-1] ** 0.5) query = query * scale v_beta = value * beta.unsqueeze(-1) k_beta = key * beta.unsqueeze(-1) # reshape to chunks query, key, value, k_beta, v_beta = [ x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta) ] g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) # chunk decay g = g.cumsum(dim=-1) decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) for i in range(1, chunk_size): row = attn[..., i, :i].clone() sub = attn[..., :i, :i].clone() attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) value = attn @ v_beta k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) last_recurrent_state = ( torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) if initial_state is None else initial_state.to(value) ) core_attn_out = torch.zeros_like(value) mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) # for each chunk for i in range(0, total_sequence_length // chunk_size): q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state v_new = v_i - v_prime attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state core_attn_out[:, :, i] = attn_inter + attn @ v_new last_recurrent_state = ( last_recurrent_state * g[:, :, i, -1, None, None].exp() + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new ) if not output_final_state: last_recurrent_state = None core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]) core_attn_out = core_attn_out[:, :, :sequence_length] core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) return core_attn_out, last_recurrent_state def torch_recurrent_gated_delta_rule( query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False ): initial_dtype = query.dtype if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) query, key, value, beta, g = [ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) ] batch_size, num_heads, sequence_length, k_head_dim = key.shape v_head_dim = value.shape[-1] scale = 1 / (query.shape[-1] ** 0.5) query = query * scale core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value) last_recurrent_state = ( torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) if initial_state is None else initial_state.to(value) ) for i in range(sequence_length): q_t = query[:, :, i] k_t = key[:, :, i] v_t = value[:, :, i] g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1) beta_t = beta[:, :, i].unsqueeze(-1) last_recurrent_state = last_recurrent_state * g_t kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) delta = (v_t - kv_mem) * beta_t last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2) if not output_final_state: last_recurrent_state = None core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) return core_attn_out, last_recurrent_state def torch_gated_delta_rule( query, # [B, L, H, K] query vectors key, # [B, L, H, K] key vectors value, # [B, L, H, V] value vectors g, # [B, L, H] log-space decay (negative values) beta, # [B, L, H] learning rate for delta update initial_state=None, # [B, H, K, V] initial recurrent state output_final_state=False, use_qk_l2norm_in_kernel=False, ): """ Non-chunk implementation of Gated Delta Rule (Linear Attention). Mathematically equivalent to torch_chunk_gated_delta_rule. Maintains a key-value memory (state S) of shape [K, V] per head, updated at each timestep using the Delta Learning Rule with gating. Per-step formula (for each head independently): ───────────────────────────────────────────────── S_t = S_{t-1} * exp(g_t) # 1. decay old memory v_pred = S_t^T @ k_t # 2. predict value for current key delta = beta_t * (v_t - v_pred) # 3. prediction error * learning rate S_t = S_t + k_t @ delta^T # 4. update memory (outer product) o_t = S_t^T @ (q_t / sqrt(d_k)) # 5. query the memory ───────────────────────────────────────────────── Shapes: S: [B, H, K, V] recurrent state (key-value memory) k_t: [B, H, K] key at timestep t v_t: [B, H, V] value at timestep t q_t: [B, H, K] query at timestep t g_t: [B, H] log-decay scalar per head beta_t:[B, H] learning rate scalar per head o_t: [B, H, V] output at timestep t """ initial_dtype = query.dtype # Optional: L2 normalize Q, K before computation if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) key = l2norm(key, dim=-1, eps=1e-6) # [B, L, H, D] -> [B, H, L, D], all in float32 query, key, value, beta, g = [ x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) ] B, H, L, K = key.shape V = value.shape[-1] # Scale query: q = q / sqrt(d_k) query = query * (K ** -0.5) # Initialize recurrent state S: [B, H, K, V] if initial_state is None: S = torch.zeros(B, H, K, V, device=query.device, dtype=torch.float32) else: S = initial_state.to(torch.float32) outputs = [] for t in range(L): q_t = query[:, :, t] # [B, H, K] k_t = key[:, :, t] # [B, H, K] v_t = value[:, :, t] # [B, H, V] g_t = g[:, :, t] # [B, H] beta_t = beta[:, :, t] # [B, H] # ── Step 1: Decay ── # S = S * exp(g_t), g_t < 0 so this shrinks old memory S = S * g_t[:, :, None, None].exp() # ── Step 2: Read (predict value for current key) ── # v_pred = S^T @ k_t : [B, H, K, V]^T @ [B, H, K, 1] -> [B, H, V] v_pred = (S.transpose(-1, -2) @ k_t.unsqueeze(-1)).squeeze(-1) # ── Step 3: Delta (prediction error * learning rate) ── # delta = beta_t * (v_t - v_pred) : [B, H, V] delta = beta_t[:, :, None] * (v_t - v_pred) # ── Step 4: Write (update memory with outer product) ── # S += k_t @ delta^T : [B, H, K, 1] @ [B, H, 1, V] -> [B, H, K, V] S = S + k_t.unsqueeze(-1) @ delta.unsqueeze(-2) # ── Step 5: Query (read output from updated memory) ── # o_t = S^T @ q_t : [B, H, K, V]^T @ [B, H, K, 1] -> [B, H, V] o_t = (S.transpose(-1, -2) @ q_t.unsqueeze(-1)).squeeze(-1) outputs.append(o_t) # Stack: list of [B, H, V] -> [B, H, L, V] core_attn_out = torch.stack(outputs, dim=2) if not output_final_state: S = None # [B, H, L, V] -> [B, L, H, V], restore original dtype core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) return core_attn_out, S class ShortConvAttention(torch.nn.Module): def __init__(self, attn, layer_id, config, mapper): super().__init__() self.layer_id = layer_id self.hidden_size = config.hidden_size self.conv_kernel_size = config.conv_L_cache ModelMapper.do_map(self, attn, mapper['linear_attention']) self.fused_attn = FusedLinearAttention( name=f'/layers.{layer_id}/self_attn/FusedLinearAttention', attn_type="short_conv", num_k_heads=1, num_v_heads=1, head_k_dim=self.hidden_size, head_v_dim=self.hidden_size, use_qk_l2norm=False ) self.conv_state = None def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: # Note: ShortConvAttention is mask-free; `attention_mask` is accepted # only to keep the call signature uniform with `Attention.forward` and # is intentionally unused. batch_size, seq_len, _ = hidden_states.shape # in_proj: [B, L, H] -> [B, L, 3H] bcx = self.in_proj(hidden_states) if torch.onnx.is_in_onnx_export(): # ONNX path: pass through FusedLinearAttention custom op bcx_t = bcx.transpose(1, 2) # [B, 3H, L] gate = torch.zeros(batch_size, seq_len, 1, dtype=bcx.dtype, device=bcx.device) beta = torch.zeros(batch_size, seq_len, 1, dtype=bcx.dtype, device=bcx.device) attn_out = self.fused_attn(bcx_t, gate, beta, self.conv.weight.data.detach()) # attn_out: [B, L, 1, H] -> [B, L, H] attn_out = attn_out.view(batch_size, seq_len, -1) output = self.out_proj(attn_out) return output # Test path: manual computation # Split into B_, C_, x_ each [B, L, H] B_, C_, x_ = bcx.chunk(3, dim=-1) # Bx = B_ * x_ Bx = B_ * x_ # Transpose for conv: [B, H, L] Bx = Bx.transpose(1, 2) conv_state_size = self.conv_kernel_size - 1 if self.conv_state is not None: conv_input = torch.cat([self.conv_state, Bx], dim=-1) conv_out = F.conv1d(conv_input, self.conv.weight, padding=0, groups=self.hidden_size) new_conv_state = conv_input[:, :, -conv_state_size:] else: new_conv_state = F.pad(Bx, (conv_state_size - Bx.shape[-1], 0)) conv_out = self.conv(Bx)[:, :, :seq_len] # No SiLU for short_conv (unlike gated_delta_rule) # Transpose back: [B, H, L] -> [B, L, H] conv_out = conv_out.transpose(1, 2) # y = C_ * conv_out y = C_ * conv_out output = self.out_proj(y) self.conv_state = new_conv_state return output class LinearAttention(torch.nn.Module): def __init__(self, attn, layer_id, config, rotary, mapper): super().__init__() self.layer_id = layer_id self.config = config self.hidden_size = config.hidden_size self.head_dim = config.head_dim self.conv_kernel_size = config.linear_conv_kernel_dim self.conv_state_size = self.conv_kernel_size - 1 self.head_k_dim = config.linear_key_head_dim self.head_v_dim = config.linear_value_head_dim self.num_k_heads = config.linear_num_key_heads self.num_v_heads = config.linear_num_value_heads self.key_dim = self.head_k_dim * self.num_k_heads self.value_dim = self.head_v_dim * self.num_v_heads self.conv_dim = self.key_dim * 2 + self.value_dim ModelMapper.do_map(self, attn, mapper['linear_attention']) original_norm = self.norm self.norm = RMSNorm(self.head_v_dim, eps=config.rms_norm_eps) self.norm.weight.data = original_norm.weight.data self.fused_attn = FusedLinearAttention( name=f'/layers.{layer_id}/self_attn/FusedLinearAttention', attn_type="gated_delta_rule", num_k_heads=self.num_k_heads, num_v_heads=self.num_v_heads, head_k_dim=self.head_k_dim, head_v_dim=self.head_v_dim, use_qk_l2norm=True ) self.conv_state = None self.rnn_state = None def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: # Note: LinearAttention is mask-free; `attention_mask` is accepted # only to keep the call signature uniform with `Attention.forward` and # is intentionally unused. batch_size, seq_len, _ = hidden_states.shape # 1. Linear Projections # mixed_qkv: [B, L, 2*key_dim + value_dim] mixed_qkv = self.in_proj_qkv(hidden_states) # Transpose for Conv1d: [B, Dim, L] mixed_qkv = mixed_qkv.transpose(1, 2) # Gate, Beta, Z projections z = self.in_proj_z(hidden_states) # [B, L, value_dim] b = self.in_proj_b(hidden_states) # [B, L, num_v_heads] a = self.in_proj_a(hidden_states) # [B, L, num_v_heads] # 2. Pre-compute gates beta = torch.sigmoid(b) gate = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) if torch.onnx.is_in_onnx_export(): attn_out = self.fused_attn(mixed_qkv, gate, beta, self.conv1d.weight.data.detach()) attn_out = attn_out.reshape(-1, self.head_v_dim) z = z.reshape(-1, self.head_v_dim) attn_out = self.norm(attn_out, z) attn_out = attn_out.view(batch_size, seq_len, -1) output = self.out_proj(attn_out) return output # === Normal path: full computation for testing === # 3. State Management (Conv State & Recurrent State) if self.conv_state is not None: conv_state = self.conv_state conv_input = torch.cat([conv_state, mixed_qkv], dim=-1) mixed_qkv = F.silu(F.conv1d(conv_input, self.conv1d.weight, self.conv1d.bias, padding=0, groups=self.conv_dim)) new_conv_state = conv_input[:, :, -self.conv_state_size:] else: new_conv_state = F.pad(mixed_qkv, (self.conv_state_size - mixed_qkv.shape[-1], 0)) mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, :seq_len]) # 4. Split Q, K, V mixed_qkv = mixed_qkv.transpose(1, 2) query, key, value = torch.split( mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1 ) query = query.view(batch_size, seq_len, self.num_k_heads, self.head_k_dim) key = key.view(batch_size, seq_len, self.num_k_heads, self.head_k_dim) value = value.view(batch_size, seq_len, self.num_v_heads, self.head_v_dim) # 5. GQA Expansion if self.num_v_heads > self.num_k_heads: factor = self.num_v_heads // self.num_k_heads query = query.repeat_interleave(factor, dim=2) key = key.repeat_interleave(factor, dim=2) # 6. Gated Delta Rule if self.rnn_state is None: attn_out, last_recurrent_state = torch_gated_delta_rule( query, key, value, g=gate, beta=beta, initial_state=None, output_final_state=True, use_qk_l2norm_in_kernel=True, ) else: recurrent_state = self.rnn_state attn_out, last_recurrent_state = torch_recurrent_gated_delta_rule( query, key, value, g=gate, beta=beta, initial_state=recurrent_state, output_final_state=True, use_qk_l2norm_in_kernel=True, ) # 7. Post-process attn_out = attn_out.reshape(-1, self.head_v_dim) z = z.reshape(-1, self.head_v_dim) attn_out = self.norm(attn_out, z) attn_out = attn_out.view(batch_size, seq_len, -1) output = self.out_proj(attn_out) # Update internal state self.conv_state = new_conv_state self.rnn_state = last_recurrent_state return output def create_linear_attention(attn, layer_id, config, rotary, mapper): """Factory function for creating LinearAttention variants based on config.""" if hasattr(config, 'conv_L_cache') and config.conv_L_cache > 0: return ShortConvAttention(attn, layer_id, config, mapper) return LinearAttention(attn, layer_id, config, rotary, mapper) def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def _compute_yarn_parameters(rotary_dim, base_theta, scaling_config, max_position_embeddings): """ 计算 YaRN (Yet another RoPE extensioN method) 的参数。 此函数等价于 Hugging Face Transformers 中的 YaRN 实现。 Args: rotary_dim (int): RoPE 的维度。 base_theta (float): RoPE 的基础 theta 值。 scaling_config (dict): 包含 YaRN 特定配置的字典。 max_position_embeddings (int): 模型的最大位置编码。 Returns: tuple[torch.Tensor, float]: - inv_freq (torch.Tensor): 计算好的、用于 RoPE 的逆频率 (即 theta)。 - attention_scaling (float): 应用于 Query 向量的缩放因子。 """ def get_mscale(scale, m_scale): if scale <= 1: return 1.0 return 0.1 * m_scale * math.log(scale) + 1.0 def find_correction_dim(num_rotations, d, b, max_pos): return (d * math.log(max_pos / (num_rotations * 2 * math.pi))) / (2 * math.log(b)) def find_correction_range(low_rot, high_rot, d, b, max_pos): low = find_correction_dim(low_rot, d, b, max_pos) high = find_correction_dim(high_rot, d, b, max_pos) return max(0, math.floor(low)), min(d - 1, math.ceil(high)) def linear_ramp_factor(mn, mx, d): if mn == mx: mx += 0.001 linear_func = (torch.arange(d, dtype=torch.float32) - mn) / (mx - mn) return torch.clamp(linear_func, 0, 1) # 1. 提取 YaRN 参数 factor = scaling_config['factor'] beta_fast = scaling_config.get("beta_fast", 32) beta_slow = scaling_config.get("beta_slow", 1) original_max_pos = scaling_config.get("original_max_position_embeddings", max_position_embeddings) mscale = scaling_config.get("mscale", 1.0) # 2. 计算 attention_scaling (即 attention_factor) attention_scaling = get_mscale(factor, mscale) # 3. 计算 inv_freq (即 theta) dim = rotary_dim # 计算插值和外推的频率 pos_freqs = base_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) inv_freq_extrapolation = 1.0 / pos_freqs inv_freq_interpolation = 1.0 / (factor * pos_freqs) # 找到需要修正的维度范围 low, high = find_correction_range(beta_fast, beta_slow, dim, base_theta, original_max_pos) # 创建维度混合的 ramp (作用于 dim//2 的频率上) inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2) # 混合插值和外推频率，得到最终的 inv_freq inv_freq = ( inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + inv_freq_extrapolation * inv_freq_extrapolation_factor ) return inv_freq, attention_scaling class Rotary(torch.nn.Module): def __init__(self, config): super().__init__() if config is None: return self.rope_theta = config.rope_theta self.rope_ratio = config.rope_ratio if self.rope_ratio is not None: self.rope_theta *= self.rope_ratio self.rotary_dim = config.head_dim self.model_type = config.model_type if hasattr(config, 'rotary_dim'): self.rotary_dim = config.rotary_dim if self.model_type == 'chatglm': self.rotary_dim = config.head_dim // 2 # Qwen3.5 / LFM2 style: flat rope_parameters dict if hasattr(config, 'rope_parameters') and config.rope_parameters is not None: rp = config.rope_parameters # Detect gemma4-style per-layer-type rope_parameters (dict of dicts) is_per_layer_type = any(isinstance(v, dict) for v in rp.values()) if not is_per_layer_type: if 'rope_theta' in rp: self.rope_theta = rp['rope_theta'] if 'partial_rotary_factor' in rp: self.partial_rotary_factor = rp['partial_rotary_factor'] self.rotary_dim = int(self.rotary_dim * self.partial_rotary_factor) config.rope_scaling = rp self.mrope_section = None self.theta_sections = None self.attention_scaling = 1.0 self.is_scaled = False self.mrope_interleaved = False def get_theta(): return 1.0 / (self.rope_theta ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.float32) / self.rotary_dim)) # default rope type's theta self.theta = get_theta() # other type if hasattr(config, 'rope_scaling') and config.rope_scaling is not None: scaling_config = config.rope_scaling # get rope_type rope_type = 'default' if 'type' in config.rope_scaling: rope_type = config.rope_scaling['type'] elif 'rope_type' in config.rope_scaling: rope_type = config.rope_scaling['rope_type'] # gen theta for rope_type if rope_type == 'dynamic': # NTK if 'alpha' in config.rope_scaling: # NTKAlpha in Hunyuan self.rope_theta *= (config.rope_scaling['alpha'] ** (self.rotary_dim / (self.rotary_dim - 2))) else: # NTKScaling pass self.theta = get_theta() elif rope_type == 'yarn': self.is_scaled = True self.theta, self.attention_scaling = _compute_yarn_parameters( rotary_dim=self.rotary_dim, base_theta=self.rope_theta, scaling_config=scaling_config, max_position_embeddings=config.max_position_embeddings ) elif rope_type == 'longrope': # longrope in MiniCPM self.is_scaled = True original_max_position_embeddings = config.rope_scaling['original_max_position_embeddings'] scale = (config.max_position_embeddings / original_max_position_embeddings) self.attention_scaling = math.sqrt(1 + math.log(scale) / math.log(original_max_position_embeddings)) # long_factor = config.rope_scaling['long_factor'] short_factor = config.rope_scaling['short_factor'] self.theta = get_theta() / torch.tensor(short_factor, dtype=torch.float32) # mrope for multimode if 'mrope_section' in scaling_config: self.mrope_interleaved = scaling_config.get('mrope_interleaved', False) self.mrope_section = scaling_config['mrope_section'] self.theta = get_theta().unsqueeze(0) self.theta_sections = self.theta.split(self.mrope_section, dim=-1) def apply_interleaved_mrope(freqs, mrope_section): # mrope apply func from qwen3-vl freqs_t = freqs[0] # just overwrite the first dimension T for dim, offset in enumerate((1, 2), start=1): # H, W length = mrope_section[dim] * 3 idx = slice(offset, length, 3) freqs_t[..., idx] = freqs[dim, ..., idx] return freqs_t if self.mrope_interleaved: half_rotary = self.rotary_dim // 2 freq_idx = torch.arange(0, 3 * half_rotary).reshape(3, 1, half_rotary) self.mrope_reindex = apply_interleaved_mrope(freq_idx, self.mrope_section).flatten() self.is_mrope = self.theta_sections is not None or self.mrope_interleaved def forward(self, position_ids): if self.is_mrope: return self.mrope_forward(position_ids) position_ids = position_ids.float().reshape(-1, 1) idx_theta = position_ids * self.theta.to(position_ids.device) rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)]) if self.model_type == 'ernie4_5': rotary_pos_emb = torch.stack((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.reshape(*rotary_pos_emb.shape[:-2], -1) elif self.model_type != 'chatglm2': rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1) if self.is_scaled: rotary_pos_emb *= self.attention_scaling return rotary_pos_emb def mrope_forward(self, position_ids): position_ids = position_ids.float().unsqueeze(-1) if self.mrope_interleaved: idx_theta = position_ids * self.theta.to(position_ids.device) idx_theta = idx_theta.transpose(1, 0).reshape(-1, 3 * self.rotary_dim // 2) idx_theta = idx_theta[:, self.mrope_reindex] else: idx_theta = torch.concat([ position_ids[0] * self.theta_sections[0], position_ids[1] * self.theta_sections[1], position_ids[2] * self.theta_sections[2] ], dim=-1) rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)]) if self.model_type in ['glm_ocr']: # interleaved doubling: [c0,c0,c1,c1,...,cn,cn] rotary_pos_emb = torch.stack((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.reshape(*rotary_pos_emb.shape[:-2], -1) else: rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1) return rotary_pos_emb def apply_rotary_pos(self, x, cos, sin): if self.model_type == 'chatglm': return self.chatglm_rotary_pos(x, cos, sin) if self.model_type == 'chatglm2': return self.chatglm2_rotary_pos(x, cos, sin) if self.model_type in ['phi-msft', 'qwen3_5', 'qwen3_5_moe']: return self.phi_rotary_pos(x, cos, sin) if self.model_type in ['ernie4_5', 'glm_ocr']: return self.ernie_rotary_pos(x, cos, sin) # Auto-detect partial rotary: cos/sin dim < x dim if cos.shape[-1] < x.shape[-1]: return self.phi_rotary_pos(x, cos, sin) return self.llama_rotary_pos(x, cos, sin) def llama_rotary_pos(self, x, cos, sin): x = (x * cos) + (rotate_half(x) * sin) return x def ernie_rotary_pos(self, x, cos, sin): rotate_half_x = torch.stack( [-x[:, :, :, 1::2], x[:, :, :, 0::2]], dim=-1 ).reshape(x.shape) x = (x * cos) + (rotate_half_x * sin) return x def phi_rotary_pos(self, x, cos, sin): # Use cos dim to determine rotary_dim (handles per-layer different rotary_dim) rotary_dim = cos.shape[-1] x, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] x = (x * cos) + (rotate_half(x) * sin) return torch.cat((x, x_pass), dim=-1) def chatglm2_rotary_pos(self, x, cos, sin): x, x_pass = x[..., :self.rotary_dim], x[..., self.rotary_dim:] b, s, n, h = x.shape xshaped = x.view(b, s, n, h//2, 2) x = torch.concat( [ xshaped[..., 0] * cos - xshaped[..., 1] * sin, xshaped[..., 1] * cos + xshaped[..., 0] * sin, ], -1, ) return torch.cat((x, x_pass), dim=-1) def chatglm_rotary_pos(self, x, cos, sin): seq = x.shape[1] x1, x2 = x[..., :self.rotary_dim], x[..., self.rotary_dim:] cos1, sin1 = cos[:, :seq, ...], sin[:, :seq, ...] cos2, sin2 = cos[:, seq:, ...], sin[:, seq:, ...] x1 = (x1 * cos1) + (rotate_half(x1) * sin1) x2 = (x2 * cos2) + (rotate_half(x2) * sin2) return torch.cat((x1, x2), dim=-1) class VisionRotary(Rotary): def __init__(self, config): super().__init__(config) # support [h_pos, w_pos] def forward(self, position_ids): # [2, patch_len, 1] position_ids = position_ids.float().unsqueeze(-1) idx_theta = position_ids * self.theta # [patch_len, rotary_dim] idx_theta = idx_theta.permute(1, 0, 2).reshape(-1, self.rotary_dim) rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)]) rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1) return rotary_pos_emb def apply_rotary_pos(self, x, cos, sin): x = (x * cos) + (rotate_half(x) * sin) return x class GptOssExpert(torch.nn.Module): def __init__(self, hidden_size, expert_dim): super().__init__() self.expert_dim = expert_dim self.gate_up_proj_linear = torch.nn.Linear(hidden_size, 2 * expert_dim) self.down_proj_linear = torch.nn.Linear(expert_dim, hidden_size) self.alpha = 1.702 self.limit = 7.0 def forward(self, hidden_states: torch.Tensor, debug=False) -> torch.Tensor: gate_up = self.gate_up_proj_linear(hidden_states) gate, up = gate_up[..., ::2], gate_up[..., 1::2] # gate = gate.clamp(min=None, max=self.limit) limit_tensor = torch.tensor(self.limit, device=gate.device, dtype=gate.dtype) gate = torch.min(gate, limit_tensor) up = up.clamp(min=-self.limit, max=self.limit) glu = gate * torch.sigmoid(gate * self.alpha) gated_output = (up + 1) * glu out = self.down_proj_linear(gated_output) return out class Qwen3Expert(torch.nn.Module): def __init__(self, hidden_size, expert_dim, act_fn): super().__init__() self.expert_dim = expert_dim self.gate_up_proj_linear = torch.nn.Linear(hidden_size, 2 * expert_dim, bias=False) self.down_proj_linear = torch.nn.Linear(expert_dim, hidden_size, bias=False) self.act_fn = act_fn def forward(self, hidden_states: torch.Tensor, debug=False) -> torch.Tensor: gate_up = self.gate_up_proj_linear(hidden_states) # gate, up = gate_up[..., ::2], gate_up[..., 1::2] gate, up = gate_up.chunk(2, dim=-1) out = self.down_proj_linear(up * self.act_fn(gate)) return out class Mlp(torch.nn.Module): def __init__(self, mlp, mapper, layer_id): super().__init__() self.layer_id = layer_id ModelMapper.do_map(self, mlp, mapper['mlp']) self.is_moe = hasattr(self, 'experts') self.export_moe = False self.custom_moe = MoE(self.num_experts, self.top_k, layer_id) if isinstance(self.experts, torch.nn.ModuleList): self.moe_type = 'qwen3_moe' else: self.moe_type = 'qwen3_5_moe' self.norm_topk_prob = True # refacte experts to qwen3_experts original_experts = self.experts hidden_size = getattr(original_experts, 'hidden_dim', None) or getattr(original_experts, 'hidden_size') expert_dim = getattr(original_experts, 'intermediate_dim', None) or getattr(original_experts, 'intermediate_size') act_fn = original_experts.act_fn new_experts_list = torch.nn.ModuleList() for i in range(self.num_experts): expert_mlp = Qwen3Expert(hidden_size, expert_dim, act_fn) expert_mlp.gate_up_proj_linear.weight.data = original_experts.gate_up_proj.data[i] expert_mlp.down_proj_linear.weight.data = original_experts.down_proj.data[i] new_experts_list.append(expert_mlp) self.experts = new_experts_list if not isinstance(self.gate, torch.nn.Linear): gate = torch.nn.Linear(hidden_size, self.num_experts, bias=False) gate.weight.data = self.gate.weight.data self.gate = gate if hasattr(self, 'expert_bias') and self.expert_bias is not None: self.moe_type = 'lfm2_moe' if hasattr(self, 'router'): self.moe_type = 'gpt_oss' hidden_dim = self.router.weight.shape[1] self.gate = torch.nn.Linear(hidden_dim, self.num_experts, bias=True) self.gate.weight.data = self.router.weight.data self.gate.bias.data = self.router.bias.data # refacte experts to qwen3_experts original_experts = self.experts expert_dim = original_experts.expert_dim new_experts_list = torch.nn.ModuleList() for i in range(self.num_experts): expert_mlp = GptOssExpert(hidden_dim, expert_dim) expert_mlp.gate_up_proj_linear.weight.data = original_experts.gate_up_proj.data[i].transpose(0, 1) expert_mlp.gate_up_proj_linear.bias.data = original_experts.gate_up_proj_bias.data[i] expert_mlp.down_proj_linear.weight.data = original_experts.down_proj.data[i].transpose(0, 1) expert_mlp.down_proj_linear.bias.data = original_experts.down_proj_bias.data[i] new_experts_list.append(expert_mlp) self.experts = new_experts_list del self.router def forward(self, hidden_states: torch.Tensor): if not self.is_moe: # general Mlp return self.down_proj(self.act_fn(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) # MoE Mlp batch_size, sequence_length, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) if hasattr(self, 'shared_expert'): shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * self.shared_expert(hidden_states) shared_expert_output = shared_expert_output.reshape(batch_size, sequence_length, hidden_dim) else: shared_expert_output = None if self.moe_type == 'lfm2_moe': router_logits = self.gate(hidden_states) routing_weights = router_logits.sigmoid() scores_for_routing = routing_weights + self.expert_bias _, selected_experts = torch.topk(scores_for_routing, self.top_k, dim=-1) routing_weights = torch.gather(routing_weights, dim=-1, index=selected_experts) if self.norm_topk_prob: routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6) routing_weights = (routing_weights * self.routed_scaling_factor).to(hidden_states.dtype) elif self.moe_type == 'gpt_oss': router_logits = self.gate(hidden_states) routing_weights, selected_experts = torch.topk(router_logits, self.top_k, dim=-1) routing_weights = F.softmax(routing_weights, dim=-1, dtype=torch.float).to(hidden_states.dtype) else: router_logits = self.gate(hidden_states) routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) if self.norm_topk_prob: # only diff with mixtral sparse moe block! routing_weights /= routing_weights.sum(dim=-1, keepdim=True) # we cast back to the input dtype routing_weights = routing_weights.to(hidden_states.dtype) if self.export_moe: expert_output = self.custom_moe(hidden_states, routing_weights, selected_experts) if shared_expert_output is not None: expert_output = expert_output + shared_expert_output return expert_output final_hidden_states = torch.zeros( (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device ) if False: # cpp impl seqlen, topk = selected_experts.shape if seqlen == 1: expert_idx = int(selected_experts[0, 0]) scale = float(routing_weights[0, 0]) output = self.experts[expert_idx](hidden_states) * scale for i in range(1, topk): expert_idx = int(selected_experts[0, i]) scale = float(routing_weights[0, i]) output += self.experts[expert_idx](hidden_states) * scale return output hss = torch.split(hidden_states, 1) expertWorks = [[] for i in range(self.num_experts)] for i in range(seqlen): for j in range(topk): expert_idx = int(selected_experts[i, j]) scale = float(routing_weights[i, j]) expertWorks[expert_idx].append((i, scale)) for i in range(self.num_experts): if len(expertWorks[i]) == 0: continue input_hs = [] for token_id, scale in expertWorks[i]: input_hs.append(hss[token_id]) output_hs = self.experts[i](torch.concat(input_hs)) output_hss = torch.split(output_hs, 1) for j in range(len(expertWorks[i])): token_id, scale = expertWorks[i][j] scale_hs = output_hss[j] * scale final_hidden_states[token_id] += scale_hs.squeeze(0) return final_hidden_states # One hot encode the selected experts to create an expert mask # this will be used to easily index which expert is going to be sollicitated expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() # Loop over all available experts in the model and perform the computation on each expert # for expert_idx in range(self.num_experts): for expert_idx in expert_hit: expert_idx = expert_idx[0] idx, top_x = torch.where(expert_mask[expert_idx]) expert_layer = self.experts[expert_idx] # Index the correct hidden states and compute the expert hidden state for # the current expert. We need to make sure to multiply the output hidden # states by `routing_weights` on the corresponding tokens (top-1 and top-2) current_state = hidden_states[None, top_x].reshape(-1, hidden_dim) current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None] # However `index_add_` only support torch tensors for indexing so we'll use # the `top_x` tensor here. final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) if shared_expert_output is not None: final_hidden_states = final_hidden_states + shared_expert_output return final_hidden_states class Decoder(torch.nn.Module): def __init__(self, decoder, layer_id, config, rotary=None, mapper=None): super().__init__() if rotary is None: rotary = config.rotary if mapper is None: mapper = config.model_map ModelMapper.do_map(self, decoder, mapper['decoder']) if 'mlp' in mapper and hasattr(self.mlp, 'experts'): self.mlp = Mlp(self.mlp, mapper, layer_id) # gemma4 MoE: router and experts are at decoder layer level (parallel to dense MLP) self.has_gemma4_moe = hasattr(self, 'experts') and self.experts is not None if self.has_gemma4_moe: original_experts = self.experts num_experts = original_experts.num_experts hidden_size = original_experts.hidden_dim expert_dim = original_experts.intermediate_dim act_fn = original_experts.act_fn # Refactor 3D Parameter experts into ModuleList of Qwen3Expert new_experts_list = torch.nn.ModuleList() for i in range(num_experts): expert_mlp = Qwen3Expert(hidden_size, expert_dim, act_fn) expert_mlp.gate_up_proj_linear.weight.data = original_experts.gate_up_proj.data[i] expert_mlp.down_proj_linear.weight.data = original_experts.down_proj.data[i] new_experts_list.append(expert_mlp) self.experts = new_experts_list # Extract gate Linear from router (router has norm+scale+proj+per_expert_scale) self.moe_gate = torch.nn.Linear(hidden_size, num_experts, bias=False) self.moe_gate.weight.data = self.router.proj.weight.data self.moe_router_norm = self.router.norm self.moe_router_scale = self.router.scale.data self.moe_router_scalar_root = self.router.scalar_root_size self.moe_per_expert_scale = self.router.per_expert_scale.data self.moe_num_experts = num_experts self.moe_top_k = config.origin_config.text_config.top_k_experts self.custom_moe = MoE(num_experts, self.moe_top_k, layer_id) self.export_moe = False del self.router self.layer_type = 'full_attention' if hasattr(self, 'self_attn') and self.self_attn is not None: self.self_attn = Attention(self.self_attn, layer_id, config, rotary, mapper) if hasattr(self, 'linear_attn') and self.linear_attn is not None: self.self_attn = create_linear_attention(self.linear_attn, layer_id, config, rotary, mapper) self.layer_type = 'linear_attention' self.hidden_size = config.hidden_size if hasattr(config, 'num_hidden_layers'): # minicpm self.num_hidden_layers = config.num_hidden_layers # chatglm self.alpha = (2 * config.num_hidden_layers) ** 0.5 if config.model_type == 'chatglm' else 1.0 def forward( self, hidden_states: torch.Tensor, rotary_pos_emb: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: hidden_states = hidden_states.view(1, -1, self.hidden_size) residual = hidden_states hidden_states = self.input_layernorm(hidden_states) norm_hidden_states = hidden_states # Self Attention or Linear Attention if self.layer_type == 'full_attention': hidden_states = self.self_attn( hidden_states=hidden_states, rotary_pos_emb=rotary_pos_emb, attention_mask=attention_mask, ) elif self.layer_type == 'linear_attention': hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, ) # Fully Connected if not hasattr(self, 'post_attention_layernorm'): # phi feed_forward_hidden_states = self.mlp(norm_hidden_states) hidden_states = hidden_states + feed_forward_hidden_states + residual elif hasattr(self, 'alpha') and self.alpha != 1.0: # chatglm-6b hidden_states = norm_hidden_states * self.alpha + hidden_states mlp_input = self.post_attention_layernorm(hidden_states) mlp_output = self.mlp(mlp_input) hidden_states = mlp_input * self.alpha + mlp_output elif hasattr(self, 'pre_feedforward_layernorm'): # gemma2 / gemma4 hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.pre_feedforward_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) if self.has_gemma4_moe: # gemma4 MoE: dense MLP + MoE experts in parallel mlp_output = self.post_feedforward_layernorm_1(hidden_states) # Router uses residual (pre-MLP hidden states) router_input = residual.reshape(-1, residual.shape[-1]) # Routing: norm -> scale -> proj -> softmax -> topk -> normalize -> per_expert_scale normed = self.moe_router_norm(router_input) normed = normed * self.moe_router_scale * self.moe_router_scalar_root router_logits = self.moe_gate(normed) routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) routing_weights, selected_experts = torch.topk(routing_weights, self.moe_top_k, dim=-1) routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6) routing_weights = (routing_weights * self.moe_per_expert_scale[selected_experts]).to(router_input.dtype) if self.export_moe: expert_input = self.pre_feedforward_layernorm_2(router_input) expert_output = self.custom_moe(expert_input, routing_weights, selected_experts) else: # Expert computation expert_input = self.pre_feedforward_layernorm_2(router_input) batch_size, sequence_length = residual.shape[0], residual.shape[1] hidden_dim = residual.shape[-1] expert_output = torch.zeros_like(router_input) expert_mask = F.one_hot(selected_experts, num_classes=self.moe_num_experts).permute(2, 1, 0) expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit: expert_idx = expert_idx[0] idx, top_x = torch.where(expert_mask[expert_idx]) current_state = expert_input[top_x] current_hidden = self.experts[expert_idx](current_state) * routing_weights[top_x, idx, None] expert_output.index_add_(0, top_x, current_hidden.to(expert_output.dtype)) expert_output = expert_output.reshape(residual.shape) expert_output = self.post_feedforward_layernorm_2(expert_output) hidden_states = mlp_output + expert_output hidden_states = self.post_feedforward_layernorm(hidden_states) hidden_states = residual + hidden_states elif hasattr(self, 'scale_depth'): # minicpm hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)) residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states * (self.scale_depth / math.sqrt(self.num_hidden_layers)) else: # general hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states # gemma4 PLE (Per-Layer Embeddings) if hasattr(self, 'per_layer_input_gate') and self.per_layer_input_gate is not None: per_layer_input = getattr(self, '_per_layer_input', None) if per_layer_input is not None: residual = hidden_states hidden_states = self.per_layer_input_gate(hidden_states) hidden_states = self.act_fn(hidden_states) hidden_states = hidden_states * per_layer_input hidden_states = self.per_layer_projection(hidden_states) hidden_states = self.post_per_layer_input_norm(hidden_states) hidden_states = residual + hidden_states # gemma4 layer_scalar if hasattr(self, 'layer_scalar') and self.layer_scalar is not None: hidden_states = hidden_states * self.layer_scalar return hidden_states class Lm(torch.nn.Module): def __init__(self, lm_, final_logit_softcapping=None): super().__init__() self.lm = lm_ self.final_logit_softcapping = final_logit_softcapping def forward(self, hidden_states): m_logits = self.lm(hidden_states) if self.final_logit_softcapping is not None: m_logits = m_logits / self.final_logit_softcapping m_logits = torch.tanh(m_logits) m_logits = m_logits * self.final_logit_softcapping return m_logits