Upload 9 files

README.md

@@ -1,136 +1,16 @@
----
-license: mit
-language:
-- zh
-- en
-datasets:
-- inclusionAI/Ling-Coder-SFT
-- amd/Instella-GSM8K-synthetic
-- tuanha1305/DeepSeek-R1-Distill
-- YeungNLP/school_math_0.25M
-- jingyaogong/minimind_dataset
-- Roman1111111/claude-sonnet-4.6-120000x
-- SnifferCaptain/z1-1m
-pipeline_tag: text-generation
----
-
-
-## 模型描述
-YModel3是SnifferCaptain训练的到目前为止（5/4/2026）最新的大语言模型。模型相比YModel2，支持了如混合思考、可调思考深度的功能，在回答的质量上有一定的进步。
-
-## 模型细节
-- 模型借鉴了MLA（ Deepseek， https://arxiv.org/pdf/2405.04434 ）的优化思路，与Gated Attention（ Qwen， https://arxiv.org/pdf/2505.06708 ）的优化思路，将ymodel2的PEGA2替换为容量更大，scaling更强的MLGA模块。显著增加了模型的参数效率。
-- 模型在FFN部分采用了SwiGLU。
-- 模型在每层RMSNorm后，增加了一层SEBlock。
-| key | value |
-|:--|:--|
-| 参数量（训练时） | 13.51M |
-| 层数 | 4 |
-| 隐藏层大小 | 512 |
-| 词表大小 | 6400 |
-| 前馈层激活函数 | SwiGLU |
-| 前馈层扩张维度 | 1024 |
-| 归一化层 | RMSNorm + SEBlock |
-| 注意力机制 | MLGA |
-| 注意力头数 | 4 |
-| 注意力头维度 | 64 |
-| KV潜在维度 | 128 |
-| RoPE嵌入维度 | 64 |
-
-## 训练细节
-- 模型全程采用与YModel2相近的SiMuon优化器训练，其中将NS迭代从2步增加到3步。SiMuon部分参数在预训练中，学习率为默认学习率的10倍。在其余阶段，学习率为默认学习率的66倍。AdamW部分学习率为默认参数的1倍。使用0.2*sqrt(max(fan in, fan out))进行学习率缩放。
-- 模型的tokenlizer与词来自MiniMind3-v（ https://github.com/jingyaogong/minimind ）
-- 在预训练阶段，模型在512长度的上下文充分使用**5B tokens**，1e-4带warmup的余弦退火到1e-5的学习率下，完成预训练。
-- 在全参微调阶段，模型在1024到4096上下文训练**2.5B tokens**，1e-5带warmup的余弦退火到1e-6的学习率下，完成后训练。最终的ppl为7.82（6400BPE词表长度）。
-- 预训练阶段采用的batch size为 65536 tokens/step。微调阶段采用 131072 tokens/step。
-
-## 补充细节
-- 由于预训练数据集原因，预训练模型在英文输入后有很大概率生成与各大跑分测试类似的四选一选择题格式
-- 关于调节思考深度，需要遵循特定的思考模板实现：
-```text
-<|im_start|>user
-[user content]<|im_end|>
-<|im_start|>assistant
-<think>juice = 1.14
-[thinking content]</think>
-[reply content]<|im_end|>
-```
-其中，juice值需要使用两位小数表示，且等号前后必须有空格，否则其模型性能会下降，其值过小或者过大可能会出现长度不符合预期的情况，这是因为训练数据的覆盖问题导致的。juice值的大小可以参考以下公式计算：
-$$\max\left(0.0,\ \log_2\left(\frac{\text{token\_count}}{128} + 1\right)\right)$$
-典型值：
-|juice|token count|
-|---|---|
-|0.59|64|
-|1.00|128|
-|2.00|384|
-|3.00|896|
-
-## 模型性能
-### 输出示例
-**关闭思考：**
-```text
-用户: 请介绍一下自己。
-助手: 我是由SnifferCaptain创建的语言模型，主要用于提供信息、解答问题和辅助学习。我的设计目标是通过高效的信息处理与用户交互来提升用户体验并增强用户的参与感。
-```
-**juice = 0.50：**
-```text
-用户: 请介绍一下自己。
-assistant: <think>juice = 0.5
-好的，用户让我简要介绍自己的身份和能力。首先，我需要确认自己是否具备一定的专业知识，但根据规则，只有在明确回答时才透露过多细节。所以这里应该保持简洁，不添加额外信息。接下来是自我反思，这属于简单问题，因为这是一个复杂的问题，可能需要多方面的分析。
-</think>
-作为AI助手，我的存在目的是提供准确、高效的信息与帮助。无论是在日常任务中还是在学习新技能或解决问题的能力上，我都以高效的方式处理各种请求。我的设计目标在于帮助用户解决他们的问题，无论是解答问题还是执行决策，都能提升效率。
-```
-
-**juice = 1.50：**
-```text
-user: 请介绍一下自己。
-assistant: <think>juice = 1.5
-嗯，用户让我介绍自己的个人简介，包括他的基本信息、背景和见解。首先，我需要确认用户的需求是什么。看起来他们可能是在学习中文还是其他类型的专业人士交流，但根据问题本身，应该保持简洁明了的语言。接下来要考虑用户的身份，可能是普通用户想了解我的能力范围或者兴趣爱好。不过作为AI助手，我不能直接回答这个问题，所以得用一般性的方式回应。同时，要确保不涉及任何敏感内容，只能提供准确的信息。另外，检查是否有遗漏的重要点，比如教育或工作坊的内容是否覆盖了不同领域的知识？这些都可以加以扩展可信度。最后，确保符合所有规则：使用���文，结构清晰，没有冗余分析，并且用自然的口语化表达。这样既符合动态响应深度的要求，又不会过于复杂。
-</think>
-您好！很高兴为您介绍我作为AI助手，以下是其基本信息及特点：
-
----
-
-### **核心信息**  
-1. **技术基础**：
-    *   **知识库：** 具备以下基础知识（如计算机科学/数据科学）与现有系统（如GPT-4），对人工智能（ML）、机器学习等概念感兴趣。例如：
-  - “人工智能”是当前最热门的技术之一，尤其在医疗诊断中占比达数百万倍；
-  - “深度学习模型”（CNN）则用于构建大规模数据集。
-     - “自然语言处理（NLP）”是当前最广泛使用的一种技术，特别适用于处理文本情感、语音识别、机器翻译等领域的应用场景。
-
-2. **应用场景：**
-        *   **智能客服：** 能够通过语音输入服务提供24/7解答疑问，解决客户投诉时出现的问题。
-        *   **金融风险管理：** 支持信用评分、欺诈检测、投资决策制定等金融风险评估指标，帮助银行预测市场波动趋势。
-        *   **多模态融合：** 利用大数据集进行实时数据处理和分析，实现自动化交易流程优化。
-        *   **虚拟现实 (VR)** 是目前最流行的智能家居解决方案之一，结合3D打印、ARIMA等全息影像设备，提升用户体验。
-
-3. **主要优势：**
-        *   **精准定位：** 基于大数据的算法推荐更优质的产品或服务。
-        *   **无人机驱动：** 为家庭中的小型飞行器，配备高效的灵活性训练软件，适合日常办公室环境。
-        *   **自动驾驶汽车：** 针对城市交通拥堵情况，减少事故发生率。
-        *   **能源效率：** 采用节能材料制成的建筑设计，降低碳排放。
-        *   **生物降解/食品回收：** 开发智能垃圾桶，提高公众环保意识。
-        *   **个性化服务：** 为用户定制一套个性化的旅游路线图，满足不同消费者的需求。
-
-4. **未来发展路径：**
-        *   **预计2030年前至2023年的科技创新：** 随着技术的进步和应用的深入，预计到2025年将会从零开始走向数字时代。这将极大地改变人们的生活方式，推动社会变革。
-        *   **全球化转型：** 在全球范围内开展跨国合作项目，促进国际间的沟通与协作，共同应对全球挑战。
-
----
-
-### **总结**  
-- **核心目标：** 成为当前最热门的技术，并强调关键价值——即如何有效利用资源、建立可持续发展的理念，以及持续学习和适应新技术的发展方向。
-- **结论：** 作为AI助手，我将专注于快速发展的行业应用，重塑人类社会的整体面貌。
-
----
-
-### **补充说明**  
-- **核心目标：** 掌握哪些技能（如编程、数据分析、网络安全、伦理道德等）。  
-- **行动建议：** 关注企业内部或政府部门的具体职责（如财务规划、政策执行、产品研发等）。  
-- **持续学习：** 不断更新最新技术、工具和技术，保持竞争力。
-
----
-
-这个信息完整展示了你的核心能力和潜力，希望对你有所启发。如果你有具体问题或需要进一步的帮助，随时告诉我！
-
-```
+# YModel3 Export Package
+
+这个目录是纯净发布包，**不依赖 `model/ymodel3.py`**。
+
+## 文件
+- `model.safetensors`
+- `config.json`
+- `tokenizer.json` / `tokenizer_config.json`
+- `ymodel3_eval.py`
+- `configuration_ymodel3.py`
+- `modeling_ymodel3.py`
+
+## 加载方式
+- `transformers.AutoConfig.from_pretrained(..., trust_remote_code=True)`
+- `transformers.AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True)`
+- 直接使用 `ymodel3_eval.py`

