Prhub

#40780 [CI/Build] Add e2e test for ViT CUDA graph

原始 PR 作者 shen-shanshan 合并时间 2026-04-24 18:12 文件变更 2 提交数 4 评论 3 代码增减 +167 / -0

执行摘要

为 ViT CUDA graph 添加端到端测试

ViT CUDA graph是vLLM多模态推理的重要优化,但缺少端到端测试确保其正确性。本PR旨在通过添加e2e测试来验证CUDA graph在ViT编码器上的功能,防止后续修改导致回归。

建议阅读本PR的测试框架设计,尤其是VitCudagraphTestConfig数据类和params_with_marks的使用方式,为多模态测试提供良好模板。后续可考虑采纳reviewer建议改进测试健壮性。

讨论亮点

在review中,gemini-code-assist指出测试资产获取方式脆弱,建议使用IMAGE_ASSETS.get_asset直接获取所需资产而非依赖fixture索引。该建议未被作者回应或采纳,但PR最终获得批准合并。此讨论体现了对测试稳定性的关注,当前实现仍保留基于索引的方式,可能在fixture变动时产生问题。

实现拆解

实现拆解

  1. 创建测试配置类:在tests/models/multimodal/generation/test_vit_cudagraph.py中定义VitCudagraphTestConfig数据类,包含模型、模态、提示、编译配置等参数,方便扩展新模型。

  2. 设计辅助函数params_with_marks将配置字典转pytest参数;qwen_vl_chat_template生成Qwen3-VL聊天模板;get_compilation_config返回启用CUDA graph的配置。

  3. 编写端到端测试test_vit_cudagraph_imagetest_vit_cudagraph_video分别测试图片和视频生成。使用vllm_runner启动推理,验证输出有效。

  4. 集成到CI:在.buildkite/test_areas/models_multimodal.yaml的qwen3+gemma步骤中添加pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model,确保测试自动运行。

关键代码片段

文件 模块 状态 重要度
tests/models/multimodal/generation/test_vit_cudagraph.py ViT-CUDA 图 added 7.84
.buildkite/test_areas/models_multimodal.yaml CI 配置 modified 2.97

关键符号

VitCudagraphTestConfig params_with_marks qwen_vl_chat_template get_compilation_config test_vit_cudagraph_image test_vit_cudagraph_video

关键源码片段

tests/models/multimodal/generation/test_vit_cudagraph.py test-coverage

新增测试文件,包含所有 ViT CUDA graph 端到端测试逻辑,是 PR 的核心变更。

@dataclass
class VitCudagraphTestConfig:
    model: str
    modalities: list[str] = field(default_factory=lambda: ["image", "video"])
    image_prompt: str | None = None
    video_prompt: str | None = None
    dtype: str = "bfloat16"
    max_model_len: int = 4096
    max_tokens: int = 64
    max_num_seqs: int = 2
    num_video_frames: int = 16
    needs_video_metadata: bool = False
    vllm_runner_kwargs: dict = field(default_factory=dict)
    marks: list = field(default_factory=list)
​
​
def get_compilation_config():
    """返回启用ViT CUDA graph的编译配置"""
    return {
        "cudagraph_mm_encoder": True, # 启用多模态编码器 CUDA graph
        "encoder_cudagraph_max_vision_items_per_batch": 1, # 每批最大视觉项数
        "encoder_cudagraph_max_frames_per_batch": 16, # 每批最大视频帧数
    }
​
​
@create_new_process_for_each_test()
def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):
    config = MODEL_CONFIGS[model_id]
    if "image" not in config.modalities:
        pytest.skip(f"{model_id} does not support the image modality.")
    image_prompts = IMAGE_ASSETS.prompts({
        "stop_sign": config.image_prompt,
        "cherry_blossom": config.image_prompt,
    })
    images = [[asset.pil_image] for asset in image_assets]
    with vllm_runner(
        config.model,
        dtype=config.dtype,
        max_model_len=config.max_model_len,
        max_num_seqs=config.max_num_seqs,
        limit_mm_per_prompt={"image": 1},
        compilation_config=get_compilation_config(),
        **config.vllm_runner_kwargs,
    ) as vllm_model:
        outputs = vllm_model.generate_greedy(
            image_prompts, config.max_tokens, images=images
        )
        assert len(outputs) == 2
        output_ids, output_text = outputs[0]
        assert len(output_ids) > 0
        assert len(output_text) > 0
        assert isinstance(output_text, str)

评论区精华

测试图像资产访问脆弱性 测试

gemini-code-assist 建议使用 IMAGE_ASSETS.get_asset 显式获取资产,而非依赖 fixture 索引,提高测试健壮性。

结论:评论未得到作者回复,但 PR 最终获批合并,当前实现保留。 · unresolved

测试视频资产访问脆弱性 测试

gemini-code-assist 建议优化视频资产获取,避免处理所有资产只使用第一个的浪费。

结论:评论未得到作者回复,但 PR 最终获批合并,当前实现保留。 · unresolved

风险与影响

低风险。本PR仅新增测试文件并修改CI配置,未改动任何生产代码。潜在风险包括:测试需要下载Qwen3-VL-2B-Instruct模型,网络波动可能导致超时;测试运行在GPU上,会增加CI资源消耗。但这些风险处于可接受范围。

对用户无直接影响。对团队,测试覆盖确保ViT CUDA graph功能稳定,有利于后续优化或模型扩展。当前仅覆盖Qwen3-VL,但配置类设计使得添加新模型非常容易。

测试依赖模型下载 CI 耗时增加

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论