首页 / 资讯中心 / 文章详情

vllm-ascend 实验性 Model Runner V2 适配指南:Ascend 特有逻辑与上游 vLLM 差距全解析

vllm-ascend 实验性 Model Runner V2 适配指南:Ascend 特有逻辑与上游 vLLM 差距全解析 ★ FEATURED ARTICLE
人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载导读本文以 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类继承上游GPUModelRunnerattn_utils.pyKV Cache 的规格推导、分配allocate与 reshape 工具aclgraph_utils.pyModelAclGraphManager等 ACL 图对应上游 CUDA Graph管理逻辑input_batch.pyAscendInputBatch/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_stepL57在步骤体执行前调用guard.sync_for_step配置总线、wave 装载、回收并在finally中对无采样路径提前返回/异常/dummy wave调用end_of_wave_sync(allow_manual_dumpFalse)确保 last-PP 的 lockstep 集合通信永远不会被跳过runtime_guard_sample_tokensL134在方法执行前清空上一步的 stash、peek 采样前的execute_model_state然后用guard.run_sample_phase包住原始sample_tokens_fn采样期间通过wrap_compute_logits_for_pre_sample与wrap_postprocess_sampled两个上下文管理器注入观测逻辑AscendAsyncOutputL324用于异步输出路径在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_sinupdate_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_cache4.1 差距内容_allocate_kv_cache_reshape_kv_cache文档给出的两个原因连续空间要求KV cache 需要连续空间因此在 Ascend 侧被拆分为独立的 K cache 与 V cache 分别存放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_specsL630把kv_cache_config.kv_cache_groups按层名展开成{layer_name: KVCacheSpec}字典处理UniformTypeKVCacheSpecs的拆包_get_attention_kv_cache_dimsL648对AscendMLAAttentionSpec返回(kv_lora_rank, qk_rope_head_dim)普通注意力返回(head_size, head_size_v)_align_memoryL791通过data_ptr计算偏移返回对齐后的张量切片_allocate_kv_cacheL878整体分配逻辑注释明确写道 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_tensorL802在启用了 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.pytorch_cuda_wrapperL24一个上下文管理器把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_wrapperL64在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 图FIAFused 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_wrapper6.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_cachespeculator.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_fullgraphL183在重放 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_statesMRV27.1 差距清单中唯一已勾选项原文档中extract_hidden_statesMRV2条目标注为[x]属于已解决项。要点如下上游 vLLM 为 MRV2 增加了extract_hidden_states支持Ascend 通过init_speculator分派直接复用上游的ExtractHiddenStatesSpeculatorHiddenStateCacheSpec走独立的单张量 allocate/reshape 路径use_aux_hidden_state_outputs由上游GPUModelRunner.__init__开启在 vLLM #517180828 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-1483HiddenStateCacheSpec子类化自MLAAttentionSpec因此必须在通用 MLA reshape 路径之前特判要求单张量、按page_size_bytes整数分页。八、MRV2 的其他 Ascend 关键适配补充阅读除差距清单外MRV2 还有若干值得了解的 Ascend 定制点可作为理解差距的上下文AscendInputBatch/AscendInputBuffersinput_batch.py为 NPU 注意力后端补充seq_lens_cpu等字段make_dummy负责构造图捕获与 dummy run 使用的假批次profile_run的 MC2 预留model_runner.py 第 668 行MoE 模型在标准profile_run之前需要一次额外的 dummy run按mc2_tokens_capacity预留 HCCL bufferfine-grained TP 的图强制model_runner.py 第 730 行o_proj/MLP 跨 DP 的细粒度 TP 交换要求每一步都运行在捕获的图上eager 分派会直接抛错防止 HCCL 集合失步草稿模型分派总表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 的勾选情况并对照本文给出的源码位置逐项核对。赞分享人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载相关推荐vLLM Ascend Model Runner V2 的 EPLB 异步架构深度解析vLLM Ascend Model Runner V2 的 EPLB 异步架构深度解析 导读 本文基于 vLLM Ascend 仓库的设计文档 model_ru人工智能大模型模型推理服务AscendCANNvLLM Ascend 安装指南在昇腾 NPU 上安装 vLLM 与 vLLM Ascend 的完整实操手册vLLM Ascend 安装指南在昇腾 NPU 上安装 vLLM 与 vLLM Ascend 的完整实操手册 本篇围绕 安装文档 https://link.g人工智能大模型模型推理服务AscendCANNvLLM Ascend 插件vllm-ascend完全指南在昇腾 NPU 上无缝运行 vLLM 的社区硬件插件vLLM Ascend 插件vllm ascend完全指南在昇腾 NPU 上无缝运行 vLLM 的社区硬件插件 本文档系统介绍 vllm ascend 这人工智能大模型模型推理服务AscendCANN上一篇RISC-V 操作码项目教程下一篇推荐使用 Laravel Report Generators 快速生成报告创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
阅读完成 · 觉得有帮助?
咨询建站