vLLM Ascend 源码讲解

结合流程图与源码，逐模块解析项目实现

一、项目文件结构

vLLM Ascend 是 vLLM 的硬件插件，核心代码位于 vllm_ascend/ 目录。

vllm_ascend/ ├── __init__.py # 插件入口，注册NPU平台 ├── platform.py # NPUPlatform类定义 ├── envs.py # 环境变量集中管理 ├── ascend_config.py # Ascend配置解析 ├── utils.py # 通用工具函数 ├── core/ # 核心模块（调度、KV缓存） ├── models/ # 模型定义 ├── attention/ # 注意力后端 ├── ops/ # 自定义算子 ├── worker/ # Worker执行器 ├── distributed/ # 分布式通信 ├── quantization/ # 量化支持 ├── patch/ # 上游补丁 ├── compilation/ # 编译与图优化 ├── sample/ # 采样器 ├── lora/ # LoRA适配器 ├── spec_decode/ # 投机解码 ├── device/ # 设备管理 ├── model_loader/ # 模型加载器 ├── profiler/ # 性能分析 └── _310p/ # Atlas 300I Duo 兼容层
目录功能关系
graph LR
    A[vllm_ascend] --> B[core]
    A --> C[models]
    A --> D[attention]
    A --> E[ops]
    A --> F[worker]
    A --> G[distributed]
    A --> H[quantization]
    A --> I[patch]
    A --> J[compilation]

二、插件入口 __init__.py

vLLM 通过 register() 函数发现并加载硬件插件。

插件注册流程
graph TD
    A[vLLM Engine启动] --> B[扫描插件]
    B --> C[调用 register]
    C --> D[返回 NPUPlatform 路径]
    D --> E[实例化 NPUPlatform]
    E --> F[NPU设备可用]
vllm_ascend/__init__.py
def register():
    """Register the NPU platform."""
    return "vllm_ascend.platform.NPUPlatform"

def register_model():
    from .models import register_model
    register_model()

def register_model_loader():
    from .model_loader.netloader import register_netloader
    register_netloader()

说明

register() 返回类路径字符串，vLLM 会动态导入该类并实例化，完成 NPU 平台注册。后续 register_model、register_model_loader 等分别注册模型和加载器。

三、NPUPlatform 平台类

platform.py 是插件核心，继承自 vLLM 的 Platform 基类。

NPUPlatform 关键方法
graph TD
    A[NPUPlatform] --> B[check_and_update_config]
    A --> C[get_attn_backend_cls]
    A --> D[import_kernels]
    A --> E[pre_register_and_update]
    A --> F[get_device_name]
    A --> G[get_current_memory_usage]
vllm_ascend/platform.py
class NPUPlatform(Platform):
    _enum = PlatformEnum.OOT
    device_name: str = "npu"
    device_type: str = "npu"
    ray_device_key: str = "NPU"
    device_control_env_var: str = "ASCEND_RT_VISIBLE_DEVICES"
    simple_compile_backend: str = "eager"

    supported_quantization: list[str] = [
        ASCEND_QUANTIZATION_METHOD,
        COMPRESSED_TENSORS_METHOD,
        FP8_METHOD,
    ]

3.1 配置检查流程

check_and_update_config 是最重要的方法，在引擎启动时被调用。

check_and_update_config 执行步骤
graph TD
    A[开始] --> B[配置日志]
    B --> C[验证并行配置]
    C --> D[自动检测量化]
    D --> E[修复不兼容参数]
    E --> F[初始化Ascend配置]
    F --> G[更新编译模式]
    G --> H[设置Worker类]
    H --> I[验证SP/DCP一致性]
    I --> J[设置NPU分配器]
    J --> K[完成]

3.2 注意力后端选择

根据模型特性自动选择最优注意力后端。

注意力后端选择逻辑
graph TD
    A[get_attn_backend_cls] --> B{use_mla?}
    B -->|Yes| C{use_sparse?}
    B -->|No| D[AscendAttentionBackend]
    C -->|Yes| E[AscendSFABackend]
    C -->|No| F[AscendMLABackend]
vllm_ascend/platform.py - 后端映射
backend_map = {
    (True, False, False):  "vllm_ascend.attention.mla_v1.AscendMLABackend",
    (False, False, False): "vllm_ascend.attention.attention_v1.AscendAttentionBackend",
    (True, True, False):   "vllm_ascend.attention.sfa_v1.AscendSFABackend",
    (True, False, True):   "vllm_ascend.attention.dsa_v1.AscendDSABackend",
}

