Prhub

#50330 [CI] Organize speculative decoding E2E tests by coverage

原始 PR 作者 mgoin 合并时间 2026-08-01 09:05 文件变更 39 提交数 9 评论 30 代码增减 +1941 / -1863

执行摘要

拆分规格解码 E2E 测试 ,CI 改目录作用域并修 DSpark 量化

PR body 明确说明目标:"Split the monolithic speculative-decoding E2E suite into feature directories that match CI ownership." 以及 "Make Buildkite run directory scopes instead of keyword/file allowlists, so tests added to an existing area are discovered by default." 同时 "Restore nightly DFlash and DSpark acceptance/accuracy coverage, using poolside/Laguna-XS-2.1-NVFP4 to keep Laguna practical on smaller GPUs"。作者还承认这是 "intentionally opinionated: directory boundaries are the CI contract.",并明确三个高 footprint Eagle 案例与四-GPU job 不回归。

值得精读。这是"目录即 CI 契约"的完整落地范例,对测试组织、CI 配置和源码 bug 联动修复都有参考价值。重点看三点:spec_decode.yaml 的目录作用域写法、utils.py 中 compute_acceptance_rate/compute_acceptance_len 的增量统计设计(支持逐 prompt 评估取平均)、以及 dspark/utils.py 里 quant_config 覆写与 #49133 的关联。若团队也在治理巨型 E2E 测试文件,本 PR 的拆分顺序和 follow-up 边界控制值得借鉴。

讨论亮点

评审人 AndreasKaratzas 在多个新测试文件(dspark/test_dspark.py、medusa、synthetic、eagle/utils.py、mtp、ngram_suffix)反复建议改用 vllm_runner 替代裸 LLM,理由是 "i would encourage vllm runners for this. This test does not run on AMD hardware, but trust me it will be more stable even on CUDA CI.";mgoin 回应:"Yes these are valid improvements but I'd like to scope this PR to just moving around existing tests and getting them working again. Improving or refactoring them should be a separate change"。关于 gemma4_dspark.py 的 causal 修复,mgoin 自述 "This is a genuine fix to Gemma4 DSpark since #48167 landed";对 load_dspark_model 的 quant_config 覆写,mgoin 表示 "I believe this will be a more general and simple solution to the problem #49133 tries to solve",并注明 "dsv4 needs more changes, will reserve that for a followup"。AndreasKaratzas 还建议 test-amd.yaml 增加 AMD mirror 镜像。最终其 APPROVED:"Let's merge this assuming that CI passes and work on the modifications that I proposed in a follow-up PR, to keep things containerized"。

实现拆解

  1. 拆分巨型测试文件:删除 tests/v1/e2e/spec_decode/test_spec_decode.py(-1679 行),将用例按功能迁入 draft_model/test_draft_model.py、ngram_suffix/test_ngram_suffix.py、eagle/test_eagle_correctness.py、mtp/test_mtp.py、speculators/test_speculators.py 以及 acceptance_rates/{dflash,dspark,mtp_other} 等新文件;每个文件保持原有用例语义,风险是拆分过程中可能丢用例或破坏 import。
  2. 收敛共享工具与 fixture:新增 tests/v1/e2e/spec_decode/utils.py,承载 get_test_prompts、greedy_sampling、stochastic_sampling、evaluate_llm_for_gsm8k、compute_acceptance_rate、compute_acceptance_len 等 helper;新增 conftest.py 提供 sampling_config、model_name、reset_torch_dynamo 等跨目录 fixture;acceptance_rates/utils.py 再抽出 load_and_process_dataset、run_acceptance_length_eval,供 DFlash 与 DSpark 复用。draft_model 里 ArgsTest dataclass 统一了正确性/接受率/GSM8K 精度多个维度的参数描述,后续加组合只需追加一条数据。
  3. CI 改为目录作用域:.buildkite/test_areas/spec_decode.yaml 用 pytest -v -s v1/e2e/spec_decode/speculators/ 等目录级命令替代旧的 -k "speculators or mtp_correctness" 关键词过滤;engine.yaml 移除被迁移走的条目;.buildkite/test-amd.yaml 同步拆分命令,避免 AMD job 遗漏。目录成为测试归属的唯一契约。
  4. 恢复 DFlash/DSpark 覆盖:test_dflash.py 以 DFlashCorrectnessConfig 声明 Qwen3-8B 与 Laguna-XS-2.1-NVFP4 两组配置,校验 GSM8K accuracy 与 compute_acceptance_len,并保留 mrv1/mrv2 双 runner 参数化;test_dspark.py 恢复 Qwen3-4B-FP8 与 Gemma4 NVFP4 两个 probabilistic 采样路径的接受率/接受长度/GSM8K 三维回归,阈值按实测均值乘 0.9 抗 flaky。
  5. 源码配套修复两处:vllm/model_executor/models/gemma4_dspark.py 在 DFlash 层初始化处补 self.causal = _dflash_layer_causal(config, layer_idx),修复 #48167 之后的 causal 元数据回归;vllm/v1/worker/gpu/spec_decode/dspark/utils.py 的 load_dspark_model 中,在 VllmConfig post-init 后覆写 draft_vllm_config.quant_config = get_draft_quant_config(vllm_config),隔离草稿的量化配置。
  6. 配套微调:提交历史显示还修复了 large_gpu_mark 造成的误跳过、增加 speculators job 超时、从 model_runner_v2.yaml 移除 MTP 案例等,属于拆分后让 CI 变绿的必要收尾。
