执行摘要

• 一句话：预热 forward_native 采样内核，避免首次 seed 请求的 JIT 延迟
• 推荐动作：

功能与动机

PR body 指出："The default-on FlashInfer top-k/top-p sampler from #40376 exposed a warmup gap in _dummy_sampler_run: the existing dummy call uses generators={}, which under the new default routes through flashinfer_sample and leaves the forward_native Triton fallback path JIT-cold. That fallback is taken at runtime whenever a request has a seed (SamplingMetadata.generators non-empty), so the first such request pays ~1s of JIT cost on its first sampling step. This makes tests/v1/engine/test_async_llm.py::test_multi_abort fail on main." #41316 原本计划 revert #40376，此 PR 解决了根本原因。

实现拆解

1. 在 vllm/v1/worker/gpu_model_runner.py 的 _dummy_sampler_run 方法中，原有的 self.sampler 调用（使用空的 generators={}）之后，新增一个条件性的第二次采样调用。
2. 第二次调用通过 replace(dummy_metadata, generators={0: torch.Generator(device=self.device).manual_seed(0)}) 传入一个带有种子的 Generator，强制路由到 forward_native 路径（即 Triton 回退内核）。
3. 加入守卫条件 if self.sampler.logprobs_mode not in ("processed_logits", "processed_logprobs")：因为这些模式下 TopKTopPSampler 已经在构造时绑定了 forward = forward_native，额外调用只会冗余增加 profile_run 期间的峰值内存。
4. 不使用 .clone()——注释说明预热输出会被丢弃，任何 in-place 修改不影响正确性。
5. 所有变更仅涉及该文件的单个方法，无测试配置改动，但修复了 tests/v1/engine/test_async_llm.py::test_multi_abort 的确定性失败。

关键文件：

• vllm/v1/worker/gpu_model_runner.py（模块 模型运行器；类别 source；类型 core-logic；符号 _dummy_sampler_run）: 唯一修改的文件，在 _dummy_sampler_run 方法中为 forward_native 采样路径添加条件预热调用，消除首次 seed 请求的 JIT 延迟。

关键符号：_dummy_sampler_run

评论区精华

• gemini-code-assist[bot] 关于 logits in-place 修改的建议：建议克隆 logits 以防止首次 sampler 调用的 in-place 修改影响第二次调用。作者在代码注释中明确回应："No .clone() of logits: warmup output is discarded, so any in-place mutation by forward_native does not affect correctness." 最终保留原样。

• vadiklyutiy 关于 rejection sampler 的提问：询问是否需要为 rejection sampler 做同样预热。作者确认 rejection sampler 没有依赖于输入的分支，因此不需要。提问者接受此解释。

  • 无未解决的争议，PR 获得批准。

  • 是否需要克隆 logits 以防止 in-place 修改 (correctness): 作者通过代码注释说明无需克隆：预热输出被丢弃，in-place 修改不影响正确性。并且注释指出不使用 clone 是为了避免 profile_run 期间的 OOM（最后一次提交 removed .clone() 就是为了解决 processed_logprobs 模式的 OOM）。

  • rejection sampler 是否需要类似预热 (design): 作者确认 rejection sampler 没有依赖于输入的分支，因此不需要额外预热。提问者接受。

风险与影响

• 风险：风险较低。主要风险在于 logprobs_mode 的条件判断：如果未来新增类似模式，可能漏掉预热（但当前两个模式已覆盖所有不需要额外预热的情况）。另外，不使用 .clone() 在理论上存在 in-place 影响的风险，但正如注释所述，预热输出被丢弃且 forward_native 的行为已验证不会破坏后续调用。从测试结果看，修复通过且无 OOM（最后一次提交专门移除了 clone 以避免 OOM）。
• 影响：
  • 用户层面：首次带 seed 的请求不再经历约 1s 的延迟，采样行为无变化。
  • 系统层面：引擎初始化时间因额外预热稍有增加（约一次采样调用的耗时），但相对于 JIT 延迟可忽略。
  • 测试层面：修复了 test_multi_abort 的持续失败，提高了 CI 稳定性。
  • 影响范围：仅 V1 GPU 模型运行器的预热阶段，不影响模型计算或推理结果。
  • 风险标记：JIT 冷启动, 条件跳过可能遗漏新模式

关联脉络

• PR #40376 FlashInfer top-k/top-p sampler as default: 引入默认 FlashInfer 采样器，暴露了 forward_native 预热缺口，导致本 PR 需要修复的回归。
• PR #41316 Revert FlashInfer default sampler: 因本 PR 修复的相同测试失败而尝试 revert #40376，此 PR 提供了无需 revert 的修复方案。