• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN

导读

本文以 vllm-ascend 仓库中 Model Runner V2 的官方说明文档 为骨架，逐项拆解该实验性模型运行器与上游 vLLMGPUModelRunner之间的六项待处理差距（Gaps）与一项已落地能力，并深入对应的源码实现（model_runner.py、attn_utils.py、hooks.py 等），说明每项差距"为什么存在、改了什么、何时可以移除"。读完本文，你将理解 Ascend NPU 上 KV Cache 分配与 reshape 的 2 MiB 对齐要求、FULL 图模式下 FIA 算子的 workspace 生命周期管理、runtime_guard 观测装饰器的注入方式，以及 speculative decoding 各草稿模型在 MRV2 下的分派路径。

注意：MRV2 处于积极开发中，以下内容以当前仓库快照为准，后续提交可能随时变更。

一、MRV2 是什么：目录结构与整体定位

vllm_ascend/worker/v2目录存放的是 vllm-ascend 在 vLLM 新版（V1 架构）模型运行器之上的 Ascend 实现。上游 vLLM 在vllm/v1/worker/gpu/model_runner.py中提供了GPUModelRunner，而 Ascend 版本通过继承与覆写的方式注入 NPU 特有的行为。从目录结构看，该目录主要包含：

• model_runner.py：核心的NPUModelRunner类，继承上游GPUModelRunner；
• attn_utils.py：KV Cache 的规格推导、分配（allocate）与 reshape 工具；
• aclgraph_utils.py：ModelAclGraphManager等 ACL 图（对应上游 CUDA Graph）管理逻辑；
• input_batch.py：AscendInputBatch/AscendInputBuffers，为 NPU 注意力后端补充seq_lens_cpu等额外字段；
• spec_decode/：EAGLE、DFlash、DFlash2、DSpark、MTP 等草稿模型（speculator）的 Ascend 实现；
• model_states/、sample/：模型状态与采样相关的子模块。

NPUModelRunner在 model_runner.py 第 128 行 定义，其类注释明确说明这是 "Model runner for Ascend NPUs"，并且声明supports_standardized_shared_kv_backing = True，表示它已支持 vLLM 上游 #51718 引入的标准化共享 KV 后备存储布局。

原文档的核心是一份"与上游 vLLM 的差距清单"（Gaps to be Addressed），逐项说明了哪些地方仍依赖 Ascend 侧的自定义逻辑、为什么需要这些逻辑，以及未来消除差距的方向。下文按文档顺序逐条展开。

二、差距一：model_runner采样阶段的 runtime_guard 装饰器

2.1 差距内容

model_runnersample-token Ascend hooks + runtime_guard decorators（未完成）

文档指出：wave sync / sample-phase 的 guard 编排位于 vllm_ascend/observability/runtime_guard/hooks.py 的装饰器中，包括与 v1 共享的@runtime_guard_step和 MRV2 独有的@runtime_guard_sample_tokens。而 runner 自身的sample_tokens与postprocess_sampled方法体仍是原始功能逻辑，不含SamplePhaseResult/_rg_*协议、hook 拆分或 runner 侧 stash——删除装饰器即可字节级还原出未加 guard 的原始方法。

2.2 源码中的实际挂接

在 model_runner.py 中可以确认这一编排方式：

• sample_tokens使用@runtime_guard_sample_tokens装饰，见 model_runner.py 第 333 行；
• execute_model的装饰器顺序是@torch.inference_mode()在外、@runtime_guard_step在内（model_runner.py 第 540-541 行），这与 hooks.py 中"guard step 装饰器必须位于@torch.inference_mode()之内，使 wave sync 保持在 inference-mode 上下文中"的约定一致；
• runner 在初始化时通过RuntimeGuardProcessor.bind(self)绑定 guard，见 model_runner.py 第 273 行。

2.3 装饰器内部做了什么

阅读 hooks.py 可知：

• runtime_guard_step（L57）：在步骤体执行前调用guard.sync_for_step（配置总线、wave 装载、回收），并在finally中对无采样路径（提前返回/异常/dummy wave）调用end_of_wave_sync(allow_manual_dump=False)，确保 last-PP 的 lockstep 集合通信永远不会被跳过；
• runtime_guard_sample_tokens（L134）：在方法执行前清空上一步的 stash、peek 采样前的execute_model_state，然后用guard.run_sample_phase包住原始sample_tokens_fn，采样期间通过wrap_compute_logits_for_pre_sample与wrap_postprocess_sampled两个上下文管理器注入观测逻辑；
• AscendAsyncOutput（L324）：用于异步输出路径，在AsyncOutput.get_output()完成 D2H 裁剪之后才执行check_after_sample。

