Prhub

#46683 Bump flashinfer version to 0.6.13

原始 PR 作者 wzhao18 合并时间 2026-06-30 00:30 文件变更 6 提交数 13 评论 3 代码增减 +29 / -15

执行摘要

FlashInfer 升级 0.6.13 并启用持久缓存

FlashInfer 0.6.13修复了持久缓存中use_8x4_sf_layout配置冲突导致选用无效tactic的问题,因而可以默认启用持久缓存来避免每次新配置都从头autotune,显著缩短模型加载时间。同时,DeepEP MoE场景下仅rank 0执行autotune会导致其他rank超时,因此需为该场景保留非持久模式。GSM8K测试中由于autotune耗时增加,需要更大的启动超时窗口。详见Review讨论中关于fp4_gemm skip_ops的权衡。

建议仔细审阅flashinfer_autotune函数中DeepEP的判断逻辑,确认覆盖了所有相关的all2all后端。同时关注FlashInfer上游后续是否彻底解决持久缓存问题,届时可移除特判。GSM8K超时调整作为配套,无特别风险。

讨论亮点
  • [设计权衡] wzhao18询问LopezCastroRoberto是否需要skip_ops="fp4_gemm"来避免cutedsl NVFP4 GEMM的自动调优。Lopez确认这是当前workaround,但会引入首次查询50us和后续0.2us的透顶开销,建议等待更长期的解决方案。目前该PR未采用此跳过策略。
  • [测试通过] mgoin审核后批准,确认失败测试与本次变更无关。

实现拆解

  1. 版本号更新:在requirements/cuda.txtflashinfer-pythonflashinfer-cubin版本从0.6.12改为0.6.13;在docker/versions.jsonFLASHINFER_VERSION默认值更新;在docker/Dockerfiledocker/Dockerfile.nightly_torch中克隆分支或pip安装参数同步更新。

  2. Autotune持久缓存逻辑重构:在vllm/model_executor/warmup/kernel_warmup.py中:

    • 移除全局禁用标志_FLASHINFER_USE_PERSISTENT_CACHE = False
    • flashinfer_autotune函数中默认启用持久缓存(use_persistent_cache = True);
    • 新增对DeepEP后端(deepep_high_throughputdeepep_low_latencydeepep_v2)的检测:如果all2all_backend匹配,则设use_persistent_cache = False,让所有rank共同执行autotune以避免DeepEP分发/结合超时;
    • 根据use_persistent_cache走不同的autotune路径:启用时由leader rank写缓存并广播,禁用时全体rank运行autotune并barrier同步。
  3. GSM8K测试启动超时调整:在tests/evals/gsm8k/test_gsm8k_correctness.py中:

    • 定义DEFAULT_STARTUP_MAX_WAIT_SECONDS = 1200
    • 从配置中读取startup_max_wait_seconds(默认1200);
    • 构建环境字典时自动注入VLLM_ENGINE_READY_TIMEOUT_S为该值;
    • 传递给RemoteOpenAIServermax_wait_seconds也使用该值。
  4. Docker建造流程适配:更新FLASHINFER_VERSION构建参数,确保镜像使用新版本。

文件 模块 状态 重要度
vllm/model_executor/warmup/kernel_warmup.py 预热模块 modified 6.65
tests/evals/gsm8k/test_gsm8k_correctness.py GSM8K 测试 modified 4.3
docker/versions.json 容器配置 modified 2.95
docker/Dockerfile 容器构建 modified 2.78
docker/Dockerfile.nightly_torch 容器构建 modified 2.96
requirements/cuda.txt 依赖配置 modified 2.07

关键符号

flashinfer_autotune test_gsm8k_correctness

关键源码片段

vllm/model_executor/warmup/kernel_warmup.py core-logic

核心 autotune 逻辑变更:启用持久缓存并添加 DeepEP 例外

def flashinfer_autotune(runner: "GPUModelRunner") -> None:
    """
    Autotune FlashInfer operations.
    FlashInfer have many implementations for the same operation,
    autotuning runs benchmarks for each implementation and stores
    the results. The results are cached transparently and
    future calls to FlashInfer will use the best implementation.
    Without autotuning, FlashInfer will rely on heuristics, which may
    be significantly slower.    Tuning is performed only on rank 0. The resulting cache is broadcast
    to every rank so all ranks dispatch the same kernel tactic.
    """
    import vllm.utils.flashinfer as fi_utils
    from vllm.distributed.parallel_state import get_world_group
​
    # 默认启用持久缓存(FlashInfer 0.6.13 已修复冲突)
    use_persistent_cache = True
