Prhub

#43586 [MM][Perf][CG] Support dual-path ViT full CUDA graph for DeepSeek-OCR

原始 PR 作者 shen-shanshan 合并时间 2026-06-16 19:35 文件变更 16 提交数 25 评论 21 代码增减 +809 / -69

执行摘要

DeepSeek-OCR 视觉编码器双路径 CUDA 图形支持

DeepSeek-OCR 使用双塔 ViT(SAM+CLIP)和动态平铺机制,单个 CUDA 图无法适应可变 crop_shape 和 patch 数量。传统单图会导致大量零填充或图失效。双路径设计允许独立捕获全局和局部图,最大化图命中率,提升推理效率。PR body 提到依赖 #41234 和 #42288,并引用 DeepSeek-OCR 技术报告。

该 PR 值得精读,尤其是双路径图的设计模式:如何通过独立预算和路径回退避免动态形状问题。EncoderCudaGraphManager 的泛化支持(dict[str, dict])为未来多模型多路径图提供了扩展基础。建议关注 CI 测试的重新开启和显存优化。

讨论亮点

gemini-code-assist[bot] 指出了四个关键问题:(1)_get_num_input_output_tokens 中缺少对 image_spatial_crop 为 None 的空检查,会导致 TypeError;(2)EncoderCudaGraphConfig 应使用 input_key_by_modality 而非 input_keys;(3)images_crop 张量必须注册到 buffer_keys 并包含在 capture/replay 缓冲区中;(4)需要实现 get_max_frames_per_video 方法以遵循协议。这些均被作者采纳。

Isotr0py 建议将 image_side 等计算改为 property 避免重复,以及将 budget_graphs 改为 dict[str, ...] 以支持多路径,作者采纳。

关于双路径图设计,Isotr0py 最初认为局部补丁是计算瓶颈,建议优先捕获局部图;shen-shanshan 回应并实现了双路径独立选择方案,使全局和局部可同时拥有图,并且支持路径级回退。

实现拆解

  1. 模型侧实现(deepseek_ocr.py):在 DeepseekOCRForCausalLM 中添加 _get_num_input_output_tokens、get_encoder_cudagraph_config、prepare_encoder_cudagraph_capture_inputs、encoder_cudagraph_forward、prepare_encoder_cudagraph_replay_buffers、encoder_eager_forward 及 postprocess_encoder_output,实现 SupportsEncoderCudaGraph 协议。新增属性 image_side、global_image_output_token、patch_side、single_patch_output_token 用于计算固定 token 预算。get_encoder_cudagraph_config 返回配置时启用 enable_dual_path_graph=True,并注册 images_crop 为缓冲区键。
  2. 管理器扩展(encoder_cudagraph.py):将 budget_graphs 类型从 dict[int, BudgetGraphMetadata] 改为 dict[str, dict[int, BudgetGraphMetadata]],新增 global_token_budgets 和 local_token_budgets 独立预算列表。capture() 方法按路径('global'/'local')分别调用 _capture_budget_graph,新增 _get_graph_set 辅助方法安全获取图集合。新增 _execute_local_single_path 和 _execute_local_dual_path 进行推理时的路径选择。
  3. 数据契约扩展(encoder_cudagraph_defs.py):在 EncoderCudaGraphConfig 中添加 enable_dual_path_graph、global_token_per_image、local_token_per_patch 字段;在 EncoderItemSpec 中添加 global_output_tokens 和 local_output_tokens。
  4. 接口同步(interfaces.py + 7 个模型文件):在 SupportsEncoderCudaGraph 协议方法中添加 path: str = 'default' 和 local_output: Optional[torch.Tensor] = None 参数,所有现有模型实现按默认值向后兼容。
  5. 测试和文档:在 test_vit_cudagraph.py 添加 deepseek_ocr 配置项(启用 dummy weight,跳过执行);更新 docs/design/cuda_graphs_multimodal.md 在支持列表中增加 DeepSeek-OCR。
文件 模块 状态 重要度
vllm/model_executor/models/deepseek_ocr.py 模型层 modified 9.05
vllm/v1/worker/encoder_cudagraph.py CUDA 图管理器 modified 8.84
vllm/v1/worker/encoder_cudagraph_defs.py 图配置定义 modified 5.97
vllm/model_executor/models/interfaces.py 接口层 modified 5.34
tests/models/multimodal/generation/test_vit_cudagraph.py CUDA 图测试 modified 5.18

关键符号

get_encoder_cudagraph_config _get_num_input_output_tokens _capture_budget_graph _get_graph_set postprocess_encoder_output

关键源码片段

vllm/model_executor/models/deepseek_ocr.py data-contract

模型主实现文件,新增 SupportsEncoderCudaGraph 协议全部方法,实现双路径图形配置和 token 计算。