文档特别强调了一个设计意图：这些装饰器是"纯粹的观测层"（pure observability），装饰后的方法仍保留原始功能体，删除装饰器即可还原上游行为——这也是为什么它被列为"待消除的差距"：一旦上游GPUModelRunner.execute_model/sample_tokens的契约变化，需要同步跟踪这些装饰器的适配。

三、差距二：set_cos_and_sin与update_cos_sin旋转位置编码处理

3.1 差距内容

set_cos_and_sin&update_cos_sin

文档指出：DeepSeek 类模型（MLA）仍然需要在 model_runner 中设置和更新 cos/sin，未来当 MLA 能在内部自行解决 cos/sin 时应移除。

涉及位置：NPUModelRunner.__init__、NPUModelRunner.prepare_inputs、AscendInputBatch.make_dummy。

3.2 源码证据

在 model_runner.py 第 249 行 可以看到set_cos_and_sin(vllm_config, self.max_num_reqs, self.decode_query_len, self.dtype, self.device)的调用，其上方就带有一行 TODO 注释：TODO: remove set_cos_and_sin (together with update_cos_sin) when mla can properly handle cos/sin internally。

这两个函数从 vllm_ascend/ops/rotary_embedding.py 导入（见 model_runner.py 第 95 行），它们负责按max_num_reqs和decode_query_len预生成并缓存旋转位置编码的 cos/sin 表。由于 MRV2 在 decode 阶段按decode_query_len = num_speculative_steps + 1的长度推进（model_runner.py 第 246 行），cos/sin 表的容量必须与图捕获时的 query 长度对齐，这就是为什么需要显式的"设置"与"更新"两套操作贯穿__init__、prepare_inputs与 dummy batch 构造。

3.3 何时可以移除

该差距的消除前提是 MLA 注意力后端能够在层内部直接计算 cos/sin，而无需 runner 预置表。在此之前，任何启用 MLA 的 DeepSeek 类模型（如 DeepSeek-V4/V4.1 系列）在 MRV2 路径下都依赖这套 runner 级旋转位置编码管理，改动时需同步关注 model_runner.py 与 input_batch.py 中make_dummy的构造逻辑。

四、差距三：KV Cache 的分配与 reshape（_allocate_kv_cache/_reshape_kv_cache）

4.1 差距内容

_allocate_kv_cache&_reshape_kv_cache

文档给出的两个原因：

1. 连续空间要求：KV cache 需要连续空间，因此在 Ascend 侧被拆分为独立的 K cache 与 V cache 分别存放；
2. PD 分离（prefill/disaggregation）的 2 MiB 对齐要求：KV cache 张量必须 2 MiB 对齐，因此需要自定义的 KV cache 初始化。

涉及位置（均在 attn_utils.py）：_get_layer_kv_cache_specs、_get_attention_kv_cache_dims、_align_memory、_allocate_kv_cache、_reshape_kv_cache。

4.2 源码实现解读

在 attn_utils.py 中可以逐一定位这些函数：

• _get_layer_kv_cache_specs（L630）：把kv_cache_config.kv_cache_groups按层名展开成{layer_name: KVCacheSpec}字典，处理UniformTypeKVCacheSpecs的拆包；
• _get_attention_kv_cache_dims（L648）：对AscendMLAAttentionSpec返回(kv_lora_rank, qk_rope_head_dim)，普通注意力返回(head_size, head_size_v)；
• _align_memory（L791）：通过data_ptr计算偏移，返回对齐后的张量切片；
• _allocate_kv_cache（L878）：整体分配逻辑，注释明确写道 "FullAttention caches use a combined allocation for block-strided K/V views" 以及 "KV transfer aligns each raw allocation to 2 MiB"。

关键的 2 MiB 对齐常量出现在_allocate_kv_cache内部（attn_utils.py 第 910-911 行）：

# prefill disaggregation need the addr of cache tensor be aligned with 2M alignment = 2 * 1024 * 1024

而_allocate_int8_cache_tensor（L802）在启用了 KV transfer 时通过_align_memory(raw_tensor, alignment)[:numel]返回对齐后的原始 int8 缓存张量，从而保留 Mooncake/ADXL 的对齐语义。此外还包含对 DeepSeek-V4 模型的专用处理：所有KVCacheTensor描述符必须共享一个后备分配（共享 tuple backing），否则直接抛错。

4.3 reshape 侧的对应实现