__init__.py

@@ -0,0 +1 @@
+from .modeling_ymodel3 import YForCausalLM3

config.json

@@ -0,0 +1,45 @@
+{
+  "dropout": 0.0,
+  "bos_token_id": 151644,
+  "eos_token_id": 151645,
+  "pad_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 512,
+  "num_hidden_layers": 4,
+  "max_position_embeddings": 4096,
+  "vocab_size": 6400,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 50000.0,
+  "rope_scaling": null,
+  "dtype": "bfloat16",
+  "self_distill": true,
+  "intermediate_size": 1024,
+  "expert_intermediate_size": 512,
+  "n_routed_experts": 0,
+  "moe_topk": 2,
+  "score_func": "softmax",
+  "n_shared_experts": 0,
+  "top_k_layer_dense": 4,
+  "aux_loss_alpha": 0.02,
+  "seq_aux": false,
+  "norm_topk_prob": true,
+  "noisy_expert": 0.0,
+  "moe_backend": "compact",
+  "router_bias_enabled": true,
+  "router_bias_update_rate": 0.001,
+  "router_bias_clamp": 5.0,
+  "num_heads": 4,
+  "mla_kv_lora_rank": 128,
+  "mla_qk_nope_head_dim": 64,
+  "mla_qk_rope_head_dim": 64,
+  "mla_attn_impl": "absorb",
+  "qkv_lora": false,
+  "model_type": "ynet3",
+  "architectures": [
+    "YForCausalLM3"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ymodel3.YConfig3",
+    "AutoModelForCausalLM": "modeling_ymodel3.YForCausalLM3"
+  }
+}

configuration_ymodel3.py

