执行摘要

• 一句话：默认启用 FlashInfer top-k/top-p 采样器
• 推荐动作：建议审核并合并此 PR。它在充分验证（安全测试、分布测试、性能 benchmark）的基础上默认启用了更快的采样器，且提供了完善的回退和降级机制。值得关注的设计决策在于：将默认值迁移到高性能实现，同时通过环境变量允许用户 opt-out，这是良好的兼容性策略。

功能与动机

FlashInfer top-k/top-p 采样器相较于 Triton 实现有显著的性能优势。过去因 NaN 输入导致的非法内存访问问题被默认禁用（PR #26859）。FlashInfer 已在 PR #2456 修复该问题，且作者无法复现此前报告的崩溃。因此有充分理由恢复默认启用，同时保留回退和测试保障。

实现拆解

1. 修改环境变量默认值（vllm/envs.py）：将 VLLM_USE_FLASHINFER_SAMPLER 的类型从 bool | None 改为 bool = True，在 envs 字典中的默认解析从 None 改为 True。

2. 调整采样器选择逻辑（vllm/v1/sample/ops/topk_topp_sampler.py）：在 TopKTopPSampler.__init__ 中重写条件分支：若硬件支持 FlashInfer（CUDA + 相应 compute capability）则默认使用 forward_cuda 路径；若硬件不支持但用户显式要求则报错；否则静默回退到 PyTorch-native 路径并记录 warning。同时移除 forward_cuda 中过时的 CPU-GPU 同步注释。

3. 补充测试（tests/v1/sample/test_topk_topp_sampler.py）：新增 TestFlashInferTopkToppRobustness 测试类，覆盖多种 NaN/Inf 污染 pattern；新增 TestFlashInferDistributionMatch 测试类，通过卡方检验验证分布一致性。

4. 调整端到端测试（tests/models/language/generation/test_hybrid.py）：将采样参数改为贪婪解码（temperature=0.0），避免因默认采样器改变导致非确定性输出。

5. 更新 CI 配置（.buildkite/test_areas/samplers.yaml）：显式添加 VLLM_USE_FLASHINFER_SAMPLER=0 和 =1 两个步骤，确保同时覆盖 PyTorch-native 和 FlashInfer 路径。

关键文件：

• vllm/v1/sample/ops/topk_topp_sampler.py（模块 采样器；类别 source；类型 core-logic；符号 TopKTopPSampler.init, TopKTopPSampler.forward_cuda）: 核心采样器选择逻辑，决定何时使用 FlashInfer 或 PyTorch-native 路径。
• tests/v1/sample/test_topk_topp_sampler.py（模块 采样测试；类别 test；类型 test-coverage；符号 _flashinfer_topk_topp_supported, TestFlashInferTopkToppRobustness, setup, _make_logits）: 新增 FlashInfer 采样器 NaN/Inf 鲁棒性测试和分布匹配测试，是验证本次变更正确性的核心测试。
• vllm/envs.py（模块 环境配置；类别 source；类型 core-logic；符号 VLLM_USE_FLASHINFER_SAMPLER）: 配置环境变量 VLLM_USE_FLASHINFER_SAMPLER 默认值，是启用 FlashInfer 的开关。
• .buildkite/test_areas/samplers.yaml（模块 CI配置；类别 config；类型 configuration）: 更新 CI 配置，确保两种采样路径都得到测试。
• tests/models/language/generation/test_hybrid.py（模块 模型测试；类别 test；类型 test-coverage）: 修复了一个因采样器默认变更可能失败的端到端测试，确保兼容性。

关键符号：_flashinfer_topk_topp_supported, TopKTopPSampler.init, TopKTopPSampler.forward_cuda, test_flashinfer_handles_pathological_logits, test_distribution_matches_theoretical

关键源码片段

vllm/v1/sample/ops/topk_topp_sampler.py

核心采样器选择逻辑，决定何时使用 FlashInfer 或 PyTorch-native 路径。

def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs") -> None:
    super().__init__()
    self.logprobs_mode = logprobs_mode
    ...
    # 根据硬件和环境变量选择采样路径
    if envs.VLLM_USE_FLASHINFER_SAMPLER:
        # 尝试导入 flashinfer 并检查 compute capability
        try:
            import flashinfer # noqa: F401
            from vllm.v1.attention.backends.flashinfer import FlashInferBackend
            capability = current_platform.get_device_capability()
            assert capability is not None
            if FlashInferBackend.supports_compute_capability(capability):
                logger.info_once("Using FlashInfer for top-p & top-k sampling.",
                                scope="global")
                self.forward = self.forward_cuda
            elif envs.is_set("VLLM_USE_FLASHINFER_SAMPLER"):
                # 用户显式要求但硬件不支持 → 报错
                raise RuntimeError(
                    "FlashInfer does not support compute capability "
                    f"{capability.as_version_str()}, unset VLLM_USE_FLASHINFER_SAMPLER=1.")
            else:
                # 默认启 + 硬件不支持 → 静默回退到 native
                logger.warning_once(
                    "FlashInfer top-p/top-k sampling not supported on "
                    "compute capability %s; falling back to PyTorch-native "
                    "sampler. Set VLLM_USE_FLASHINFER_SAMPLER=0 to silence.",
                    capability.as_version_str())
                self.forward = self.forward_native
        except ImportError:
            # flashinfer 未安装 → 走 native
            self.forward = self.forward_native
    else:
        # 用户显式设为 0 → 使用 native
        logger.info_once("FlashInfer top-p/top-k sampling disabled via "
                         "VLLM_USE_FLASHINFER_SAMPLER=0; using PyTorch-native sampler.")
        self.forward = self.forward_native