​
    # 检测 DeepEP all2all 后端,避免仅 rank0 执行 autotune 导致超时
    deepep_a2a_backends = {
        "deepep_high_throughput", "deepep_low_latency", "deepep_v2",
    }
    if runner.vllm_config.parallel_config.all2all_backend in deepep_a2a_backends:
        use_persistent_cache = False
​
    if not use_persistent_cache:
        # 非持久缓存路径:所有 rank 执行 autotune 并 barri 同步
        with torch.inference_mode(), fi_utils.autotune():
            runner._dummy_run(
                num_tokens=runner.scheduler_config.max_num_batched_tokens,
                skip_eplb=True, is_profile=True,
            )
        get_world_group().barrier()
        return
​
    # 持久缓存路径:leader 写入缓存文件,其他 rank 等待广播
    world = get_world_group()
    is_leader = world.rank_in_group == 0
    cache_path = resolve_flashinfer_autotune_file(runner)
    if is_leader:
        logger.info("Using FlashInfer autotune cache file: %s", cache_path)
​
    dummy_run_kwargs = dict(
        num_tokens=runner.scheduler_config.max_num_batched_tokens,
        skip_eplb=True, is_profile=True,
    )
​
    with torch.inference_mode():
        if is_leader:
            with fi_utils.autotune(tune_mode=True, cache=str(cache_path)):
                runner._dummy_run(**dummy_run_kwargs)
        else:
            # 非 leader 分支:等待 leader 完成 autotune 并加载缓存
            # (原实现延续此处,未因本 PR 变更)
            pass
tests/evals/gsm8k/test_gsm8k_correctness.py test-coverage

测试启动超时调整以适应 autotune 耗时增加

DEFAULT_STARTUP_MAX_WAIT_SECONDS = 1200 # 新增默认启动超时def test_gsm8k_correctness(config_filename):
    # ... 前面的代码(跳过高)
    # 从配置读取启动超时,默认 1200s
    startup_max_wait_seconds = eval_config.get(
        "startup_max_wait_seconds", DEFAULT_STARTUP_MAX_WAIT_SECONDS
    )
    # 构建环境变量,注入引擎就绪超时
    env_dict = dict(eval_config.get("env") or {})
    env_dict["VLLM_ENGINE_READY_TIMEOUT_S"] = str(int(startup_max_wait_seconds))
​
    # 打印详细信息
    print(f"Startup max wait: {startup_max_wait_seconds}s")
    print(f"Environment variables: {env_dict}")
​
    # 启动 server(使用调整后的 env_dict 和 max_wait_seconds)
    with RemoteOpenAIServer(
        eval_config["model_name"],
        server_args,
        env_dict=env_dict,
        max_wait_seconds=startup_max_wait_seconds, # 之前是硬编码 600
    ) as remote_server:
        # ...
        results = run_gsm8k_eval(eval_config, server_url)
        # ...

评论区精华

是否需要跳过 fp4_gemm autotune 设计

wzhao18 引用上游 issue 询问是否应该用 `with flashinfer.autotune(skip_ops="fp4_gemm")` 来避免 cutedsl NVFP4 GEMM 的长时间 autotune。LopezCastroRoberto 确认这是 workaround,但会引入首次查询 50us 和后续 0.2us overhead,建议等待更好的长期方案。

结论:当前 PR 未采用 skip_ops,保留完整 autotune,但此讨论提示未来改进方向。 · 已解决

风险与影响

  • 兼容性风险:升级FlashInfer版本可能带来新的API行为变化,但当前变更已在kernel_warmup中适配。若用户自定义了FlashInfer调用,需确保接口兼容。
  • DeepEP超时风险:虽然为DeepEP场景禁用了持久缓存,但全体rank执行autotune仍可能导致部分rank落后,尤其是大规模集群。需关注autotune完成后的barrier同步是否可靠。
  • 测试超时掩蔽:GSM8K启动超时从600s提升到1200s,可能掩盖真正的服务器启动失败(如模型加载错误),需辅以其他异常检测。
  • 持久缓存一致性:之前因冲突禁用缓存,现在重新启用,若仍有未修复的冲突可能导致推理错误。依赖上游0.6.13的修复完备性。
  • 对用户:非DeepEP场景下模型加载速度提升(跳过重复autotune);DeepEP场景下无感知但保持稳定性。GSM8K测试更可靠。
  • 对系统:减少了不必要的autotune计算,降低启动阶段GPU负载。
  • 对团队:需维护DeepEP后端的特化逻辑;未来FlashInfer版本升级时需重新评估持久缓存策略。
持久缓存依赖上游修复 DeepEP 超时风险 测试超时掩蔽

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论