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