模型定义 (models/)

# 1. vLLM启动时调用插件注册（vllm/plugins/）
# 扫描已安装的插件，调用register函数

# 2. vllm-ascend插件注册（vllm_ascend/platform.py）
class NPUPlatform(Platform):
    def register():
        return NPUPlatform

# 3. 模型注册（vllm_ascend/models/__init__.py）
def register_model():
    ModelRegistry.register_model(
        "KimiLinearForCausalLM",
        "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
    )

# 4. vLLM查找模型（vllm/v1/engine/llm_engine.py）
model_cls_path = ModelRegistry.get_model_class(model_name)

# 1. 模型注册（vllm_ascend/models/__init__.py）
def register_model():
    ModelRegistry.register_model("模型名称", "模块路径:类名")

# 2. ModelRegistry存储（vllm/models/registry.py）
class ModelRegistry:
    _models = {}  # 存储模型名称到类路径的映射
    
    @staticmethod
    def get_model_class(model_name):
        return _models.get(model_name)

# 3. Worker查找并加载模型（vllm/v1/worker/gpu_worker.py）
class Worker:
    def load_model(self, model_name):
        model_cls_path = ModelRegistry.get_model_class(model_name)
        # 动态导入并实例化
        model = import_and_instantiate(model_cls_path)
        return model

# 4. NPU上前向传播（vllm_ascend/worker/model_runner_v1.py）
class NPUModelRunner:
    def execute(self, input_ids):
        return self.model(input_ids)

# 1. vLLM Engine初始化（vllm/v1/engine/llm_engine.py）
class LLMEngine:
    def __init__(self, vllm_config):
        # 初始化插件，调用register_model
        self.executor = Executor(vllm_config)

# 2. 插件注册（vllm_ascend/platform.py）
class NPUPlatform(Platform):
    @staticmethod
    def register():
        return NPUPlatform

# 3. 模型注册（vllm_ascend/models/__init__.py）
def register_model():
    ModelRegistry.register_model("模型名", "模块:类")

# 4. Worker加载模型（vllm/v1/worker/gpu_worker.py）
class Worker:
    def init_device(self):
        model_cls = ModelRegistry.get_model_class(model_name)
        self.model = model_cls(config)

# 5. NPU执行推理（vllm_ascend/worker/model_runner_v1.py）
class NPUModelRunner:
    def execute(self, input_ids):
        return self.model(input_ids)

from vllm import ModelRegistry

# 这是模型注册的入口函数
# vLLM启动时会自动调用此函数
def register_model():
    # 注册Kimi K3因果语言模型
    # 参数1: 模型名称（用于HuggingFace配置匹配）
    # 参数2: 模型类的完整路径（用于动态导入）
    ModelRegistry.register_model(
        "KimiLinearForCausalLM",  # 模型名称
        "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",  # 类路径
    )

# KimiLinearForCausalLM - 主模型
# 用于文本生成任务
ModelRegistry.register_model(
    "KimiLinearForCausalLM",
    "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
)

# KimiK3ForCausalLM - 兼容性别名
# 用于旧版本checkpoint的兼容
ModelRegistry.register_model(
    "KimiK3ForCausalLM",
    "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
)

# KimiK3ForConditionalGeneration - 条件生成模型
# 用于多模态任务（如图文理解）
ModelRegistry.register_model(
    "KimiK3ForConditionalGeneration",
    "vllm_ascend.models.kimi_k3:AscendKimiK3ForConditionalGeneration",
)

# DeepseekV4ForCausalLM - DeepSeek V4主模型
# 支持MLA多头潜在注意力
ModelRegistry.register_model(
    "DeepseekV4ForCausalLM",
    "vllm_ascend.models.deepseek_v4.model:AscendDeepseekV4ForCausalLM",
)

# DeepSeekV4MTPModel - 多token预测模型
# 用于投机解码的草稿模型
ModelRegistry.register_model(
    "DeepSeekV4MTPModel",
    "vllm_ascend.models.deepseek_v4.mtp:DeepSeekV4MTP",
)

# DSparkDraftModel - DSpark草稿模型
# 用于分布式投机解码
ModelRegistry.register_model(
    "DSparkDraftModel",
    "vllm_ascend.models.deepseek_v4.dspark:DSparkDeepseekV4ForCausalLM",
)

# MiniMaxM3SparseForCausalLM - 稀疏模型
# 使用MoE混合专家架构
ModelRegistry.register_model(
    "MiniMaxM3SparseForCausalLM",
    "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForCausalLM",
)

# MiniMaxM3SparseForConditionalGeneration - 条件生成
# 用于多模态任务
ModelRegistry.register_model(
    "MiniMaxM3SparseForConditionalGeneration",
    "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForConditionalGeneration",
)

# DeepSeekMTPModel - DeepSeek多token预测
# 用于Eagle3投机解码
ModelRegistry.register_model(
    "DeepSeekMTPModel",
    "vllm_ascend.models.deepseek_mtp:AscendDeepSeekMTP",
)

# Qwen3DSparkModel - Qwen3 DSpark模型
# 用于Qwen3的分布式投机解码
ModelRegistry.register_model(
    "Qwen3DSparkModel",
    "vllm_ascend.models.qwen3_dspark:AscendQwen3DSparkForCausalLM",
)

# DFlash2DraftModel - DFlash2草稿模型
# 用于Qwen3的Flash投机解码
ModelRegistry.register_model(
    "DFlash2DraftModel",
    "vllm_ascend.models.qwen3_dflash2:DFlash2Qwen3ForCausalLM",
)