与文档中的_reshape_kv_cache对应的是 attn_utils.py 第 1322 行的_reshape_kv_cache_v2。该函数负责把分配得到的原始张量 reshape 成注意力后端（如 FA3、SFA、MLA）所需的 block-strided K/V 视图；对于 MLA 主缓存（如AscendSFAIndexerCacheSpec）还需要附带 scale 张量。文档把分配与 reshape 一并列为待消除差距，是因为这两步都是为满足"连续 K/V 分离 + 2 MiB 对齐"这两条 NPU 硬件约束而存在的自定义逻辑，一旦上游或硬件不再需要这些约束即可移除。

五、差距四：torch_npu_graph_wrapper与 FULL 图模式下的 workspace 管理

5.1 差距内容

torch_npu_graph_wrapper

文档指出：FIA 算子在 FULL 模式下需要显式分配 workspace，且每个 workspace 对应一张图（一个特定 batch_size），必须在每次捕获恰好完成时通过weak_ref_workspaces释放，否则会 OOM。因此需要把weak_ref_workspaces作为torch.npu.graph的后处理逻辑进行 patch。

涉及位置：utils.torch_cuda_wrapper、utils.torch_npu_graph_wrapper。

5.2 源码实现解读

两个函数都位于 vllm_ascend/worker/v2/utils.py：

• torch_cuda_wrapper（L24）：一个上下文管理器，把torch.cuda.*的一系列符号替换为torch.npu.*对应物（Event、Stream、CUDAGraph = UpdatableGraph、graph = torch_npu_graph_wrapper、synchronize等），并同时替换breakable_cudagraph.weak_ref_tensor(s)，使上游 CUDA Graph 捕获代码能在 NPU 上运行；
• torch_npu_graph_wrapper（L64）：在torch.npu.graph上下文中进入，退出时调用weak_ref_workspaces(get_graph_params())与weak_ref_workspaces(get_draft_graph_params())，把每次捕获的图 workspace 转成弱引用，从而在捕获完成时及时回收，避免 OOM。若开启 super kernel 优化，还会在退出时执行graph.super_kernel_optimize。

weak_ref_workspaces、get_graph_params、get_draft_graph_params均来自 vllm_ascend/compilation/acl_graph.py（见 utils.py 第 9-13 行 的导入）。

5.3 适用场景

FULL 模式（cudagraph_mode为 FULL）下，每个 batch_size 对应一张捕获的 ACL 图；FIA（Fused Indexer Attention）类算子的 workspace 必须按图显式分配。torch_npu_graph_wrapper恰好复用了torch.npu.graph上下文管理器的退出时机来做 workspace 的弱引用后处理，避免再打一个上游 monkey patch。这正是文档所说的"把weak_ref_workspaces视为torch.npu.graph的后处理并 patch 它"。

六、差距五与六：model_runner.graph_manager_wrapper与speculator.graph_manager_wrapper

6.1 差距内容

model_runner.graph_manager_wrapper（未完成）speculator.graph_manager_wrapper（未完成）

文档指出：ModelAclGraphManager需要 model_runner 的input_buffers和model_state.attn_metadata来执行update_full_graph_params，因此 model_runner 必须被传入ModelAclGraphManager.__init__；同理，AutoRegressiveAclGraphManager需要 speculator 的input_buffers和model_state.attn_metadata，因此 speculator 也要传入其构造函数。

涉及位置：

• model_runner.NPUModelRunner.initialize_kv_cache；
• speculator.AscendEagleSpeculator.init_cudagraph_manager。

6.2 目标模型侧：graph_manager_wrapper

在 model_runner.py 第 1361 行 定义的graph_manager_wrapper(model_runner)是一个上下文管理器：进入时把vllm_model_runner.ModelCudaGraphManager替换为工厂函数，工厂函数构造ModelAclGraphManager(..., model_runner, ...)，把 runner 实例注入；退出时恢复原始类。它被用在 model_runner.py 第 419 行 的initialize_kv_cache中，包裹上游super().initialize_kv_cache(...)的调用。

ModelAclGraphManager定义于 aclgraph_utils.py 第 147 行，其run_fullgraph（L183）在重放 FULL 图时通过update_full_graph_params（来自 vllm_ascend/compilation/acl_graph.py）按每步的实际 token 数更新注意力后端参数，而这一更新需要读取self.model_runner.model_state.attn_metadata与self.model_runner.attn_groups——这正是文档所述"model_runner 必须传入__init__"的原因。

6.3 草稿模型侧：init_cudagraph_manager的 speculator 注入