tests/v1/sample/test_topk_topp_sampler.py

新增 FlashInfer 采样器 NaN/Inf 鲁棒性测试和分布匹配测试，是验证本次变更正确性的核心测试。

# 模块级别判断 FlashInfer 是否可用
FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
​
@pytest.mark.skipif(
    not FLASHINFER_TOPK_TOPP_SUPPORTED,
    reason="FlashInfer top-k/top-p sampler requires CUDA "
           "and a GPU with FlashInfer support.",
)
class TestFlashInferTopkToppRobustness:
    """验证 FlashInfer 采样器在处理 NaN/Inf logits 时的鲁棒性。
    关键约束：不崩溃、不越界、不污染 batch 中其他正常行。"""
    BATCH = 8
    VOCAB = 32768
    TOPK = 50
    TOPP = 0.9
​
    @pytest.fixture(autouse=True)
    def setup(self):
        torch.set_default_device(DEVICE_TYPE)
        self.generator = Generator(device=DEVICE_TYPE).manual_seed(1234)
​
    def _make_logits(self, pattern: str) -> torch.Tensor:
        # 生成无污染 logits，然后根据 pattern 对第 0 行施加污染
        logits = torch.randn(self.BATCH, self.VOCAB,
                             generator=self.generator,
                             dtype=torch.float32) * 5.0
        if pattern == "clean":
            return logits
        elif pattern == "nan_one_row":
            logits[0, :] = float("nan")
        elif pattern == "nan_few":
            idx = torch.randperm(self.VOCAB, generator=self.generator)[:16]
            logits[0, idx] = float("nan")
        # ... 更多 pattern
        return logits
​
    def test_flashinfer_handles_pathological_logits(self):
        for pattern in ["clean", "nan_one_row", "nan_few", "nan_at_top", ...]:
            logits = self._make_logits(pattern)
            k = torch.tensor([self.TOPK] * self.BATCH, device=DEVICE_TYPE)
            p = torch.tensor([self.TOPP] * self.BATCH, device=DEVICE_TYPE)
            # 调用 FlashInfer 采样（通过 forward_cuda 间接）
            sampler = TopKTopPSampler()
            sampler.forward = sampler.forward_cuda # 强制使用 FlashInfer
            tokens, _ = sampler.forward(logits, k, p)
            # 断言：无 NaN token、所有 token 在 [0, vocab) 内、第 0 行之后的行 token 合理

评论区精华

• WoosukKwon 建议添加卡方检验确保输出分布不变。作者随后添加了 TestFlashInferDistributionMatch 卡方检验测试类。
• vadiklyutiy 询问为何使用独立脚本测试 NaN 鲁棒性而非单元测试。作者回应已在 test_topk_topp_sampler.py 中添加单元测试。

• vadiklyutiy 最终表示 LGTM 并通过审核。

• 将 NaN 鲁棒性检查加入单元测试 (testing): arpera 回复已在 tests/v1/sample/test_topk_topp_sampler.py 中作为单元测试添加。

• 添加分布匹配测试 (testing): arpera 添加了 TestFlashInferDistributionMatch 测试类。
• 端到端性能测试建议 (other): 作者已在 PR 描述中提供了基于小模型的 microbench 结果。

风险与影响

• 风险：

  1. 回退路径正确性：当 FlashInfer 不支持时，代码应可靠回退到 PyTorch-native 路径。当前实现中，如果硬件不支持且用户未显式开启则回退；若用户显式开启但硬件不支持则报错。需要确认回退路径中的 warning 日志不会导致冗余。
  2. 分布一致性：虽然通过了卡方检验，但不同模型输入分布可能仍存在微小差异，尤其涉及 NaN 边界情况可能触发旧 bug，需要持续监控。
  3. 非法内存访问：尽管作者无法复现，但 FlashInfer 的 NaN 修复可能不覆盖所有边缘情况。现有的鲁棒性测试覆盖了多种 NaN/Inf 模式，降低了回归风险。
  4. 性能影响：FlashInfer 不需要 CPU-GPU 同步，但新的分支逻辑引入的额外开销几乎可忽略。
  5. 测试覆盖：新测试依赖 FlashInfer 可用性，在非 CUDA 下会跳过，不影响 CI 通过率。
     - 影响：用户视角：所有部署在 CUDA 且 FlashInfer 兼容 GPU 上的 vLLM 用户将自动获得更快的采样，无需任何配置变更。若遇到采样异常，可通过 VLLM_USE_FLASHINFER_SAMPLER=0 回退到旧路径。
     系统视角：采样延迟显著降低，尤其是 batch size 较大时（如 batch 1024 时 top-k 延迟从 0.559ms 降至 0.224ms），可能减少整体推理延迟，提升吞吐量。
     团队视角：需要关注用户反馈，确认无新增的采样问题。CI 中添加了双路径测试，保障持续集成覆盖。

• 风险标记：回退路径风险, 分布一致性依赖, 非法内存访问历史

关联脉络

• PR #26859 Disable FlashInfer sampler by default: 之前默认禁用 FlashInfer 采样器的 PR，此 PR 将其反转并默认启用。