执行摘要

• 一句话：为 Kimi-VL 添加 ViT 全 CUDA 图支持
• 推荐动作：值得精读。重点阅读 moonvit.py 中 prepare_encoder_metadata 的设计（如何将 grid 相关计算外提）以及 kimi_vl.py 中协议方法的实现，可作为其他视觉模型启用 ViT CUDA 图的参考蓝本。

功能与动机

Kimi-VL 的 MoonVitPretrainedModel 在前向中包含 .tolist() 调用（位置嵌入插值、RoPE 频率计算、patch merging），与 CUDA 图捕获不兼容。本 PR 通过重构 moonvit.py 添加图安全路径（预计算元数据缓冲区），然后将 kimi_vl.py 接入 SupportsEncoderCudaGraph 协议，旨在消除 ViT 编码器的内核启动开销，提升多模态推理性能。

实现拆解

1. 重构 moonvit.py 编码器层：将 Learnable2DInterpPosEmb.forward 拆分为 get_pos_embeds，接受 grid_hws_list（Python 列表）而非原始 tensor，避免 CUDA 图内部的 .tolist() 循环；在 MoonVisionPatchEmbed 中添加 pos_embeds 参数支持跳过插值；为 Rope2DPosEmb 新增 get_freqs_cis_by_seqlens_list 方法，批量计算 RoPE 频率；新增 patch_merger_packed 和 _build_merge_gather_idx 实现图安全的 patch merging。
2. 预计算元数据入口：在 MoonVitPretrainedModel 中新增 prepare_encoder_metadata，接受 grid 列表，一次性返回所有与 grid 相关的元数据字典（pos_embeds、rope_freqs_cis、cu_seqlens、merge_gather_idx 等），此方法在图外调用。
3. 实现 SupportsEncoderCudaGraph 协议：KimiVLForConditionalGeneration 新增继承 SupportsEncoderCudaGraph，并实现以下方法：
   • get_encoder_cudagraph_config：定义 buffer_keys（pixel_values、pos_embeds、rope_freqs_cis 等）和配置。
   • get_encoder_cudagraph_budget_range：返回 token 预算范围（64 ~ max_num_batched_tokens）。
   • get_encoder_cudagraph_item_specs：根据 grid 返回每个 item 的 input_size 和 output_tokens。
   • select_encoder_cudagraph_items：根据 indices 从 mm_kwargs 中选择对应 buffer。
   • prepare_encoder_cudagraph_capture_inputs 等构建捕获/回放输入。
4. 测试覆盖：tests/models/multimodal/generation/test_vit_cudagraph.py 新增 kimi_vl 配置，使用 dummy 权重（dummy_hf_overrides）和单 bucket budget（encoder_cudagraph_token_budgets=[1024]），确保 CI 兼容。
5. 示例与文档：examples/generate/multimodal/vision_language_offline.py 在 MODELS_SUPPORT_VIT_CUDA_GRAPH 中添加 kimi_vl；docs/design/cuda_graphs_multimodal.md 在模型表格中新增 KimiVLForConditionalGeneration 行。

关键文件：

• vllm/model_executor/models/moonvit.py（模块 视觉编码器；类别 source；类型 core-logic；符号 forward, get_pos_embeds, get_freqs_cis_by_seqlens, get_freqs_cis_by_seqlens_list）: 核心重构：提取 get_pos_embeds、patch_merger_packed、prepare_encoder_metadata 等图安全方法，消除 .tolist() 依赖。
• vllm/model_executor/models/kimi_vl.py（模块 多模态模型；类别 source；类型 core-logic；符号 KimiVLForConditionalGeneration, get_encoder_cudagraph_config, get_encoder_cudagraph_budget_range, _get_grid_hws）: 协议实现：KimiVLForConditionalGeneration 继承 SupportsEncoderCudaGraph 并实现所有必要方法。
• tests/models/multimodal/generation/test_vit_cudagraph.py（模块 CUDA图测试；类别 test；类型 test-coverage；符号 kimi_vl_chat_template）: 测试覆盖：新增 Kimi-VL 配置，使用 dummy 权重和单 budget 验证 CUDA 图捕获与回放。
• examples/generate/multimodal/vision_language_offline.py（模块 离线示例；类别 source；类型 core-logic）: 示例支持：将 kimi_vl 加入 MODELS_SUPPORT_VIT_CUDA_GRAPH 列表。
• docs/design/cuda_graphs_multimodal.md（模块 设计文档；类别 docs；类型 documentation）: 文档更新：在 CUDA 图支持表格中添加 Kimi-VL 行。