文件 模块 状态 重要度
tests/v1/e2e/spec_decode/test_spec_decode.py 解码测试 removed 7.49
tests/v1/e2e/spec_decode/draft_model/test_draft_model.py 草稿模型 added 8.02
tests/v1/e2e/spec_decode/utils.py 测试工具 added 7.3
vllm/v1/worker/gpu/spec_decode/dspark/utils.py 草稿加载 modified 5.63
vllm/model_executor/models/gemma4_dspark.py 模型层 modified 5.96
.buildkite/test_areas/spec_decode.yaml CI 配置 modified 5.53
tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py 匹配解码 added 7.0
tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py 接受率回归 added 7.02
tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py 接受率回归 added 6.94
tests/v1/e2e/spec_decode/acceptance_rates/utils.py 测试工具 added 6.32
tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py 草稿回归 added 6.5
tests/v1/e2e/spec_decode/mtp/test_mtp.py 多词预测 added 6.2
.buildkite/test_areas/engine.yaml CI 配置 modified 4.55

关键符号

test_draft_model_correctness test_draft_model_engine_args_tensor_parallelism compute_acceptance_rate compute_acceptance_len run_acceptance_length_eval load_and_process_dataset test_dflash_correctness test_dflash_reference_acceptance_lengths test_dspark_correctness_and_acceptance_rate test_gemma4_dspark_correctness_and_acceptance_rate test_ngram_and_suffix_correctness test_suffix_decoding_acceptance test_mtp_correctness _run_eagle_correctness test_eagle_correctness_light test_eagle_correctness_medium test_speculators_model_integration load_dspark_model

关键源码片段

tests/v1/e2e/spec_decode/draft_model/test_draft_model.py test-coverage

本 PR 最大的新增测试文件(+388 行),集中验证 draft_model 方式的正确性、量化组合、TP>1 与 EngineArgs 配置独立性,ArgsTest dataclass 是后续扩展组合的模板。

# tests/v1/e2e/spec_decode/draft_model/test_draft_model.py
# 拆分自原 test_spec_decode.py: 集中验证 draft_model 规格解码方式的正确性、
# 量化组合、TP>1 与 EngineArgs 配置独立性。
​
from dataclasses import dataclass
​
import pytest
​
from vllm import LLM, SamplingParams
from vllm.config import VllmConfig, replace
from vllm.engine.arg_utils import EngineArgs
​
from ..utils import (
    Messages,
    _skip_if_insufficient_gpus_for_tp,
    evaluate_llm_for_gsm8k,
    get_instruct_coder_messages,
    get_test_prompts,
    greedy_sampling,
    stochastic_sampling,
)
​
​
class AsyncSchedulingNotEnabledError(AssertionError):
    """draft_model 规格解码未启用 async_scheduling 时抛出。"""
​
​
@dataclass
class ArgsTest:
    """用一个 dataclass 统一描述一组规格解码验证参数。
​
    后续新增 draft_model 组合场景时,只需追加一条 ArgsTest,
    参数化会自动展开,避免为每种组合重复写引擎初始化样板。
    """
    target_model: str
    draft_model: str
    sampling_config: SamplingParams
    num_speculative_tokens: int
    expected_acceptance_rate: float
    expected_acceptance_len: float
    # 默认跳过 GSM8K 精度校验 ; 需要时显式给出阈值
    expected_gsm8k_accuracy: float = 0.0
    enforce_eager: bool = True
    parallel_drafting: bool = False
    target_tensor_parallel_size: int = 1
    draft_tensor_parallel_size: int = 1
    max_model_len: int = 2048
    gpu_memory_utilization: float = 0.5
    dataset: str = "test_prompts"
    num_prompts: int = 100