四、环境变量管理 envs.py

所有环境变量集中在 envs.py 管理，禁止在代码中硬编码。

环境变量使用流程
graph TD
    A[代码中需要配置] --> B[import envs]
            B --> C[envs.VLLM_ASCEND_XXX]
            C --> D[返回配置值]
vllm_ascend/envs.py
import os

env_variables = {
    # 是否启用NZ格式（默认启用）
    "VLLM_ASCEND_ENABLE_NZ": lambda: int(os.getenv("VLLM_ASCEND_ENABLE_NZ", "1")),

    # 是否启用图融合优化
    "VLLM_ASCEND_ENABLE_FUSION": lambda: int(os.getenv("VLLM_ASCEND_ENABLE_FUSION", "1")),

    # 调试模式开关
    "VLLM_ASCEND_DEBUG": lambda: int(os.getenv("VLLM_ASCEND_DEBUG", "0")),
}

规范要求

新增环境变量必须：1) 使用 VLLM_ASCEND_ 前缀；2) 在 env_variables 中定义并注释；3) 提供默认值。

五、模型注册机制

Ascend 插件为特定模型提供优化实现，通过注册机制接入 vLLM。

模型注册流程
graph TD
    A[register_model] --> B[注册 KimiK3]
    A --> C[注册 DeepSeekV4]
    A --> D[注册 MiniMaxM3]
    A --> E[注册 Qwen3变体]
vllm_ascend/models/__init__.py
from vllm import ModelRegistry

def register_model():
    ModelRegistry.register_model(
        "KimiLinearForCausalLM",
        "vllm_ascend.models.kimi_k3:AscendKimiLinearForCausalLM",
    )
    ModelRegistry.register_model(
        "DeepseekV4ForCausalLM",
        "vllm_ascend.models.deepseek_v4.model:AscendDeepseekV4ForCausalLM",
    )
    ModelRegistry.register_model(
        "MiniMaxM3SparseForCausalLM",
        "vllm_ascend.models.minimax_m3:MiniMaxM3SparseForCausalLM",
    )
已注册模型列表
graph TD
    A[已注册模型] --> B[KimiLinearForCausalLM]
    A --> C[DeepseekV4ForCausalLM]
    A --> D[MiniMaxM3SparseForCausalLM]
    A --> E[Qwen3DSparkModel]
    A --> F[DeepSeekMTPModel]
    A --> G[Eagle3LlamaForCausalLM]

六、注意力后端 attention/

attention 目录包含多种 NPU 优化的注意力实现。

attention/ ├── attention_v1.py # 标准注意力 ├── mla_v1.py # MLA多头潜在注意力 ├── sfa_v1.py # SFA稀疏注意力 ├── dsa_v1.py # DSA动态稀疏注意力 ├── fa3_v1.py # Flash Attention 3 ├── attention_mask.py # 注意力掩码 ├── indexer.py # 索引器 └── utils.py # 工具函数
注意力后端继承关系
graph TD
    A[AttentionBackend] --> B[AscendAttentionBackend]
    A --> C[AscendMLABackend]
    A --> D[AscendSFABackend]
    A --> E[AscendDSABackend]
    A --> F[AscendFABackend]
vllm_ascend/attention/attention_v1.py
class AscendAttentionBackend(AttentionBackend):
    """标准注意力后端，适用于普通Transformer模型"""

    def __init__(self, vllm_config: VllmConfig):
        self.vllm_config = vllm_config
        self.block_size = vllm_config.cache_config.block_size

    def forward(self, ...):
        # NPU注意力前向计算
        ...
标准注意力计算流程
graph TD
    A[输入Q/K/V] --> B[计算注意力分数]
    B --> C[应用注意力掩码]
            C --> D[Softmax归一化]
            D --> E[加权求和]
            E --> F[输出结果]

七、自定义算子 ops/

ops 目录包含 NPU 优化的自定义算子实现。

