Prhub

#48341 [Bugfix][Spec Decode] Auto-enable async scheduling for draft models

原始 PR 作者 BWAAEEEK 合并时间 2026-08-06 13:51 文件变更 3 提交数 1 评论 6 代码增减 +25 / -39

执行摘要

draft_model 默认启用异步调度并移除 xfail

issue #38929 报告 draft_model 推测解码测试的两类失败:Expected async_scheduling=True for draft_model spec decode, got False. 的断言错误,以及引擎核心初始化失败。PR body 指出根因:VllmConfig 在显式开启 async_scheduling 时允许 Draft Model,但默认解析路径(async_scheduling is None 分支)把 draft_model 当作不支持的方法而默认禁用,造成默认行为与显式行为不一致,并使相关正确性测试长期 xfail。本 PR 的目标是把默认解析与显式开启对齐,无需用户手动设置即可自动启用异步调度。

建议快速精读,3 个文件、1 行核心改动,几分钟即可读完:

  • 最值得关注的设计VllmConfig.__post_init__ 中异步调度默认解析的完整分支链(pooling → speculative 方法白名单 → disable_padded_drafter_batch → executor 支持性 → ROCm DeepEP),是理解 vLLM 配置默认值“先收窄、再放宽”策略的好样例。
  • 可借鉴的测试模式:用“定制异常 + xfail(raises=...)”在缺陷未修复时精确标记失败、避免 CI 噪音,修复后再收敛为普通断言,让回归第一时间可见。
    • 若负责 speculative-decoding 或配置系统,建议阅读 tests/test_config.py 新增用例作为默认值回归测试的模板。
讨论亮点