@property
def image_side(self) -> int:
    # 每个全局图像每维的输出网格单元数
    return math.ceil((BASE_SIZE // 16) / 4) # 16@property
def global_image_output_token(self) -> int:
    # 每个全局图像的 token 数:网格 + 每行一个换行
    return self.image_side * (self.image_side + 1) # 272@property
def patch_side(self) -> int:
    # 每个局部补丁每维的网格单元数
    return math.ceil((IMAGE_SIZE // 16) // 4) # 10@property
def single_patch_output_token(self) -> int:
    # 每个局部补丁的 token 数
    return self.patch_side * (self.patch_side + 1) # 110def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
    # 返回编码器 CUDA 图配置,启用双路径模式
    return EncoderCudaGraphConfig(
        modalities=['image'],
        input_key_by_modality={'image': 'pixel_values'},
        buffer_keys=['images_crop'],
        out_hidden_size=self.config.vision_config.hidden_size,
        enable_dual_path_graph=True,
        global_token_per_image=self.global_image_output_token,
        local_token_per_patch=self.single_patch_output_token,
    )
vllm/v1/worker/encoder_cudagraph.py core-logic

CUDA 图管理器核心逻辑,扩展支持双路径模式:budget_graphs 改为 dict[str, dict[int, ...]],新增单独 global/local 预算列表,capture/execute 分流。

def capture(self, graph_pool: Any):
    # 按路径分别捕捉 CUDA 图
    self.graph_pool = graph_pool
​
    if self.config.enable_dual_path_graph:
        # 双路径模式:先捕捉所有全局图,再捕捉局部图
        for token_budget in sorted(self.global_token_budgets, reverse=True):
            self._capture_budget_graph(token_budget, path='global')
        for token_budget in sorted(self.local_token_budgets, reverse=True):
            if token_budget == 0: # 跳过零预算(表示无局部补丁)
                continue
            self._capture_budget_graph(token_budget, path='local')
        logger.info(
            'Encoder CUDA graph capture complete. '
            'Captured %d global + %d local budget graphs.',
            len(self.budget_graphs['global']),
            len(self.budget_graphs['local']),
        )
        return
​
    # 默认单路径模式(保持向后兼容)
    for token_budget in sorted(self.token_budgets, reverse=True):
        self._capture_budget_graph(token_budget, path='default')
    logger.info(
        'Encoder CUDA graph capture complete. Captured %d budget graphs.',
        len(self.budget_graphs['default']),
    )

评论区精华

images_crop 必须注册到 CUDA 图缓冲区 正确性

gemini-code-assist[bot] 指出 images_crop 张量在 capture/replay 中未显式处理,导致回放时仍用初始化的虚拟数据。

结论:作者在 get_encoder_cudagraph_config 的 buffer_keys 中增加了 'images_crop',并确保 capture_inputs 和 replay_buffers 中包含此字段。 · 已解决

双路径预算图存储结构 设计

Isotr0py 建议将 budget_graphs 由 dict[int, BudgetGraphMetadata] 改为 dict[str, dict[int, BudgetGraphMetadata]],以支持多路径。

结论:作者采纳建议,重构 budget_graphs 结构,并新增 _get_graph_set 辅助方法。 · 已解决

局部补丁计算效率讨论及双路径方案选择 设计

Isotr0py 认为局部补丁是计算瓶颈,建议优先为局部补丁捕获图。shen-shanshan 回应并实现了双路径独立选择,使得全局和局部可同时拥有图,并支持路径级回退。

结论:双路径图设计允许每条路径独立选择最小合适预算,实现部分图回退,比仅捕获局部图更灵活。 · 已解决

风险与影响

  1. 显存增加:双路径模式需捕获两套图(全局 + 局部),CUDA 图捕获数量和总显存占用翻倍,可能对大 batch 或低显存环境造成压力。
  2. 测试覆盖不足:CI 测试因 OOM 被标记为 skip(见 commit skip deepseek-ocr ci),缺乏回归检测,可能遗漏边缘情况。
  3. 接口适配面广:interfaces.py 及 7 个模型文件同步修改协议方法签名,虽然向后兼容,但若其他分支未同步合并可能导致冲突。
  4. 动态形状交互:双路径图在 batch 内不同图像的分块状态可能不同(有/无局部补丁),图选择逻辑需要处理这种混合情况,但当前实现可能没覆盖所有组合。

对 DeepSeek-OCR 模型:端到端视觉编码延迟降低 5-17%(小图 224x224 达 13.93%,大图 1024x1024 约 4.75%)。对用户无功能影响,但需显存更多。对系统:CUDA 图管理器新增双路径模式,但仅当模型配置 enable_dual_path_graph=True 时激活,不影响其他 ViT 模型。对其他模型:接口扩展仅增加可选参数,向后兼容。

测试跳过 显存增加 核心路径变更 多模型接口适配

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论