ops/ ├── activation.py # 激活函数(SiLU等) ├── layernorm.py # LayerNorm/RMSNorm ├── linear.py # 线性层 ├── rotary_embedding.py # 旋转位置编码 ├── mla.py # MLA专用算子 ├── fused_moe/ # MoE融合算子 ├── triton/ # Triton内核 └── register_custom_ops.py # 算子注册
自定义算子分类
graph TD
    A[自定义算子] --> B[注意力相关]
    A --> C[归一化]
    A --> D[激活函数]
    A --> E[位置编码]
    A --> F[线性层]
vllm_ascend/ops/layernorm.py
class AscendRMSNorm(torch.nn.Module):
    """NPU优化的RMSNorm实现"""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 使用NPU专用内核计算RMSNorm
        return rms_norm_npu(x, self.weight, self.eps)

八、Worker执行器 worker/

Worker 负责模型的实际执行，是推理的核心组件。

worker/ ├── worker.py # Worker入口 ├── model_runner_v1.py # V1模型运行器 ├── v2/ # V2模型运行器 │ └── model_runner.py ├── npu_input_batch.py # NPU输入批处理 ├── block_table.py # 块表管理 └── utils.py # 工具函数
Worker初始化流程
graph TD
    A[Worker启动] --> B[加载配置]
    B --> C[初始化NPU设备]
    C --> D[加载模型权重]
    D --> E[初始化注意力后端]
    E --> F[分配KV缓存]
    F --> G[就绪]
推理执行流程
graph TD
    A[接收请求] --> B[预处理Token]
    B --> C[构建注意力元数据]
    C --> D[模型前向传播]
    D --> E[采样输出Token]
            E --> F{是否结束?}
            F -->|No| D
            F -->|Yes| G[返回结果]
vllm_ascend/worker/model_runner_v1.py
class NPUModelRunner(GPUModelRunner):
    """NPU模型运行器，继承自GPU实现并适配NPU"""

    def __init__(self, vllm_config: VllmConfig):
        super().__init__(vllm_config)
        self.device = torch.device("npu")

    def execute_model(self, ...):
        # 在NPU上执行模型前向
        ...

九、分布式通信 distributed/

distributed 目录实现多NPU间的通信和协调。

distributed/ ├── kv_transfer.py # KV缓存传输 ├── weight_transfer.py # 权重传输 ├── device_communicators/ # 设备通信器 │ └── npu_communicator.py └── parallel.py # 并行工具
分布式推理架构
graph TD
    A[分布式推理] --> B[张量并行 TP]
    A --> C[流水线并行 PP]
    A --> D[数据并行 DP]
    B --> E[NPU通信器]
    C --> E
    D --> E
KV缓存传输流程
graph TD
    A[P节点-预填充] --> B[生成KV缓存]
            B --> C[序列化KV数据]
            C --> D[传输到D节点]
            D --> E[D节点-解码]
            E --> F[返回结果]

十、量化支持 quantization/

量化模块提供多种模型量化方案，提升推理效率。

quantization/ ├── ascend_quant.py # Ascend专用量化 ├── fp8.py # FP8量化 ├── compressed_tensors.py # 压缩张量 └── utils.py # 量化工具
量化方法选择
graph TD
    A[量化检测] --> B{模型格式?}
    B -->|FP8| C[AscendFp8Config]
    B -->|压缩张量| D[AscendCompressedTensorsConfig]
    B -->|Ascend格式| E[AscendModelSlimConfig]

十一、补丁机制 patch/

patch 目录包含对上游 vLLM 的补丁，用于适配 NPU 特性。

patch/ ├── platform/ # 平台级补丁 ├── worker/ # 工作节点补丁 └── patch_utils.py # 补丁工具
补丁应用流程
graph TD
    A[引擎启动] --> B[adapt_patch]
            B --> C[加载平台补丁]
            C --> D[加载Worker补丁]
            D --> E[完成适配]
vllm_ascend/utils.py
def adapt_patch(is_global_patch: bool = False):
    """应用Ascend适配补丁"""
    from vllm_ascend.patch import apply_patches
    apply_patches(is_global_patch)

十二、编译优化 compilation/

compilation 目录实现 NPU 专用的编译和图优化。

compilation/ ├── acl_graph.py # ACL图模式 ├── compiler_interface.py # 编译器接口 ├── graph_fusion_pass_manager.py # 图融合 └── piecewise.py # 分段编译
编译流程
graph TD
    A[模型代码] --> B[Torch编译]
            B --> C[图优化Pass]
            C --> D[图融合]
            D --> E[生成NPU内核]
            E --> F[ACL图执行]
首页