工作节点 - vLLM Ascend

一、模块概述

worker目录包含模型执行的核心组件，负责模型加载、推理执行和分布式协调。包含V1和V2两代ModelRunner实现，以及完整的子模块体系。

图1：Worker架构

graph LR
    A[NPUWorker] --> B[ModelRunner]
    A --> C[NPUInputBatch]
    A --> D[BlockTable]
    A --> E[DCPManager]
    
    B --> F[模型前向传播]
    B --> G[ACL图捕获]
    
    C --> H[批次构建]
    C --> I[序列管理]
    
    D --> J[KV缓存管理]
    
    F --> K[NPU执行]

二、文件结构

worker/
├── __init__.py
├── worker.py                 # NPUWorker主入口
├── model_runner_v1.py        # V1 NPUModelRunner
├── npu_input_batch.py        # NPUInputBatch输入批处理
├── block_table.py            # BlockTable, MultiGroupBlockTable
├── encoder_acl_graph.py      # EncoderAclGraphManager
├── dcp_utils.py              # DCPManager（分解计算工具）
├── utils.py                  # AscendKVBlockZeroer
├── v2/                       # V2架构（全新设计）
│   ├── model_runner.py      # V2 NPUModelRunner
│   ├── input_batch.py       # AscendInputBatch, AscendInputBuffers
│   ├── block_table.py       # AscendBlockTables
│   ├── aclgraph_utils.py    # ModelAclGraphManager
│   ├── states.py            # AscendRequestState
│   ├── attn_utils.py        # get_kv_cache_spec(), build_attn_metadata()
│   ├── eplb.py              # AscendEPLBController
│   ├── pcp_manager.py       # AscendPCPManager
│   ├── pp_utils.py          # SpecPPSupport
│   ├── spec_decode/          # V2投机解码
│   ├── sample/               # V2 NPU采样
│   └── model_states/         # 模型状态管理

三、V1 vs V2架构对比

图2：V1与V2架构对比

graph LR
    subgraph V1架构
    A1[NPUWorker] --> B1[ModelRunner V1]
    B1 --> C1[GPUModelRunner继承]
    B1 --> D1[NPUInputBatch]
    B1 --> E1[BlockTable]
    end
    
    subgraph V2架构
    A2[NPUWorker] --> B2[ModelRunner V2]
    B2 --> C2[全新实现]
    B2 --> D2[AscendInputBatch]
    B2 --> E2[AscendBlockTables]
    B2 --> F2[AscendEPLBController]
    B2 --> G2[AscendPCPManager]
    end

V1 vs V2区别

1. V1：继承自vLLM的GPUModelRunner，适配NPU实现

2. V2：全新架构设计，不继承GPU实现，直接针对NPU优化

3. V2优势：更好的扩展性，支持EPLB、PCP等高级特性

四、核心组件详解

文件	关键类/函数	功能说明
worker.py	NPUWorker, parse_text_output()	NPU Worker主入口，管理Worker生命周期，协调ModelRunner和输入管理
model_runner_v1.py	NPUModelRunner, GraphCaptureContext, graph_capture(), ExecuteModelState	V1模型运行器，继承GPUModelRunner，支持ACL图捕获优化
npu_input_batch.py	NPUInputBatch	NPU输入批处理管理，优化批次构建和序列管理
block_table.py	BlockTable, MultiGroupBlockTable	KV缓存块表管理，支持多组块表（不同量化级别）
encoder_acl_graph.py	EncoderAclGraphManager, EncoderGraphParams, EncoderForwardContext	编码器ACL图优化，提升编码器推理性能
dcp_utils.py	DCPManager, DCPSpecDecodeMTPInputs	分解计算管理器，支持PD分离架构
v2/model_runner.py	NPUModelRunner (V2)	V2模型运行器，全新架构，更好的性能和扩展性
v2/input_batch.py	AscendInputBatch, AscendInputBuffers	V2输入批处理，更高效的内存管理
v2/block_table.py	AscendBlockTables	V2块表管理，支持更多块表类型
v2/eplb.py	AscendEPLBController	V2专家并行负载均衡控制器
v2/pcp_manager.py	AscendPCPManager, AscendPCPAttentionContext	V2流水线上下文并行管理器

五、Worker初始化流程

图3：Worker初始化流程

sequenceDiagram
    participant U as 上游vLLM
    participant W as NPUWorker
    participant MR as ModelRunner
    participant A as Attention Backend
    participant O as Ops
    
    U->>W: 创建Worker
    W->>MR: 初始化ModelRunner
    MR->>A: 选择注意力后端
    MR->>O: 初始化自定义算子
    W->>W: 加载模型权重
    W->>W: 分配KV缓存
    W-->>U: Worker就绪

六、推理执行流程

图4：单次推理执行流程

sequenceDiagram
    participant U as 上游vLLM
    participant W as NPUWorker
    participant MR as ModelRunner
    participant A as Attention
    participant O as Ops
    participant N as NPU
    
    U->>W: execute_model()
    W->>MR: execute()
    MR->>O: RMSNorm归一化
    O->>N: NPU执行
    N-->>O: 归一化结果
    MR->>A: 计算注意力
    A->>N: NPU执行
    N-->>A: 注意力输出
    MR->>O: FFN前馈网络
    O->>N: NPU执行
    N-->>O: FFN输出
    MR-->>W: logits
    W-->>U: 返回结果

vllm_ascend/worker/model_runner_v1.py（简化）

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

    def execute_model(self, ...):
        # 1. 预处理输入
        input_ids = self._prepare_input_ids(...)
        
        # 2. 构建注意力元数据
        attn_metadata = self._build_attn_metadata(...)
        
        # 3. 模型前向传播
        hidden_states = self.model(input_ids)
        
        # 4. 调用attention和ops模块
        for layer in self.model.layers:
            hidden_states = layer.attention(hidden_states, attn_metadata)
            hidden_states = layer.feed_forward(hidden_states)
        
        # 5. 采样输出
        logits = self.model.head(hidden_states)
        return logits

七、V2子模块详解

spec_decode/

V2投机解码支持

init_speculator() - 初始化提议器
Eagle提议器
DFlash提议器
DSpark提议器

sample/

V2 NPU优化采样

apply_top_k_top_p_npu()
apply_bad_words()
compute_token_logprobs()

model_states/

模型状态管理

AscendModelState
AscendMambaHybridModelState

pcp_manager.py

流水线上下文并行

AscendPCPManager
AscendPCPAttentionContext