草稿模型一侧由 spec_decode/autoregressive/speculator.py 第 288 行 的init_cudagraph_manager完成：在调用super().init_cudagraph_manager(cudagraph_mode)之后，把prefill_cudagraph_manager/decode_cudagraph_manager上的speculator属性指向 self，并共享update_stream，这样自动回归式（auto-regressive）草稿模型的 ACL 图管理器也能拿到 speculator 的输入缓冲与注意力元数据来完成update_full_graph_params。EAGLE 草稿模型通过 spec_decode/eagle/speculator.py 中的AscendEagleSpeculator（继承AscendAutoRegressiveSpeculator与上游EagleSpeculator）复用这条路径。

七、已落地能力：extract_hidden_states（MRV2）

7.1 差距清单中唯一已勾选项

原文档中extract_hidden_states（MRV2）条目标注为[x]，属于已解决项。要点如下：

• 上游 vLLM 为 MRV2 增加了extract_hidden_states支持，Ascend 通过init_speculator分派直接复用上游的ExtractHiddenStatesSpeculator；
• HiddenStateCacheSpec走独立的单张量 allocate/reshape 路径；
• use_aux_hidden_state_outputs由上游GPUModelRunner.__init__开启；
• 在 vLLM #51718（0828 pin）之后，每个 KV 描述符都是共享后备存储上的一个 view，采用[B, H, N, C]页布局；hidden-state dump 保留在每层私有缓冲中，避免与混合 Attention/Mamba 存储重叠，reshape 与CacheOnlyAttentionLayer.basic_cache对齐。

7.2 源码佐证

在 spec_decode/init.py 第 25 行的init_speculator中可以看到分派逻辑：当speculative_config.method == "extract_hidden_states"时，直接返回vllm.v1.worker.gpu.spec_decode.extract_hidden_states.ExtractHiddenStatesSpeculator，注释明确写着 "No Ascend-specific behavior beyond update_stream assignment in NPUModelRunner; reuse upstream ExtractHiddenStatesSpeculator as-is"。也就是说，这条路径上 Ascend 无需自定义 speculator 类，仅由 runner 统一管理update_stream。

在 attn_utils.py 中还有两处配套实现：

• 分配侧（L1047-1054）：当某组共享层中包含HiddenStateCacheSpec时，让它脱离 #51718 的混合后备存储，按各自的 page 独立分配，避免 float32 SSM 写入与 bfloat16 hidden-state dump 相互覆盖；
• reshape 侧（L1476-1483）：HiddenStateCacheSpec子类化自MLAAttentionSpec，因此必须在通用 MLA reshape 路径之前特判：要求单张量、按page_size_bytes整数分页。

八、MRV2 的其他 Ascend 关键适配（补充阅读）

除差距清单外，MRV2 还有若干值得了解的 Ascend 定制点，可作为理解差距的上下文：

1. AscendInputBatch/AscendInputBuffers（input_batch.py）：为 NPU 注意力后端补充seq_lens_cpu等字段，make_dummy负责构造图捕获与 dummy run 使用的假批次；
2. profile_run的 MC2 预留（model_runner.py 第 668 行）：MoE 模型在标准profile_run之前需要一次额外的 dummy run，按mc2_tokens_capacity预留 HCCL buffer；
3. fine-grained TP 的图强制（model_runner.py 第 730 行）：o_proj/MLP 跨 DP 的细粒度 TP 交换要求每一步都运行在捕获的图上，eager 分派会直接抛错，防止 HCCL 集合失步；
4. 草稿模型分派总表（spec_decode/init.py）：按use_dspark()/use_dflash()/method == "mtp"/use_eagle()依次分派到AscendDSparkSpeculator、AscendDFlash(2)Speculator、AscendMTPSpeculator、AscendEagleSpeculator，未支持的方法抛NotImplementedError。

九、总结：差距清单的工程价值

MRV2 的这份差距清单本质上是一份"技术债台账"：每一项都标注了 Why（为什么需要 Ascend 侧自定义）与 Location（涉及的具体函数），并用复选框标记完成状态。它对于开发者有两重价值：

• 追踪上游演进：当 vLLM 上游GPUModelRunner、ModelCudaGraphManager、speculator 契约变化时，清单中对应的条目就是需要重新审视的锚点；
• 理解 NPU 约束：K/V 分离与 2 MiB 对齐、FULL 图 workspace 回收、cos/sin 表管理、MLA hidden-state 独立缓冲等，都是硬件/算子在 MRV2 中的真实约束，是排查分配、OOM、图捕获相关问题的第一手索引。

截至当前仓库快照，extract_hidden_states已完全复用上游实现，其余六项仍为 Ascend 侧自定义逻辑；后续如需跟进最新状态，可随时查阅 vllm_ascend/worker/v2/README.md 的勾选情况，并对照本文给出的源码位置逐项核对。

• 人工智能
• 大模型
• 模型推理服务
• Ascend
• CANN