@@ -0,0 +1,48 @@
+from transformers import PretrainedConfig
+
+
+class YConfig3(PretrainedConfig):
+    model_type = "ynet3"
+
+    def __init__(self, **kwargs):
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.bos_token_id = kwargs.pop("bos_token_id", 151644)
+        self.eos_token_id = kwargs.pop("eos_token_id", 151645)
+        self.pad_token_id = kwargs.pop("pad_token_id", 151643)
+        self.hidden_act = kwargs.pop("hidden_act", "silu")
+        self.hidden_size = kwargs.pop("hidden_size", 768)
+        self.num_hidden_layers = kwargs.pop("num_hidden_layers", 8)
+        self.max_position_embeddings = kwargs.pop("max_position_embeddings", 8192)
+        self.vocab_size = kwargs.pop("vocab_size", 6400)
+        self.rms_norm_eps = kwargs.pop("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.pop("rope_theta", 5e4)
+        self.rope_scaling = kwargs.pop("rope_scaling", None)
+        self.dtype = kwargs.pop("dtype", "float32")
+        self.self_distill = kwargs.pop("self_distill", True)
+        self.intermediate_size = kwargs.pop("intermediate_size", 1536)
+        self.expert_intermediate_size = kwargs.pop("expert_intermediate_size", None) or self.intermediate_size
+        self.n_routed_experts = kwargs.pop("n_routed_experts", 0)
+        self.moe_topk = kwargs.pop("moe_topk", 2)
+        self.score_func = kwargs.pop("score_func", "softmax")
+        self.n_shared_experts = kwargs.pop("n_shared_experts", 0)
+        self.top_k_layer_dense = kwargs.pop("top_k_layer_dense", 1)
+        self.aux_loss_alpha = kwargs.pop("aux_loss_alpha", 0.02)
+        self.seq_aux = kwargs.pop("seq_aux", False)
+        self.norm_topk_prob = kwargs.pop("norm_topk_prob", True)
+        self.noisy_expert = kwargs.pop("noisy_expert", 0.0)
+        self.moe_backend = kwargs.pop("moe_backend", "compact")
+        self.router_bias_enabled = kwargs.pop("router_bias_enabled", True)
+        self.router_bias_update_rate = kwargs.pop("router_bias_update_rate", 1e-3)
+        self.router_bias_clamp = kwargs.pop("router_bias_clamp", 5.0)
+        self.num_heads = kwargs.pop("num_heads", 12)
+        self.mla_kv_lora_rank = kwargs.pop("mla_kv_lora_rank", 64)
+        self.mla_qk_nope_head_dim = kwargs.pop("mla_qk_nope_head_dim", 64)
+        self.mla_qk_rope_head_dim = kwargs.pop("mla_qk_rope_head_dim", 32)
+        self.mla_attn_impl = kwargs.pop("mla_attn_impl", "absorb")
+        self.qkv_lora = kwargs.pop("qkv_lora", False)
+        super().__init__(
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id,
+            **kwargs,
+        )

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31d5195ad56f084e67d681cbb4a2e9eedfcbea24f3889bb43c771c066c767d24
+size 33571344

modeling_ymodel3.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_ymodel3 import YConfig3
+from .ymodel3_eval import YModel3
+
+
+class YForCausalLM3(PreTrainedModel):
+    config_class = YConfig3
+    base_model_prefix = "model"
+
+    def __init__(self, config: YConfig3):
+        super().__init__(config)
+        self.model = YModel3(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+        self.lm_head.weight = value.weight
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def tie_weights(self):
+        self.model.embed_tokens.weight = self.lm_head.weight
+        return None
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+            "use_cache": use_cache,
+            "cache_position": kwargs.get("cache_position", None),
+            "position_ids": kwargs.get("position_ids", None),
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+        cache_position=None,
+        position_ids=None,
+        **kwargs,
+    ):
+        h, past_kvs, _, _ = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_ids=position_ids,
+        )
+        logits = self.lm_head(h)
+        return CausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=past_kvs,
+            hidden_states=(h,),
+        )

tokenizer_config.json

@@ -0,0 +1,335 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<|audio_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "17": {
+      "content": "<tts_pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "18": {
+      "content": "<tts_text_bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "19": {
+      "content": "<tts_text_eod>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "20": {
+      "content": "<tts_text_bos_single>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "21": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "22": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "23": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "24": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "25": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "26": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "27": {
+      "content": "<|buffer1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "28": {
+      "content": "<|buffer2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "29": {
+      "content": "<|buffer3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "30": {
+      "content": "<|buffer4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "31": {
+      "content": "<|buffer5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "32": {
+      "content": "<|buffer6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "33": {
+      "content": "<|buffer7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "34": {
+      "content": "<|buffer8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "35": {
+      "content": "<|buffer9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_pad|>",
+    "<tts_pad>",
+    "<tts_text_bos>",
+    "<tts_text_eod>",
+    "<tts_text_bos_single>"
+  ],
+  "bos_token": "<|im_start|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "unk_token": "<|endoftext|>",
+  "image_token": "<|image_pad|>",
+  "audio_token": "<|audio_pad|>",
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>",
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "chat_template": "{%- if tools %}{{- '<|im_start|>system\\n' }}{%- if messages[0].role == 'system' %}{{- messages[0].content + '\\n\\n' }}{%- endif %}{{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}{%- for tool in tools %}{{- \"\\n\" }}{{- tool | tojson }}{%- endfor %}{{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}{%- else %}{%- if messages[0].role == 'system' %}{{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}{%- endif %}{%- endif %}{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}{%- for message in messages[::-1] %}{%- set index = (messages|length - 1) - loop.index0 %}{%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}{%- set ns.multi_step_tool = false %}{%- set ns.last_query_index = index %}{%- endif %}{%- endfor %}{%- for message in messages %}{%- if message.content is string %}{%- set content = message.content %}{%- else %}{%- set content = '' %}{%- endif %}{%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}{{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}{%- elif message.role == \"assistant\" %}{%- set reasoning_content = '' %}{%- if message.reasoning_content is string %}{%- set reasoning_content = message.reasoning_content %}{%- else %}{%- if '</think>' in content %}{%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}{%- set content = content.split('</think>')[-1].lstrip('\\n') %}{%- endif %}{%- endif %}{%- if true %}{{- '<|im_start|>' + message.role + '\\n<think>' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}{%- endif %}{%- if message.tool_calls %}{%- for tool_call in message.tool_calls %}{%- if (loop.first and content) or (not loop.first) %}{{- '\\n' }}{%- endif %}{%- if tool_call.function %}{%- set tool_call = tool_call.function %}{%- endif %}{{- '<tool_call>\\n{\"name\": \"' }}{{- tool_call.name }}{{- '\", \"arguments\": ' }}{%- if tool_call.arguments is string %}{{- tool_call.arguments }}{%- else %}{{- tool_call.arguments | tojson }}{%- endif %}{{- '}\\n</tool_call>' }}{%- endfor %}{%- endif %}{{- '<|im_end|>\\n' }}{%- elif message.role == \"tool\" %}{%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}{{- '<|im_start|>user' }}{%- endif %}{{- '\\n<tool_response>\\n' }}{{- content }}{{- '\\n</tool_response>' }}{%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}{{- '<|im_end|>\\n' }}{%- endif %}{%- endif %}{%- endfor %}{%- if add_generation_prompt %}{{- '<|im_start|>assistant\\n' }}{%- set juice_value = thinking_juice if thinking_juice is defined else 2.00 %}{%- set juice_str = '%.2f' | format(juice_value) %}{%- if open_thinking is defined and open_thinking %}{{- '<think>juice = ' + juice_str + '\\n' }}{%- else %}{{- '<think>juice = ' + juice_str + '\\n</think>\\n\\n' }}{%- endif %}{%- endif %}",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