​
​
def get_messages(dataset: str, n: int) -> list[Messages]:
    """按数据集名字路由到不同的 prompt 构造器,测试数据源可插拔。"""
    if dataset == "test_prompts":
        return get_test_prompts(mm_enabled=False, num_prompts=n)
    if dataset == "likaixin/InstructCoder":
        return get_instruct_coder_messages(n=n)
    raise NotImplementedError(f"Dataset '{dataset}' not implemented")
​
​
# 覆盖两类代表性组合 : 同模型贪心 ( 接受率上限 )、小模型随机采样 ( 真实收益 )
cases = [
    ArgsTest(
        target_model="Qwen/Qwen3-0.6B",
        draft_model="Qwen/Qwen3-0.6B",
        sampling_config=greedy_sampling(),
        num_speculative_tokens=3,
        expected_acceptance_len=0.98 * (3 + 1), # K+1 的 epsilon 折扣
        expected_acceptance_rate=0.98,
        expected_gsm8k_accuracy=0.25,
    ),
    ArgsTest(
        target_model="Qwen/Qwen3-1.7B",
        draft_model="Qwen/Qwen3-0.6B",
        sampling_config=stochastic_sampling(),
        num_speculative_tokens=3,
        expected_acceptance_len=3.4,
        expected_acceptance_rate=0.80,
        expected_gsm8k_accuracy=0.5,
    ),
]
​
​
# draft_model 尚未启用 async_scheduling, 统一 xfail 到 issue #38929,
# 避免 CI 被已知短板拖红
@pytest.mark.parametrize("args", cases)
@pytest.mark.parametrize("enforce_eager", [True, False])
@pytest.mark.xfail(
    raises=AsyncSchedulingNotEnabledError,
    reason="draft_model does not yet enable async_scheduling: issue #38929",
)
def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool):
    args.enforce_eager = enforce_eager
    assert_draft_model_correctness(args)
tests/v1/e2e/spec_decode/utils.py test-coverage

共享测试工具的统一出口,compute_acceptance_rate 与 compute_acceptance_len 支持累计指标增量计算,是 DFlash/DSpark/Medusa 所有接受率回归的基础。