关键符号：Learnable2DInterpPosEmb.get_pos_embeds, MoonVisionPatchEmbed.forward, Rope2DPosEmb.get_freqs_cis_by_seqlens_list, MoonVitPretrainedModel.prepare_encoder_metadata, MoonVitPretrainedModel.patch_merger_packed, MoonVitPretrainedModel._build_merge_gather_idx, KimiVLForConditionalGeneration.get_encoder_cudagraph_config, KimiVLForConditionalGeneration.get_encoder_cudagraph_budget_range, KimiVLForConditionalGeneration.get_encoder_cudagraph_item_specs, KimiVLForConditionalGeneration.select_encoder_cudagraph_items, KimiVLForConditionalGeneration.prepare_encoder_cudagraph_capture_inputs, KimiVLForConditionalGeneration.prepare_encoder_cudagraph_replay_buffers

关键源码片段

vllm/model_executor/models/moonvit.py

核心重构：提取 get_pos_embeds、patch_merger_packed、prepare_encoder_metadata 等图安全方法，消除 .tolist() 依赖。

def get_pos_embeds(
    self,
    grid_hws_list: list[list[int]] | list[tuple[int, int]],
) -> torch.Tensor:
    """ 根据 grid 列表生成打包的位置嵌入。
    输出形状 ``(sum(h * w), dim)``，通过对每个 grid 的
    学习权重 ``(height, width, dim)`` 进行插值并拼接。
    此方法在 CUDA 图外部调用，因此 per-grid Python 循环是安全的。
    """
    weight_shape = list(self.weight.shape[:-1]) # [H, W]
    pos_embs: list[torch.Tensor] = []
​
    for shape in grid_hws_list:
        shape_list = [int(shape[0]), int(shape[1])]
        if shape_list == weight_shape:
            # 无需插值，直接使用预训练权重
            pos_embs.append(self.weight.flatten(end_dim=1))
        else:
            # 对权重进行插值以适应不同大小的 grid
            pos_embs.append(
                F.interpolate(
                    self.weight.permute((2, 0, 1)).unsqueeze(0),
                    size=tuple(shape_list),
                    mode=self.interpolation_mode,
                )
                .squeeze(0)
                .permute((1, 2, 0))
                .flatten(end_dim=1)
            )
​
    if not pos_embs:
        return self.weight.new_zeros((0, self.weight.shape[-1]))
    return torch.cat(pos_embs)
​
​
def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
    # forward 现在调用 get_pos_embeds，传入 grid_hws 的 .tolist()
    pos_embs = self.get_pos_embeds(grid_hws.tolist())
    out = x + pos_embs
    return out

vllm/model_executor/models/kimi_vl.py

协议实现：KimiVLForConditionalGeneration 继承 SupportsEncoderCudaGraph 并实现所有必要方法。

def get_encoder_cudagraph_config(self):
    from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphConfig
    return EncoderCudaGraphConfig(
        modalities=["image"], # 当前仅支持图像模态
        buffer_keys=[ # CUDA 图中使用的 buffer 键
            "pixel_values",
            "pos_embeds",
            "rope_freqs_cis",
            "cu_seqlens",
            "max_seqlen",
            "merge_gather_idx",
        ],
        out_hidden_size=self.hidden_size, # 编码器输出隐藏层大小
    )