ymodel3_eval.py

@@ -0,0 +1,755 @@
+"""Self-contained ymodel3 inference module.
+
+Only depends on: torch, safetensors.
+No dependency on kernel.*, model.ymodel3, transformers.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file as load_safetensors
+
+
+# ── Config ──────────────────────────────────────────────────────────
+
+
+class YConfig3:
+    model_type = "ynet3"
+
+    def __init__(self, **kwargs):
+        self.dropout = kwargs.get("dropout", 0.0)
+        self.bos_token_id = kwargs.get("bos_token_id", 151644)
+        self.eos_token_id = kwargs.get("eos_token_id", 151645)
+        self.pad_token_id = kwargs.get("pad_token_id", 151643)
+        self.hidden_act = kwargs.get("hidden_act", "silu")
+        self.hidden_size = kwargs.get("hidden_size", 768)
+        self.num_hidden_layers = kwargs.get("num_hidden_layers", 8)
+        self.max_position_embeddings = kwargs.get("max_position_embeddings", 8192)
+        self.vocab_size = kwargs.get("vocab_size", 6400)
+        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.get("rope_theta", 5e4)
+        self.rope_scaling = kwargs.get("rope_scaling", None)
+        self.dtype = kwargs.get("dtype", "float32")
+        self.self_distill = kwargs.get("self_distill", True)
+        self.intermediate_size = kwargs.get("intermediate_size", 1536)
+        self.expert_intermediate_size = kwargs.get("expert_intermediate_size", None) or self.intermediate_size
+        self.n_routed_experts = kwargs.get("n_routed_experts", 0)
+        self.moe_topk = kwargs.get("moe_topk", 2)
+        self.score_func = kwargs.get("score_func", "softmax")
+        self.n_shared_experts = kwargs.get("n_shared_experts", 0)
+        self.top_k_layer_dense = kwargs.get("top_k_layer_dense", 1)
+        self.aux_loss_alpha = kwargs.get("aux_loss_alpha", 0.02)
+        self.seq_aux = kwargs.get("seq_aux", False)
+        self.norm_topk_prob = kwargs.get("norm_topk_prob", True)
+        self.noisy_expert = kwargs.get("noisy_expert", 0.0)
+        self.moe_backend = kwargs.get("moe_backend", "compact")
+        self.router_bias_enabled = kwargs.get("router_bias_enabled", True)
+        self.router_bias_update_rate = kwargs.get("router_bias_update_rate", 1e-3)
+        self.router_bias_clamp = kwargs.get("router_bias_clamp", 5.0)
+        self.num_heads = kwargs.get("num_heads", 12)
+        self.mla_kv_lora_rank = kwargs.get("mla_kv_lora_rank", 64)
+        self.mla_qk_nope_head_dim = kwargs.get("mla_qk_nope_head_dim", 64)
+        self.mla_qk_rope_head_dim = kwargs.get("mla_qk_rope_head_dim", 32)
+        self.mla_attn_impl = kwargs.get("mla_attn_impl", "absorb")
+        self.qkv_lora = kwargs.get("qkv_lora", False)
+
+    @property
+    def head_dim(self) -> int:
+        return self.mla_qk_nope_head_dim + self.mla_qk_rope_head_dim
+
+    def scale_lvl(self, lvl: int = 0):
+        if lvl == 0:
+            self.hidden_size = 1024
+            self.num_hidden_layers = 8
+            self.num_heads = 8
+            self.mla_kv_lora_rank = 256
+            self.mla_qk_nope_head_dim = 192
+            self.mla_qk_rope_head_dim = 64
+            self.intermediate_size = 2048
+            self.expert_intermediate_size = 512
+            self.n_routed_experts = 16
+            self.moe_topk = 1
+            self.n_shared_experts = 0
+            self.top_k_layer_dense = 1
+            self.router_bias_update_rate = 1e-3
+        elif lvl == -1:
+            self.hidden_size = 768
+            self.num_hidden_layers = 8
+            self.num_heads = 6
+            self.mla_kv_lora_rank = 128
+            self.mla_qk_nope_head_dim = 64
+            self.mla_qk_rope_head_dim = 64
+            self.intermediate_size = 1536
+            self.expert_intermediate_size = 768
+            self.n_routed_experts = 0
+            self.moe_topk = 2
+            self.n_shared_experts = 0
+            self.top_k_layer_dense = 8
+        elif lvl == -2:
+            self.hidden_size = 512
+            self.num_hidden_layers = 4
+            self.num_heads = 4
+            self.mla_kv_lora_rank = 128
+            self.mla_qk_nope_head_dim = 64
+            self.mla_qk_rope_head_dim = 32
+            self.intermediate_size = 1024
+            self.expert_intermediate_size = 512
+            self.n_routed_experts = 0
+            self.moe_topk = 2
+            self.n_shared_experts = 0
+            self.top_k_layer_dense = 4
+        else:
+            raise ValueError(f"invalid ymodel3 scale level: {lvl}")
+        return self
+
+    @classmethod
+    def from_json_file(cls, path: str) -> "YConfig3":
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls(**data)
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "YConfig3":
+        return cls(**data)
+
+
+# ── Basic modules ──────────────────────────────────────────────────
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (out * self.weight.float()).to(x.dtype)
+
+
+class SEBlock(nn.Module):
+    def __init__(self, dim: int, reduction: int = 16, act: Optional[nn.Module] = None):
+        super().__init__()
+        reduction = max(reduction, dim // reduction)
+        self.se = nn.Sequential(
+            nn.Linear(dim, reduction, bias=False),
+            act or nn.SiLU(),
+            nn.Linear(reduction, dim, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.se(x)
+
+
+# ── RoPE helpers ──────────────────────────────────────────────────
+
+
+def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor:
+    if low == high:
+        high += 0.001
+    linear = (torch.arange(dim, dtype=torch.float32) - low) / (high - low)
+    return torch.clamp(linear, 0.0, 1.0)
+
+
+def _yarn_correction_dim(num_rotations: float, dim: int, theta: float, max_position_embeddings: int) -> float:
+    return dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) / (2 * math.log(theta))
+
+
+def precompute_freqs_cis(
+    dim: int,
+    end: int,
+    theta: float,
+    rope_scaling: Optional[dict] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    attention_factor = 1.0
+    if rope_scaling and str(rope_scaling.get("type", "yarn")).lower() == "yarn":
+        factor = float(rope_scaling.get("factor", 1.0))
+        if factor > 1.0:
+            original = int(rope_scaling.get("original_max_position_embeddings", end))
+            beta_fast = float(rope_scaling.get("beta_fast", 32.0))
+            beta_slow = float(rope_scaling.get("beta_slow", 1.0))
+            low = math.floor(_yarn_correction_dim(beta_fast, dim, theta, original))
+            high = math.ceil(_yarn_correction_dim(beta_slow, dim, theta, original))
+            ramp = _yarn_linear_ramp(low, high, dim // 2)
+            freqs = freqs / factor * (1.0 - ramp) + freqs * ramp
+            attention_factor = float(rope_scaling.get("attention_factor", 1.0))
+    t = torch.arange(end)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor
+    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor
+    return freqs_cos, freqs_sin
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    return torch.cat((-x[..., x.shape[-1] // 2 :], x[..., : x.shape[-1] // 2]), dim=-1)
+
+
+def apply_rope_to_single(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    if cos.dim() == 2:
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+    elif cos.dim() == 3:
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+# ── Attention ──────────────────────────────────────────────────────
+
+
+class MLGA(nn.Module):
+    """Multihead Latent Gated Attention"""
+
+    def __init__(self, config: YConfig3, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.dropout = config.dropout
+        self.kv_lora_rank = config.mla_kv_lora_rank
+        self.qk_nope_head_dim = config.mla_qk_nope_head_dim
+        self.qk_rope_head_dim = config.mla_qk_rope_head_dim
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.attn_impl = config.mla_attn_impl
+        self.softmax_scale = self.qk_head_dim ** -0.5
+        self.out_dim = self.num_heads * self.kv_lora_rank
+
+        self.wq = nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+        self.wkv_a = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
+        self.kv_norm = RMSNorm(self.kv_lora_rank, config.rms_norm_eps)
+        self.wkv_b = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, bias=False)
+        self.z_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False)
+        self.o_proj = nn.Linear(self.out_dim, self.hidden_size, bias=False)
+
+    def _project_q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_len, _ = x.shape
+        q = self.wq(x)
+        q = q.reshape(bsz, seq_len, self.num_heads, self.qk_head_dim)
+        return q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+    def _project_kv(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        raw = self.wkv_a(x)
+        c_kv, k_pe = raw.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        c_kv = self.kv_norm(c_kv)
+        k_pe = apply_rope_to_single(k_pe.unsqueeze(1), cos, sin).permute(0, 2, 1, 3)
+        return c_kv, k_pe
+
+    def _explicit_kv(self, c_kv: torch.Tensor, k_pe: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_len, _ = c_kv.shape
+        k_nope = self.wkv_b(c_kv).reshape(bsz, seq_len, self.num_heads, self.qk_nope_head_dim)
+        k = torch.cat([k_nope, k_pe.expand(-1, -1, self.num_heads, -1)], dim=-1)
+        v = c_kv.unsqueeze(2).expand(-1, -1, self.num_heads, -1)
+        return k, v
+
+    def _attention_mask(self, attention_mask: Optional[torch.Tensor], bsz: int, seq_len: int, total_len: int):
+        if attention_mask is None:
+            return None
+        if attention_mask.shape[-1] != total_len:
+            attention_mask = attention_mask[..., -total_len:]
+        mask = attention_mask.reshape(bsz, 1, 1, total_len).bool()
+        return mask.expand(bsz, self.num_heads, seq_len, total_len)
+
+    def _forward_sdpa(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        c_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        z: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, seq_len, _, _ = q_nope.shape
+        total_len = c_kv.shape[1]
+        k, v = self._explicit_kv(c_kv, k_pe)
+        q = torch.cat([q_nope, q_pe], dim=-1).permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len)
+        is_causal = attention_mask is None and seq_len == total_len
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout if self.training else 0.0,
+            is_causal=is_causal,
+            scale=self.softmax_scale,
+        )
+        out = out.permute(0, 2, 1, 3).reshape(bsz, seq_len, self.out_dim)
+        out = out * torch.sigmoid(z)
+        return self.o_proj(out)
+
+    def _forward_absorb(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        c_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        z: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, seq_len, _, _ = q_nope.shape
+        total_len = c_kv.shape[1]
+        w = self.wkv_b.weight.reshape(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank)
+        q_nope_c = torch.einsum("bshd,hdc->bshc", q_nope, w)
+        scores = torch.einsum("bshc,btc->bsht", q_nope_c, c_kv)
+        scores = scores + torch.einsum("bshr,btr->bsht", q_pe, k_pe.squeeze(2))
+        scores = scores * self.softmax_scale
+
+        causal = torch.full((seq_len, seq_len), float("-inf"), device=scores.device, dtype=scores.dtype)
+        causal = torch.triu(causal, diagonal=1).reshape(1, seq_len, 1, seq_len)
+        scores = scores + F.pad(causal, (total_len - seq_len, 0), value=0.0)
+        if attention_mask is not None:
+            if attention_mask.shape[-1] != total_len:
+                attention_mask = attention_mask[..., -total_len:]
+            scores = scores + (1.0 - attention_mask.reshape(bsz, 1, 1, total_len).float()) * -1e9
+        probs = torch.softmax(scores.float(), dim=-1).to(q_nope.dtype)
+        out = torch.einsum("bsht,btc->bshc", probs, c_kv).reshape(bsz, seq_len, self.out_dim)
+        out = out * torch.sigmoid(z)
+        return self.o_proj(out)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        bsz, seq_len, _ = x.shape
+        cos, sin = position_embeddings
+        if cos.dim() == 2:
+            cos = cos[:seq_len, : self.qk_rope_head_dim]
+            sin = sin[:seq_len, : self.qk_rope_head_dim]
+        else:
+            cos = cos[:, :seq_len, : self.qk_rope_head_dim]
+            sin = sin[:, :seq_len, : self.qk_rope_head_dim]
+        q_nope, q_pe = self._project_q(x)
+        q_pe = apply_rope_to_single(q_pe.permute(0, 2, 1, 3), cos, sin).permute(0, 2, 1, 3)
+        c_kv, k_pe = self._project_kv(x, cos, sin)
+        z = self.z_proj(x)
+
+        if past_key_values is not None:
+            past_c, past_pe = past_key_values
+            c_kv = torch.cat([past_c, c_kv], dim=1)
+            k_pe = torch.cat([past_pe, k_pe], dim=1)
+        new_past = (c_kv, k_pe) if use_cache else None
+
+        if self.attn_impl == "naive":
+            out = self._forward_sdpa(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
+        else:
+            out = self._forward_absorb(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
+        out = F.dropout(out, p=self.dropout, training=self.training)
+        return out, new_past
+
+
+# ── FFN / MoE ──────────────────────────────────────────────────────
+
+
+_ACT_FNS = {
+    "silu": F.silu,
+    "swish": F.silu,
+    "relu": F.relu,
+    "gelu": lambda x: F.gelu(x, approximate="tanh"),
+    "sigmoid": torch.sigmoid,
+}
+
+_ACT_MODULES = {
+    "silu": nn.SiLU,
+    "swish": nn.SiLU,
+    "relu": nn.ReLU,
+    "gelu": lambda: nn.GELU(approximate="tanh"),
+    "sigmoid": nn.Sigmoid,
+}
+
+
+class DenseFFN(nn.Module):
+    def __init__(self, config: YConfig3, intermediate_size: Optional[int] = None):
+        super().__init__()
+        inter = intermediate_size or config.intermediate_size
+        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)
+        self.hidden_act = config.hidden_act
+        self.act = _ACT_FNS.get(config.hidden_act, F.silu)
+        self.dropout = config.dropout
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, gate = self.up_proj(x), self.gate_proj(x)
+        up = self.act(gate) * up
+        up = F.dropout(up, p=self.dropout, training=self.training)
+        return self.down_proj(up)
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: YConfig3):
+        super().__init__()
+        self.n_routed_experts = config.n_routed_experts
+        self.topk = min(config.moe_topk, max(1, config.n_routed_experts))
+        self.score_func = config.score_func
+        self.norm_topk_prob = config.norm_topk_prob
+        self.aux_loss_alpha = config.aux_loss_alpha
+        self.seq_aux = config.seq_aux
+        self.router_bias_enabled = config.router_bias_enabled
+        self.router_bias_update_rate = config.router_bias_update_rate
+        self.router_bias_clamp = config.router_bias_clamp
+        self.weight = nn.Linear(int(config.hidden_size), int(self.n_routed_experts), bias=False)
+        if self.router_bias_enabled:
+            self.register_buffer("router_bias", torch.zeros(self.n_routed_experts), persistent=True)
+        else:
+            self.register_buffer("router_bias", None, persistent=False)
+
+    def forward(self, x: torch.Tensor, aux_mask: Optional[torch.Tensor] = None):
+        bsz, seq_len, hidden = x.shape
+        flat = x.reshape(-1, hidden)
+        route_logits = self.weight(flat)
+        if self.score_func == "softmax":
+            route_scores = torch.softmax(route_logits.float(), dim=-1).to(x.dtype)
+        elif self.score_func == "sigmoid":
+            route_scores = torch.sigmoid(route_logits.float()).to(x.dtype)
+        else:
+            raise ValueError(f"unsupported MoE score_func: {self.score_func}")
+
+        choice_scores = route_scores
+        if self.router_bias is not None:
+            choice_scores = choice_scores + self.router_bias.to(dtype=choice_scores.dtype).unsqueeze(0)
+
+        topk_idx = torch.topk(choice_scores, k=self.topk, dim=-1, sorted=False).indices
+        topk_weight = route_scores.gather(1, topk_idx)
+        if self.topk > 1 and self.norm_topk_prob:
+            denom = topk_weight.float().sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = (topk_weight.float() / denom).to(x.dtype)
+
+        aux_loss = x.new_zeros((), dtype=x.dtype)
+        return (
+            topk_idx.reshape(bsz, seq_len, self.topk),
+            topk_weight.reshape(bsz, seq_len, self.topk),
+            aux_loss,
+        )
+
+
+def _torch_moe_swiglu(
+    x: torch.Tensor,
+    topk_idx: torch.Tensor,
+    topk_weight: torch.Tensor,
+    w_up: torch.Tensor,
+    w_down: torch.Tensor,
+    activation: str = "silu",
+) -> torch.Tensor:
+    """Pure PyTorch MoE SwiGLU forward (inference only, no noisy_expert)."""
+    original_shape = x.shape
+    x_flat = x.reshape(-1, x.shape[-1])
+    idx = topk_idx.reshape(x_flat.shape[0], -1)
+    weight = topk_weight.reshape(x_flat.shape[0], -1)
+    y = torch.zeros_like(x_flat)
+    n_experts = w_up.shape[0]
+    inter = w_down.shape[-1]
+    act_fn = _ACT_FNS.get(activation, F.silu)
+    for expert_id in range(n_experts):
+        token_pos, choice_pos = torch.where(idx == expert_id)
+        if token_pos.numel() == 0:
+            continue
+        inp = x_flat[token_pos]
+        uv = F.linear(inp, w_up[expert_id])
+        up, gate = uv.split(inter, dim=-1)
+        hidden = act_fn(gate) * up
+        out = F.linear(hidden, w_down[expert_id])
+        route_w = weight[token_pos, choice_pos].unsqueeze(-1)
+        y.index_add_(0, token_pos, out * route_w)
+    return y.reshape(original_shape)
+
+
+class YMoE(nn.Module):
+    """Pure PyTorch eval MoE (no Triton dependency)."""
+
+    def __init__(self, config: YConfig3, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        self.hidden_size = config.hidden_size
+        self.expert_intermediate_size = config.expert_intermediate_size
+        self.intermediate_size = self.expert_intermediate_size
+        self.n_routed_experts = config.n_routed_experts
+        self.use_moe = self.n_routed_experts > 0 and layer_id >= config.top_k_layer_dense
+        self.noisy_expert = config.noisy_expert
+        if not self.use_moe:
+            self.dense = DenseFFN(config)
+            self.gate = None
+            self.w_up = None
+            self.w_down = None
+            return
+        self.dense = None
+        self.gate = MoEGate(config)
+        self.w_up = nn.Parameter(torch.empty(self.n_routed_experts, 2 * self.expert_intermediate_size, self.hidden_size))
+        self.w_down = nn.Parameter(torch.empty(self.n_routed_experts, self.hidden_size, self.expert_intermediate_size))
+        nn.init.kaiming_uniform_(self.w_up, a=math.sqrt(5))
+        nn.init.kaiming_uniform_(self.w_down, a=math.sqrt(5))
+
+    def forward(self, x: torch.Tensor, aux_mask: Optional[torch.Tensor] = None):
+        if not self.use_moe:
+            return self.dense(x), None
+        topk_idx, topk_weight, aux_loss = self.gate(x, aux_mask)
+        y = _torch_moe_swiglu(x, topk_idx, topk_weight, self.w_up, self.w_down, activation="silu")
+        return y, aux_loss
+
+
+# ── Transformer block ──────────────────────────────────────────────
+
+
+class YBlock3(nn.Module):
+    def __init__(self, config: YConfig3, layer_id: int):
+        super().__init__()
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attn = MLGA(config, layer_id)
+        self.ffn = YMoE(config, layer_id)
+        act_module = _ACT_MODULES.get(config.hidden_act, nn.SiLU)
+        self.se1 = SEBlock(config.hidden_size, act=act_module() if isinstance(act_module, type) else act_module())
+        self.se2 = SEBlock(config.hidden_size, act=nn.SiLU())
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        past_key_values=None,
+        use_cache: bool = False,
+        attention_mask: Optional[torch.Tensor] = None,
+        aux_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        x0 = self.se1(self.input_layernorm(x))
+        attn_out, past = self.attn(
+            x0,
+            position_embeddings,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            use_cache=use_cache,
+        )
+        x = x + attn_out
+        x0 = self.se2(self.post_attention_layernorm(x))
+        ffn_out, aux_loss = self.ffn(x0, aux_mask)
+        x = x + ffn_out
+        return x, past, aux_loss
+
+
+# ── Full model ────────────────────────────────────────────────────
+
+
+class YModel3(nn.Module):
+    def __init__(self, config: YConfig3):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.num_layers = config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = config.dropout
+        self.use_self_distill = config.self_distill
+        self.layers = nn.ModuleList([YBlock3(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        freqs_cos, freqs_sin = precompute_freqs_cis(
+            dim=config.mla_qk_rope_head_dim,
+            end=config.max_position_embeddings,
+            theta=config.rope_theta,
+            rope_scaling=config.rope_scaling,
+        )
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[list] = None,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        bsz, seq_len = input_ids.shape
+        if use_cache and past_key_values is None:
+            past_key_values = [None] * self.num_layers
+        if cache_position is None:
+            if past_key_values is not None and past_key_values[0] is not None:
+                past_seen = past_key_values[0][0].shape[1]
+            else:
+                past_seen = 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, device=input_ids.device)
+
+        x = F.dropout(self.embed_tokens(input_ids), p=self.dropout, training=self.training)
+        if position_ids is None:
+            position_ids = cache_position
+        position_embeddings = (self.freqs_cos[position_ids].to(x.device), self.freqs_sin[position_ids].to(x.device))
+        aux_mask = None
+        new_past = [] if use_cache else None
+        aux_loss = None
+
+        for i, layer in enumerate(self.layers):
+            past = past_key_values[i] if past_key_values is not None else None
+            x, layer_past, layer_aux = layer(
+                x,
+                position_embeddings=position_embeddings,
+                past_key_values=past,
+                attention_mask=attention_mask,
+                use_cache=use_cache,
+                aux_mask=aux_mask,
+            )
+            if use_cache:
+                new_past.append(layer_past)
+            if self.training and layer_aux is not None:
+                aux_loss = layer_aux if aux_loss is None else aux_loss + layer_aux
+
+        return self.norm(x), new_past, None, aux_loss
+
+
+class _InferenceOutput:
+    """Simple container for model outputs (replaces transformers CausalLMOutputWithPast)."""
+
+    __slots__ = ("last_hidden_state", "logits", "past_key_values", "dist_loss", "aux_loss")
+
+    def __init__(self):
+        self.last_hidden_state = None
+        self.logits = None
+        self.past_key_values = None
+        self.dist_loss = None
+        self.aux_loss = None
+
+    def __setitem__(self, key, value):
+        setattr(self, key, value)
+
+
+class YForCausalLM3(nn.Module):
+    """Pure PyTorch CausalLM wrapper for ymodel3 inference (no transformers dependency)."""
+
+    config_class = YConfig3
+
+    def __init__(self, config: Optional[YConfig3] = None):
+        super().__init__()
+        self.config = config or YConfig3()
+        self.model = YModel3(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.OUT = _InferenceOutput()
+        dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(self.config.dtype)
+        if dtype is not None:
+            self.to(dtype)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[list] = None,
+        use_cache: bool = False,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        h, past_kvs, dist_loss, aux_loss = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_ids=kwargs.get("position_ids", None),
+        )
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(h[:, slice_indices, :])
+        self.OUT.__setitem__("last_hidden_state", h)
+        self.OUT.__setitem__("logits", logits)
+        self.OUT.__setitem__("past_key_values", past_kvs)
+        self.OUT.__setitem__("dist_loss", dist_loss)
+        self.OUT.__setitem__("aux_loss", aux_loss)
+        return self.OUT
+
+    def generate(
+        self,
+        inputs,
+        attention_mask=None,
+        max_new_tokens=8192,
+        temperature=0.85,
+        top_p=0.85,
+        top_k=50,
+        eos_token_id=None,
+        streamer=None,
+        use_cache=True,
+        num_return_sequences=1,
+        do_sample=True,
+        repetition_penalty=1.0,
+        **kwargs,
+    ):
+        input_ids = kwargs.get("input_ids", inputs).repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
+        past_key_values = None
+        if streamer:
+            streamer.put(input_ids.cpu())
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                if use_cache and past_key_values is not None:
+                    outputs = self.forward(input_ids[:, -1:], None, past_key_values, use_cache=use_cache)
+                else:
+                    outputs = self.forward(input_ids, attention_mask, past_key_values, use_cache=use_cache)
+                logits = outputs.logits[:, -1, :] / temperature
+                if repetition_penalty != 1.0:
+                    for i in range(input_ids.shape[0]):
+                        logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+                if top_k > 0:
+                    logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float("inf")
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                    mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
+                    logits[mask.scatter(1, sorted_indices, mask)] = -float("inf")
+                next_token = torch.multinomial(torch.softmax(logits, dim=-1), 1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+                input_ids = torch.cat([input_ids, next_token], dim=-1)
+                past_key_values = outputs.past_key_values if use_cache else None
+                if streamer:
+                    streamer.put(next_token.cpu())
+                if eos_token_id and (next_token == eos_token_id).any():
+                    break
+        if streamer:
+            streamer.end()
+        return input_ids
+
+
+# ── Loading utilities ──────────────────────────────────────────────
+
+
+def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]:
+    path = Path(path)
+    if path.is_dir():
+        safetensors_path = path / "model.safetensors"
+        bin_path = path / "pytorch_model.bin"
+        if safetensors_path.exists():
+            path = safetensors_path
+        elif bin_path.exists():
+            path = bin_path
+        else:
+            raise FileNotFoundError(f"no model.safetensors or pytorch_model.bin found in {path}")
+    if path.suffix == ".safetensors":
+        return load_safetensors(str(path), device="cpu")
+    return torch.load(path, map_location="cpu", weights_only=True)
+
+
+def load_ymodel3_eval(path: Union[str, Path], config: Optional[YConfig3] = None, strict: bool = True) -> YForCausalLM3:
+    if config is None:
+        config_path = Path(path) / "config.json" if Path(path).is_dir() else Path(path).with_name("config.json")
+        if not config_path.exists():
+            raise FileNotFoundError("config is required when config.json is not next to the checkpoint")
+        config = YConfig3.from_json_file(str(config_path))
+    model = YForCausalLM3(config)
+    state = _load_state_dict(path)
+    model.load_state_dict(state, strict=strict)
+    model.eval()
+    return model
+
+
+# ── Backward-compatible aliases ────────────────────────────────────
+
+YModel3Eval = YModel3
+YForCausalLM3Eval = YForCausalLM3