执行摘要

• 一句话：修复 gumbel_sample 缓存列 stride 错误，避免错位写入
• 推荐动作：此 PR 值得精读，它展示了一个典型的 stride 计算 bug 的修复过程，并且测试用例设计良好，覆盖了跨步写入隔离和边界拒绝场景。对于理解 Triton kernel 中张量步长处理很有参考价值。建议关注后续是否有更多针对 gumbel_sample 的优化。

功能与动机

PR 描述明确指出：“Currently, we are striding along columns in the draft logits cache using vocab_size instead of the actual stride value for the tensor. This is dangerous because in cases where vocab_size != draft_logits.size(-1), we can end up writing to the wrong slot.” 当 draft 模型（如自回归模型）在输出层添加额外的 mask 或噪声列时，缓存宽度会大于实际 logits 宽度，原代码按 vocab_size 步进会导致后续步骤写入错位，产生静默错误。

实现拆解

本 PR 的核心修复分三步：

1. 修改 gumbel_block_argmax 函数签名与实现（vllm/v1/worker/gpu/sample/gumbel.py）：将单一的 logits_cache_stride 参数替换为 logits_cache_stride_0 和 logits_cache_stride_1 两个参数，分别表示 cache 张量第 0 维和第 1 维的步长。原代码中 col * vocab_size 被替换为 col * logits_cache_stride_1，从而使用 cache 实际列宽进行偏移。
2. 同步更新 _gumbel_sample_kernel 与 gumbel_sample 函数（同文件）：_gumbel_sample_kernel 传递两个步长参数；gumbel_sample 调用 kernel 时，传入 logits_cache.stride(0) 和 logits_cache.stride(1)，并在入口处增加断言，确保 cache 的最后一个维度不小于 vocab_size，否则抛出异常提示会被截断。
3. 更新调用方（vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py）：_resample_kernel 中原本传 0 作为单一步长，现改为两个 0，以匹配新签名。
4. 补充测试（tests/v1/worker/test_gpu_gumbel_sample.py）：新增 test_logits_cache_columns_stay_separate_across_steps 测试多步写入时各步列保持独立；新增 test_logits_cache_narrower_than_logits_is_rejected 测试 cache 过窄时正确抛错。

关键文件：

• vllm/v1/worker/gpu/sample/gumbel.py（模块 采样内核；类别 source；类型 core-logic；符号 gumbel_block_argmax, _gumbel_sample_kernel, gumbel_sample）: 核心修复文件，修改 stride 计算逻辑，新增断言，是本次 PR 的关键改动。
• tests/v1/worker/test_gpu_gumbel_sample.py（模块 采样测试；类别 test；类型 test-coverage；符号 test_logits_cache_columns_stay_separate_across_steps, test_logits_cache_narrower_than_logits_is_rejected）: 新增两个测试，覆盖跨步隔离和过窄拒绝，验证修复正确性。
• vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py（模块 拒绝采样；类别 source；类型 core-logic；符号 _resample_kernel）: 同步更新调用点，传递两个步长参数，保持接口一致。

关键符号：gumbel_block_argmax, _gumbel_sample_kernel, gumbel_sample, _resample_kernel

关键源码片段

tests/v1/worker/test_gpu_gumbel_sample.py

新增两个测试，覆盖跨步隔离和过窄拒绝，验证修复正确性。

# tests/v1/worker/test_gpu_gumbel_sample.py
# 测试：跨步写入时各步列保持独立，防止错位
@pytest.mark.parametrize("extra_cache_cols", [0, 1])
def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int):
    """Each drafting step must land in its own cache column."""
    # extra_cache_cols=1 模拟 draft 模型添加额外列的场景
    torch.manual_seed(0)
    num_reqs, vocab_size, num_steps = 4, 1031, 3
    ...
    cache = torch.zeros(num_reqs, num_steps, vocab_size + extra_cache_cols, device=DEVICE)
    cols = torch.arange(num_steps, dtype=torch.int32, device=DEVICE)
    per_step = [torch.randn(num_reqs, vocab_size, device=DEVICE) for _ in range(num_steps)]
    # 依次模拟多个步的写入
    for step, logits in enumerate(per_step):
        gumbel_sample(logits, idx_mapping, temp, seed, pos, apply_temperature=True,
                      logits_cache=cache, logits_cache_col=cols[step])
    # 验证每步写入未互相覆盖
    for step, logits in enumerate(per_step):
        stored = cache[:, step, :vocab_size]
        assert torch.equal(_float_bits(stored), _float_bits(logits)), ...
    # 额外列保持未写入
    assert not cache[:, :, vocab_size:].any()
​
# 测试：缓存过窄时直接拒绝，避免静默截断
def test_logits_cache_narrower_than_logits_is_rejected():
    # 构造宽度小于 vocab_size 的缓存
    cache = torch.zeros(num_reqs, num_steps, vocab_size - 1, device=DEVICE)
    with pytest.raises(AssertionError, match="narrower"):
        gumbel_sample(logits, idx_mapping, temp, seed, pos, apply_temperature=True,
                      logits_cache=cache, ...)

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

同步更新调用点，传递两个步长参数，保持接口一致。

# vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
# 更新调用：将原来的单一 stride 替换为两个 0
value, idx = gumbel_block_argmax(
    residual_logits,
    block,
    mask,
    resample_token_idx,
    expanded_idx_mapping_ptr,
    temp_ptr,
    seed_ptr,
    pos_ptr,
    None, # logits_cache_ptr
    0, # logits_cache_stride_0
    0, # logits_cache_stride_1
    None, # logits_cache_col_ptr
    vocab_size,
    ...
)

评论区精华

Review 评论区有两条评论：

• 作者 TheEpicDolphin 触发 CI 运行，github-actions 自动回复已触发 Buildkite CI #84704。

• 维护者 WoosukKwon 直接批准（APPROVED），没有提出问题。
  没有其他实质性讨论，该修复方案清晰，风险低。

• 修复讨论 (design): 修复被维护者 WoosukKwon 批准，方案合理。

风险与影响

• 风险：本 PR 修改了核心采样内核 gumbel_block_argmax，但改动仅涉及 stride 计算，逻辑简单，风险较低。主要风险点：
  1. 调用兼容性：由于函数签名变更，其他调用 gumbel_block_argmax 的位置（如 rejection_sampler_utils.py 的 _resample_kernel）已同步更新，但需检查是否有遗漏。
  2. 性能影响：新增断言在 gumbel_sample 入口处，每次调用都会检查 cache 宽度，但断言开销极小，不影响性能。
  3. 回归风险：如果 cache 形状从未出现宽度大于 vocab_size 的情况，旧代码正确；修复后行为不变，但新增断言可能暴露一些原本未考虑到的情况，需要关注。
     - 影响：影响范围集中在 Model Runner V2 的 spec decode 路径，尤其是 gumbel_sample 用于 draft logits 缓存写入时。修复后，当 draft 模型输出的 logits 宽度小于缓存宽度时，各步的写入位置正确，避免了数据错位导致的采样错误。对于用户而言，修复了潜在的静默错误，提升了解码结果的可靠性。对系统无性能影响，对团队而言，此修复为 MRV2 的稳定性贡献了一次重要修正。
     - 风险标记：核心采样路径变更, 函数签名变更需同步所有调用方, 缺少对非 Triton 后端的兼容性验证

关联脉络

• PR #52987 Revert "[Kernel] Gemma-4 FA4 FP8 Kernel": 涉及 spec decode 路径，相关 gumbel_sample 的调用，可能相互影响。