# tests/v1/e2e/spec_decode/utils.py
# 公共测试工具 : 原先散落在巨型测试文件里的 helper 全部收敛到这里 ,
# 供 draft_model、eagle、mtp、ngram_suffix 等子目录共享。
​
​
def compute_acceptance_rate(
    metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
    """计算投机解码接受率(acceptance rate)。
​
    vLLM 暴露的指标是累计值,传入 prev_metrics 时计算两次采样之间的
    增量,适配逐条 prompt 评估后取平均的场景。
    """
    name2metric = {metric.name: metric for metric in metrics}
    n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value
    if n_draft_toks == 0:
        return float("nan") # 无草稿 token 时无法定义接受率
    n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
    if prev_metrics is not None:
        prev_name2metric = {metric.name: metric for metric in prev_metrics}
        n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value
        n_accepted_toks -= prev_name2metric[
            "vllm:spec_decode_num_accepted_tokens"
        ].value
        if n_draft_toks <= 0:
            return float("nan")
    return n_accepted_toks / n_draft_toks
​
​
def compute_acceptance_len(
    metrics: list[Metric], prev_metrics: list[Metric] | None = None
) -> float:
    """计算平均接受长度:期望每轮投机接受的 token 数,含第一个确定 token。"""
    name2metric = {metric.name: metric for metric in metrics}
    n_drafts = name2metric["vllm:spec_decode_num_drafts"].value
    n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value
    if n_drafts == 0:
        return 1 # 没有投机轮次时退化为单 token 生成
    if prev_metrics is not None:
        prev_name2metric = {metric.name: metric for metric in prev_metrics}
        n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value
        n_accepted_toks -= prev_name2metric[
            "vllm:spec_decode_num_accepted_tokens"
        ].value
        if n_drafts <= 0:
            return 1
    return 1 + (n_accepted_toks / n_drafts)
vllm/v1/worker/gpu/spec_decode/dspark/utils.py dependency-wiring

源码级修复的关键点:在 VllmConfig post-init 后覆写草稿量化配置,隔离 NVFP4 target 与 BF16 草稿,是激活 Gemma4 DSpark 测试后暴露的真实 bug。

# vllm/v1/worker/gpu/spec_decode/dspark/utils.py
# load_dspark_model 中的关键修复 : 草稿模型的量化配置必须与 target 隔离。
​
# VllmConfig 的 post-init 逻辑会把 target 的 quant_config 带回草稿配置 ,
# 而 target 配置需要保留用于 DSpark 的目标层元数据 , 因此这里必须
# 在构造完成后覆写一次草稿的量化配置 , 否则 NVFP4 target + BF16 草稿
# 这类组合会错误地按 NVFP4 加载草稿权重。
draft_vllm_config.quant_config = get_draft_quant_config(vllm_config)

评论区精华

改用 vllm_runner 提升 E2E 稳定性 测试

AndreasKaratzas 在 dspark、medusa、synthetic、eagle、mtp、ngram_suffix 多处建议用 vllm_runner 替代裸 LLM,称 "i would encourage vllm runners for this. This test does not run on AMD hardware, but trust me it will be more stable even on CUDA CI."

结论:mgoin 认可为有效改进,但将本 PR 范围限定为 " 移动现有测试并恢复运行 ",统一改造留 follow-up。 · 已解决

Gemma4 DSpark causal 元数据修复 正确性

mgoin 在 gemma4_dspark.py 评论中自评 self.causal = _dflash_layer_causal(config, layer_idx) 是 #48167 落地之后引入的真实回归修复。

结论:确认合入该修复,由 Gemma4 DSpark E2E 用例长期守护。 · 已解决

草稿 quant_config 覆写与 #49133 的关系 正确性

mgoin 认为在 load_dspark_model 中覆写 draft_vllm_config.quant_config 是比 #49133 更通用简单的解法,"I believe this will be a more general and simple solution to the problem #49133 tries to solve";同时注明 DSv4 NVFP4 还需 follow-up。

结论:本 PR 覆盖 Gemma4 NVFP4/BF16 场景,DSv4 NVFP4 后续处理。 · 已解决

AMD 平台镜像与平台差异处理 question

AndreasKaratzas 建议在 .buildkite/test-amd.yaml 增加 AMD mirror 镜像,并建议 draft_model 的 moe_backend 测试按 CUDA/ROCm 分支选择 backend 或跳过。

结论:平台兼容性改进未在本 PR 落实,留 follow-up;mgoin 表示保持本 PR 聚焦。 · 待处理

验收阈值、tqdm 与随机种子细节 设计

AndreasKaratzas 质疑 CI 中 tqdm 的必要性("Do we really need tqdm in CI?"),建议 speculators 阈值更激进、utils.py 使用全局 vllm 随机种子以保证 CI 稳定性。

结论:保持现状合并,阈值保守化以抗 flaky;改进项列入 follow-up。 · 待处理

风险与影响

  • 大文件删除风险:1679 行文件整体删除,仓库内虽已清理 import,但外部 fork 或未搜索到的引用会直接断裂;此类重构合入后应尽快让全量 CI 跑一遍。
  • CI 目录作用域静默覆盖风险:Buildkite 改为按目录收集测试后,新测试若放错目录不会被 spec_decode job 发现,造成覆盖静默丢失;"目录即契约"依赖命名纪律。
  • 模型语义修复风险:vllm/model_executor/models/gemma4_dspark.py 的 causal 赋值改变注意力 mask 语义,虽经 E2E 验证,但属于推理路径正确性变更,需 DSpark 测试长期守护。
  • 量化配置隔离不完整:vllm/v1/worker/gpu/spec_decode/dspark/utils.py 的覆写只验证了 Gemma4 NVFP4/BF16 场景,mgoin 自述 DSv4 NVFP4 仍需 follow-up,存在同路径其他模型组合未覆盖的隐患。
  • 阈值 flaky 风险:多个 acceptance rate/len 阈值是经验值乘 0.9 裕量,在 probabilistic 采样路径下仍可能因 GPU、驱动或模型下载差异抖动。
  • CI 基础设施:spec_decode.yaml、engine.yaml、test-amd.yaml 等多个面积从关键词匹配改为目录作用域,后续新增测试的文件位置直接决定 CI 归属,影响所有 spec decode 相关开发者的提交流程。
  • 测试维护者:需要遵循新的目录约定,把用例放入对应 feature 目录;共享 utils.py 降低了跨用例的重复代码。
  • 产品功能:无用户可见功能变化;间接获得 DFlash、DSpark 与 Gemma4 NVFP4 DSpark 的持续回归保护。
  • 团队协作:9 个 commit 中多次 merge main、多次针对 CI 失败打补丁(超时、skip、平台差异),展示了"先拆分、再恢复绿灯"的渐进式迁移节奏;AI 辅助(OpenAI Codex)参与生成,作者声明已人工复核。
巨型测试文件拆分 (1679 行删除 ) CI 目录作用域静默覆盖风险 模型语义修复 (causal 元数据 ) DSpark NVFP4 仅覆盖部分场景 接受率阈值经验值可能 flaky

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论