# Eagle3LlamaForCausalLM - LLaMA Eagle3
# 用于LLaMA模型的Eagle3投机解码
ModelRegistry.register_model(
    "Eagle3LlamaForCausalLM",
    "vllm_ascend.models.llama_eagle3:AscendEagle3LlamaForCausalLM",
)

# 1. ModelRegistry查找（vllm/models/registry.py）
class ModelRegistry:
    @staticmethod
    def get_model_class(model_name):
        # 返回类路径字符串，如 "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM"
        return _models[model_name]

# 2. 动态导入（vllm/v1/worker/gpu_worker.py）
import importlib
module_path, class_name = model_cls_path.split(":")
module = importlib.import_module(module_path)
model_cls = getattr(module, class_name)

# 3. 实例化模型（vllm_ascend/models/kimi_k3.py）
class AscendKimiLinearForCausalLM(torch.nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([...])
        self.norm = RMSNorm(config.hidden_size)

# 4. 加载权重（vllm_ascend/models/kimi_k3.py）
def load_weights(self, weights_path):
    state_dict = torch.load(weights_path)
    self.load_state_dict(state_dict)

class AscendKimiLinearForCausalLM(torch.nn.Module):
    """Kimi K3因果语言模型的NPU优化实现"""
    
    def __init__(self, config):
        super().__init__()
        # 初始化模型层
        self.embed_tokens = nn.Embedding(...)
        self.layers = nn.ModuleList([...])
        self.norm = RMSNorm(...)
        
    def forward(self, input_ids, positions, ...):
        # 前向传播
        hidden_states = self.embed_tokens(input_ids)
        for layer in self.layers:
            hidden_states = layer(hidden_states, ...)
        return hidden_states

from vllm import ModelRegistry


def register_model():
    # ==================== Kimi K3 系列 ====================
    # 主模型：用于文本生成
    ModelRegistry.register_model(
        "KimiLinearForCausalLM",
        "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
    )
    # 兼容性别名：用于旧版本checkpoint
    ModelRegistry.register_model(
        "KimiK3ForCausalLM",
        "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
    )
    # 条件生成模型：用于多模态任务
    ModelRegistry.register_model(
        "KimiK3ForConditionalGeneration",
        "vllm_ascend.models.kimi_k3:AscendKimiK3ForConditionalGeneration",
    )
    # MTP模型：用于投机解码
    ModelRegistry.register_model(
        "KimiK3MTPModel",
        "vllm_ascend.models.kimi_k3_mtp:AscendKimiK3MTP",
    )
    # DSpark草稿模型：用于分布式投机解码
    ModelRegistry.register_model(
        "K3DSparkModel",
        "vllm_ascend.models.kimi_k3_dspark:AscendK3DSparkForCausalLM",
    )

    # ==================== DeepSeek V4 系列 ====================
    # 主模型：支持MLA多头潜在注意力
    ModelRegistry.register_model(
        "DeepseekV4ForCausalLM",
        "vllm_ascend.models.deepseek_v4.model:AscendDeepseekV4ForCausalLM",
    )
    # MTP模型：用于投机解码
    ModelRegistry.register_model(
        "DeepSeekV4MTPModel",
        "vllm_ascend.models.deepseek_v4.mtp:DeepSeekV4MTP",
    )
    # DSpark草稿模型：用于分布式投机解码
    ModelRegistry.register_model(
        "DSparkDraftModel",
        "vllm_ascend.models.deepseek_v4.dspark:DSparkDeepseekV4ForCausalLM",
    )

    # ==================== MiniMax M3 系列 ====================
    # 稀疏模型：使用MoE混合专家架构
    ModelRegistry.register_model(
        "MiniMaxM3SparseForCausalLM",
        "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForCausalLM",
    )
    # 条件生成模型：用于多模态任务
    ModelRegistry.register_model(
        "MiniMaxM3SparseForConditionalGeneration",
        "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForConditionalGeneration",
    )

    # ==================== 投机解码模型 ====================
    # DeepSeek MTP：用于Eagle3投机解码
    ModelRegistry.register_model(
        "DeepSeekMTPModel",
        "vllm_ascend.models.deepseek_mtp:AscendDeepSeekMTP",
    )
    ModelRegistry.register_model(
        "DeepseekV32MTPModel",
        "vllm_ascend.models.deepseek_mtp:AscendDeepSeekMTP",
    )
    ModelRegistry.register_model(
        "GlmMoeDsaForCausalLM",
        "vllm_ascend.models.deepseek_mtp:AscendGlmMoeDsaForCausalLM",
    )
    # LLaMA Eagle3：用于LLaMA的投机解码
    ModelRegistry.register_model(
        "LlamaForCausalLMVwnEagle3",
        "vllm_ascend.models.llama_eagle3_vwn:Eagle3VwnLlamaForCausalLM",
    )
    ModelRegistry.register_model(
        "Eagle3LlamaForCausalLM",
        "vllm_ascend.models.llama_eagle3:AscendEagle3LlamaForCausalLM",
    )
    # Qwen3 DSpark：用于Qwen3的分布式投机解码
    ModelRegistry.register_model(
        "Qwen3DSparkModel",
        "vllm_ascend.models.qwen3_dspark:AscendQwen3DSparkForCausalLM",
    )
    # Qwen3 DFlash2：用于Qwen3的Flash投机解码
    ModelRegistry.register_model(
        "DFlash2DraftModel",
        "vllm_ascend.models.qwen3_dflash2:DFlash2Qwen3ForCausalLM",
    )