【vLLM-Ascend】 ModelRunner架构解析

作者：昇腾实战派

ModelRunner总览结构图

1.核心功能概述

• ​模型加载与编译 ​​：支持模型的加载、LoRA适配、以及使用torch.compile进行优化。
• ​KV缓存管理 ​​：初始化和管理键值（KV）缓存，用于加速注意力机制中的计算。
• ​图捕获与优化 ​​：通过torchair或ACL Graph进行NPU图捕获，以提升推理效率。
• ​请求调度与状态更新 ​​：处理来自调度器的输出，并更新模型内部的状态。
• ​多模态支持 ​​：支持多模态输入（如图像、文本等），并进行相应的嵌入处理。

2.类结构分析

class NPUModelRunner

该类是整个模块的核心，继承自LoRAModelRunnerMixin，提供了以下关键方法：

• ​**__init__**​：初始化配置参数、设备设置、模型相关参数、缓存配置等。
• ​**_update_states**​：根据调度器输出更新模型内部状态。
• ​**execute_model**​：执行模型推理流程，包括前向传播、采样生成、日志概率计算等。
• ​**load_model**​：加载模型权重并应用LoRA适配。
• ​**initialize_kv_cache**​：初始化KV缓存。
• ​**capture_model**​：捕获NPU图，优化推理性能。
• ​**profile_run**​：对模型进行性能分析，特别是多模态输入的处理。

3. 主要方法解析

1._update_states(self, scheduler_output: "SchedulerOutput")

• 根据调度器输出更新模型内部状态。
• 移除已完成的请求，添加新请求。
• 更新每个请求的已生成 token 数量、块 ID 等信息。
• 刷新采样元数据（sampling metadata）。

2.execute_model(...)

这是整个推理流程的核心函数，负责：

1. 更新状态

   self._update_states(scheduler_output)

2. 处理请求

   (attn_metadata, hidden_states, spec_decode_metadata, positions,

   num_scheduled_tokens, sample_indices) = self._process_reqs(…)

3. 前向传播

   logits = self.model.compute_logits(hidden_states[sample_indices],None)

4. 应用结构化输出约束

   ifscheduler_output.grammar_bitmaskisnotNone:

   logits = self.apply_grammar_bitmask(scheduler_output, logits)

5. 采样下一个 token
   sampling_metadata = self.input_batch.sampling_metadata
   sampler_output = self.sampler(logits=bonus_logits, sampling_metadata=sampling_metadata)

6. 构建输出结果
   model_runner_output = ModelRunnerOutput(…)

---

3.apply_grammar_bitmask(...)

• ​作用 ​​：对 logits 应用结构化输出约束（bitmask），防止生成非法 token。
• ​原理 ​​：
  • 输入一个 bitmask，表示哪些 token 是允许的。
  • 将不符合语法规则的 token 对应的 logit 设为负无穷。
  • 使用 xgrammar 提供的 C++ 实现，保证性能。
• ​数据拷贝 ​​：
  • xgrammar 只支持 CPU 和 float32 类型的 logits。
  • 所以先将 logits 拷贝到 CPU 并转换为 float32，处理完后再拷贝回 GPU。

---

4._process_reqs(...)

• ​作用 ​​：准备输入数据（位置编码、注意力掩码等）。
• ​流程 ​​：
  1. 获取每个请求的 scheduled token 数量。
  2. 构造位置编码（positions）。
  3. 构造 slot_mapping（用于索引缓存块）。
  4. 构造 attention mask。
  5. 构造 attention metadata。
  6. 执行模型前向传播。

---

5.initialize_kv_cache(...)/get_kv_cache_spec(...)

• ​作用 ​​：初始化 KV 缓存。
• ​实现方式 ​​：
  • 每个注意力层分配内存空间。
  • 支持不同的 KV 缓存格式（如FullAttentionSpec）。
  • 使用bind_kv_cache将缓存绑定到模型中。

---

6.capture_model()

• ​作用 ​​：通过torchair或ACL Graph进行图捕获，提升推理效率。
• ​流程 ​​：
  • 对不同 batch size 运行 dummy run。
  • 编译图模式。
  • 记录消耗时间和显存。

---

7.profile_run()/_dummy_run(...)

• ​作用 ​​：对模型进行性能分析，特别是多模态输入的处理。
• ​流程 ​​：
  1. 构造 dummy 输入。
  2. 执行一次前向传播。
  3. 清理缓存。

---

4. 其他重要方法

方法名	说明
load_model()	加载模型权重并应用 LoRA 适配
get_model()	返回当前模型
_get_spec_token_ids(...)	生成推测解码所需的 token ID
_generate_draft_token_ids(...)	使用 ngram 等策略生成 draft tokens
_generate_mtp_token_ids(...)	使用 deepseek_mtp 等策略生成 draft tokens

5. NPU modelrunner 流程图

这是整个NPU modelrunner 的流程图，上图中步骤1-5为主模型推理，6为主模型验证，7.0-7.3为MTP模型推理。此次主要适配内容为主模型验证和mtp模型推理。
如果未配置MTP可以忽略整个第7阶段。
整个MTP的推理流程主要分为三个阶段：主模型推理、主模型验证、mtp模型推理。