​
​
def get_encoder_cudagraph_item_specs(self, mm_kwargs):
    """ 根据每个 grid 的 (h, w) 计算 input_size 和 output_tokens。 """
    from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
    kh, kw = self.config.vision_config.merge_kernel_size
    return [
        EncoderItemSpec(
            input_size=h * w, # patch 总数 = h * w
            output_tokens=(h // kh) * (w // kw), # merge 后的 token 数
        )
        for h, w in self._get_grid_hws(mm_kwargs)
    ]

评论区精华

• 文档符号一致性：shen-shanshan 指出表格中 ✗ 应统一为 ❌︎，作者已修复。
• 测试模型规模：shen-shanshan 担忧 MoE 模型超过 CI GPU 显存，作者参考 #43082 改用 dummy_hf_overrides 和单 budget，保持测试轻量。
• 协议接口更新：shen-shanshan 指出 get_input_modality 已在 #44484 中废除，作者删除该方法。
• 变量命名：shen-shanshan 认为 normalized 参数名易混淆，作者将其重命名为 grid_pairs。

• prepare_encoder_metadata 用途：shen-shanshan 询问是否仅用于 CG 路径，作者确认并更新文档说明。

• 文档符号一致性 (style): 作者已采纳并修复。

• 测试模型规模 (testing): 作者改用 dummy_hf_overrides 和单 budget [1024]，参考 #43082 模式，避免下载完整 checkpoint。
• 协议方法废弃 (design): 作者已删除该方法。
• 变量命名清晰度 (style): 作者将 normalized 重命名为 grid_pairs。
• prepare_encoder_metadata 用途 (question): 作者确认仅用于 CUDA 图路径，并修正文档字符串，表示 eager 路径未改动（类似 InternVL #41759）。

风险与影响

• 风险：
  • 重构 moonvit.py 核心前向路径（如 forward 跳转至 get_pos_embeds），若图安全路径与 eager 路径产生数值差异可能导致回归；但测试验证了 eager 与 CUDA 图输出匹配，且遵循已有模型（Qwen3-VL）模式。
  • 当前仅支持 image 模态，video 模态未处理（modalities=["image"]），若未来 Kimi-VL 支持视频需扩展。
  • 测试使用 dummy 权重，不覆盖真实权重下 RoPE 频率计算等数值路径，存在遗漏精度问题的风险。
  • max_seqlen_override 参数可能被误用，但已在 CG 路径中强制设置。
• 影响：
  • 用户：Kimi-VL 用户可通过 --enable-vit-cuda-graph 获得显著 TTFT 降低（P99 降幅 33.5%），尤其在高并发场景下收益明显。
  • 系统：首次 CUDA 图捕获增加预热开销，但后续推理内核启动开销被消除；需额外显存存储 capture graph buffer（由 budget 控制）。
  • 团队：为后续模型（如 GLM-V、DeepSeek VL）接入 ViT CUDA 图提供可复现的模式（moonvit.py 的 prepare_encoder_metadata 模式）。
  • 风险标记：核心路径重构需验证数值对齐, 仅支持 image 模态, 测试使用 dummy 权重未覆盖真实路径, CUDA 图预热增加首次推理延迟, 依赖 SupportsEncoderCudaGraph 协议接口稳定性

关联脉络

• PR #38175 [RFC]: Support ViT Full CUDA Graph (Tracker): 本 PR 是该 tracker issue 的具体实现项之一，目标为实现 ViT 全 CUDA 图。
• PR #38061 [MM][Perf][CG] Support ViT full CUDA graph for Qwen3-VL: 本 PR 遵循 Qwen3-VL 的 ViT CUDA 图模式，许多设计（如 prepare_encoder_metadata、buffer_keys）与之相似。
• PR #44484 [MM] Remove get_input_modality from SupportsEncoderCudaGraph (cleanup): 该 PR 移除了 get_input_modality 方法，本 PR 作者据此删除了 Kimi-VL 的对应实现。
• PR #42288 [MM][CG] Migrate encoder CUDA graph to single values dict API: 本 PR 的 select_encoder_cudagraph_items 和 buffer 管理参考了该 API 变更。
• PR #43082 [MM][CG] Step3-VL ViT CUDA graph test with dummy weights: 本 PR 的测试模式（dummy_hf_overrides + 单 budget）借鉴了 Step3-VL 的 CI 友好测试配置。