该 PR 没有任何 inline review 评论,时间线主要被自动化机器人占据,技术讨论几乎为零:

  • mergify[bot] 曾提示存在 merge conflict,需要 rebase 后才能合并。
  • mgoin 通过 /ci run 批准并触发 Buildkite CI(#82523),作者随后对最新提交再次触发 CI(#82616)。
  • 唯一有价值的设计信息来自被删除的代码注释:原 AsyncSchedulingNotEnabledError 继承 AssertionError 的目的,是让 xfail(raises=AsyncSchedulingNotEnabledError) 只捕获“异步调度未开启”这一种失败,而正确性、接受率等其他断言失败仍会作为真实失败暴露。修复落地后,这个“精确标记已知缺陷”的机制与 xfail 一起被移除,收敛为普通 assert has_async
  • PR body 强调作者已检索 issue #38929 评论及所有相关 PR,确认没有其他 PR 处理同一默认值修复,避免了重复工作。

实现拆解

本变更通过 5 个步骤完成配置默认值修复与测试恢复:

  1. 定位默认解析入口vllm/config/vllm.pyVllmConfig.__post_init__async_scheduling 的唯一默认值解析点。当用户未显式设置(值为 None)时,代码按优先级依次检查 pooling 模型、speculative 方法白名单、disable_padded_drafter_batch、executor 后端支持性和 ROCm DeepEP 这几个不兼容条件,任何一条命中就关闭异步调度。
  2. 修改核心判断:在“方法不在 EAGLE / NGram GPU / DSpark 白名单则禁用”的条件里新增 self.speculative_config.method != "draft_model",使 draft_model 不再落入禁用分支。由于显式开启分支本来就允许 Draft Model,修复后默认解析与显式行为保持一致(对应 issue #38929 中的 async_scheduling=False 断言失败)。
  3. 保留兜底保护:新增判断只影响 None 分支;disable_padded_drafter_batchexecutor_supports_async_scheduses_rocm_deepep_ht_dbo 等后续检查原样保留,默认开启不等于无条件开启。
  4. 补充配置层回归测试tests/test_config.py 新增 test_draft_model_enables_async_scheduling_by_default,用 ParallelConfig(distributed_executor_backend="uni")SpeculativeConfig(method="draft_model", ...) 构造最小配置,断言 cfg.scheduler_config.async_scheduling is True
  5. 清理 xfail 并恢复强断言tests/v1/e2e/spec_decode/draft_model/test_draft_model.py 删除 AsyncSchedulingNotEnabledError 异常类、5 处 xfail 标记及相关 TODO 注释,把“抛定制异常”改为普通 assert has_async,让正确性与接受率断言重新成为真实回归信号。

验证配套:本地单卡 11 项测试与双卡 TP=2 测试通过(2x NVIDIA B200),GSM8K 准确率 0.673,接受率 0.93、接受长度 3.78,pre-commit 与 DCO 检查通过。

文件 模块 状态 重要度
vllm/config/vllm.py 配置解析 modified 4.98
tests/test_config.py 配置测试 modified 4.94
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py 推测解码 modified 5.73

关键符号

VllmConfig.__post_init__ test_draft_model_enables_async_scheduling_by_default assert_draft_model_correctness

关键源码片段

vllm/config/vllm.py core-logic

核心修复文件:在 `VllmConfig.__post_init__` 的默认解析分支中新增 `method != "draft_model"` 判断,使 `draft_model` 推测解码默认启用异步调度,与显式开启行为对齐。1 行新增,全部逻辑价值所在。

        elif self.scheduler_config.async_scheduling is None:
            # 未显式指定 async_scheduling 时,按优先级自动解析默认值。
            # 分支顺序很关键:pooling 模型、不支持的 speculative 方法、
            # disable_padded_drafter_batch、executor 后端、ROCm DeepEP 依次拦截。
            if (
                self.model_config is not None
                and self.model_config.runner_type == "pooling"
            ):
                # 异步调度会拖慢 pooling 模型,因此默认关闭。
                logger.debug(
                    "Disabling asynchronous scheduling by default for pooling model."
                )
                self.scheduler_config.async_scheduling = False
            elif (
                self.speculative_config is not None
                and self.speculative_config.method not in get_args(EagleModelTypes)
                and self.speculative_config.method not in get_args(NgramGPUTypes)
                # 本次修复的关键一行:把 draft_model 从“不支持清单”中排除,
                # 使 draft_model 推测解码默认启用异步调度,与显式开启行为一致。
                and self.speculative_config.method != "draft_model"
                and self.speculative_config.method != "dspark"
            ):
                logger.warning_once(
                    "Async scheduling not supported with %s-based "
                    "speculative decoding and will be disabled.",
                    self.speculative_config.method,
                )
                self.scheduler_config.async_scheduling = False
            elif (
                self.speculative_config is not None
                and self.speculative_config.disable_padded_drafter_batch
            ):
                logger.warning_once(
                    "Async scheduling is not compatible with "
                    "disable_padded_drafter_batch=True and will be disabled.",
                )
                self.scheduler_config.async_scheduling = False
            elif not executor_supports_async_sched:
                # 后端不支持时仍然兜底关闭,保证默认路径与显式路径行为安全。
                logger.warning_once(
                    "Async scheduling will be disabled because it is not supported "
                    "with the `%s` distributed executor backend. ",
                    executor_backend,
                )
                self.scheduler_config.async_scheduling = False
            elif uses_rocm_deepep_ht_dbo:
                logger.warning_once(
                    "Async scheduling is disabled for ROCm DeepEP "
                    "high-throughput DBO because that combination can corrupt "
                    "DP+EP generation accuracy."
                )
                self.scheduler_config.async_scheduling = False
tests/test_config.py test-coverage

新增纯配置层回归测试 `test_draft_model_enables_async_scheduling_by_default`,防止将来再次把 `draft_model` 归入不支持列表而回归。

def test_draft_model_enables_async_scheduling_by_default():
    # 构造最小但完整的配置组合:draft_model 推测解码 + uni 分布式后端,
    # 不显式设置 async_scheduling,走默认解析路径。
    parallel_config = ParallelConfig(distributed_executor_backend="uni")
    model_config = ModelConfig("Qwen/Qwen3-0.6B", max_model_len=2048)
    speculative_config = SpeculativeConfig(
        method="draft_model",
        model="Qwen/Qwen3-0.6B",
        num_speculative_tokens=3,
        target_model_config=model_config,
        target_parallel_config=parallel_config,
    )
    cfg = VllmConfig(
        model_config=model_config,
        scheduler_config=SchedulerConfig(
            max_model_len=2048,
            is_encoder_decoder=False,
        ),
        parallel_config=parallel_config,
        speculative_config=speculative_config,
    )
​
    # 回归断言:默认解析后 async_scheduling 必须为 True,
    # 防止将来再次把 draft_model 归入“不支持”列表而禁用异步调度。
    assert cfg.scheduler_config.async_scheduling is True
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py test-coverage

删除 5 处 xfail、`AsyncSchedulingNotEnabledError` 定制异常及相关 TODO,将异步调度断言收敛为普通 assert,恢复 e2e 测试的真实回归能力。

    # 原实现用 AsyncSchedulingNotEnabledError 包装该断言,配合 xfail 精确捕获
    # issue #38929 的已知缺陷;修复后收敛为普通 assert,任何回归都会直接失败。
    has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling
    del spec_llm # CLEANUP
    torch.accelerator.empty_cache()
    cleanup_dist_env_and_memory()
    assert has_async, "Expected async_scheduling=True for draft_model spec decode"

评论区精华

无实质 review 讨论,仅 CI 与冲突处理交互 other

PR 没有任何 inline review 评论。评论时间线仅包含 mergify[bot] 的冲突提示、github-actions[bot] 的 CI 触发确认、mgoin 的 `/ci run` 批准以及作者对最新提交的二次 `/ci run`。

结论:无技术性分歧;mgoin 批准后合入。值得留意的是测试代码中曾经存在的设计:用 `AsyncSchedulingNotEnabledError`(继承 `AssertionError`)配合 `xfail(raises=...)` 精确标记已知缺陷,修复后该机制随 xfail 一起移除。 · 已合入

风险与影响

风险点集中在默认行为变更的波及面:

  • 默认行为变更:凡是使用 draft_model 且未显式设置 async_scheduling 的部署,默认值从 False 变为 True。虽然这是 issue #38929 声明的预期行为,但仍属于全局配置默认值的变化;好在该变更处在 VllmConfig.__post_init__ 的 elif 链中,disable_padded_drafter_batch、executor 后端支持性、ROCm DeepEP 等兜底检查(vllm/config/vllm.py)仍然生效,不支持异步调度的环境不会受影响。
  • 显式路径不受影响:用户显式设置 async_scheduling=FalseTrue 时走的是另外两个分支,本次改动不涉及。
  • e2e 测试恢复强断言:5 个测试从 xfail 恢复为真实执行,依赖 GPU 与 HuggingFace 模型下载的 CI 任务可能因环境波动出现偶发失败;这是修复后的正常成本,但值得 CI 关注。
    • 无安全、性能相关风险;异步调度的运行时实现未改动。

影响评估如下:

  • 用户影响:使用 draft_model 推测解码(如 Qwen/Qwen3-0.6B 作 draft)的用户无需再手动开启异步调度,默认行为与显式设置一致,消除了 issue #38929 报告的行为不一致。
  • 系统影响:仅触及 v1 引擎配置解析默认值(vllm/config/vllm.py 一处条件),调度与执行运行时代码零改动。
  • 团队与工程影响:5 个 GPU e2e 测试从 xfail 恢复为强断言,speculative-decoding 回归防护显著增强;tests/test_config.py 新增的配置层测试为后续 spec decode 默认行为变更提供了可复用的回归模板。
  • 影响程度:中低。改动面小(1 行核心逻辑)、有兜底分支保护,但触点位于全局默认配置,且恢复了多个依赖真实模型下载的 e2e 测试。
默认行为变更 e2e 测试依赖 GPU 与模型下载

关联 Issue

#38929 [Bug]: Draft model speculative decoding tests failing: async_scheduling not enabled and engine core initialization errors

完整报告

参与讨论