Prhub

#44042 [CI] Reject out-of-vocabulary before they reach the GPU logprob path

原始 PR 作者 AndreasKaratzas 合并时间 2026-06-03 11:27 文件变更 4 提交数 8 评论 7 代码增减 +82 / -9

执行摘要

提前拒绝越界 token ID,稳定 ROCm CI

ROCm CI 中由 Schemathesis 生成的包含越界 logprob_token_ids(如负数、超大数)的请求,会触发 ROCm/HSA 硬件异常并杀死引擎,导致后续测试级联失败。需要在请求验证阶段提前拒绝这些无效参数。

值得精读。该 PR 展示了早期验证如何防御 GPU 异常,以及如何针对平台差异做最小侵入性 workaround。其中验证插入位置和 platform check 的使用方式可为类似问题提供参考。

讨论亮点
  1. 验证时机讨论:hclsys 担心 verify() 可能在模型加载前被调用,导致 get_vocab_size 不可用。作者回应 verify 依赖 ModelConfig,且已用于 logprobs=-1,不会引入新假设。
  2. PR 范围讨论:njhill 建议将 ROCm 特定修复(thinking budget、shutdown)拆分为独立 PR。作者解释这些修复是使 CI 通过的必要条件,部分问题因重构才暴露。
  3. 技术说明:作者在评论中解释了 thinking budget 修复的动机——避免 ROCm 上 2D 高级索引写入的潜在问题。

实现拆解

  1. 在 SamplingParams._validate_logprobs 中新增词汇范围校验:获取 vocab_size 后,对 logprob_token_ids 中每个 token_id 检查是否 <0 或 >= vocab_size,若是则抛出 VLLMValidationError。
  2. 在 test_logprobs.py 中添加单元测试,利用 SimpleNamespace 模拟 ModelConfig,验证有效和无效 token_id 的通过/拒绝行为。
  3. 在 gpu_model_runner.py 的 shutdown 方法中,为 ROCm 增加 CUDA 图显式清理(CUDAGraphWrapper.clear_all_graphs)及额外的 gc.collect/empty_cache/synchronize,防止下次启动时 HSA 故障。
  4. 在 thinking_budget_state.py 的 apply_forcing_to_logits 中,将 ROCm 平台上的 2D index_put 替换为 1D index_fill_ 或逐行赋值,以避免 ROCm 在高级索引写入时的内存访问错误。
文件 模块 状态 重要度
vllm/sampling_params.py 采样参数 modified 6.47
vllm/v1/sample/thinking_budget_state.py 思考预算 modified 6.5
tests/v1/sample/test_logprobs.py 测试 modified 5.89
vllm/v1/worker/gpu_model_runner.py 模型运行器 modified 5.72

关键符号

_validate_logprobs shutdown _apply_forcing_to_logits _model_config test_logprob_token_ids_validate_vocab_bounds_valid test_logprob_token_ids_validate_vocab_bounds_invalid

关键源码片段

vllm/v1/sample/thinking_budget_state.py core-logic

修复 ROCm 上 thinking budget 写入路径的内存访问错误,是 CI 稳定的关键一环。

if active_indices_cpu:
    device = logits.device
    if current_platform.is_rocm() and logits.is_contiguous():
        # Flattened 1D index_fill 避免 ROCm 上 2D 高级索引写错误
        vocab_size = logits.shape[1]
        flat_indices_cpu = [
            row * vocab_size + token
            for row, token in zip(active_indices_cpu, force_tokens_cpu)
        ]
        flat_indices = async_tensor_h2d(
            flat_indices_cpu, dtype=torch.long, device=device
        )
        logits.view(-1).index_fill_(0, flat_indices, 1e9)
    elif current_platform.is_rocm():
        # 非连续张量退化为逐行赋值
        fill = logits.new_tensor(1e9)
        for row, token in zip(active_indices_cpu, force_tokens_cpu):
            logits[row, token] = fill
    else:
        # 非 ROCm 平台保持原 index_put_ 路径
        active_indices = async_tensor_h2d(
            active_indices_cpu, dtype=torch.long, device=device
        )
        force_tokens = async_tensor_h2d(
            force_tokens_cpu, dtype=torch.long, device=device
        )
        fill = logits.new_full((len(active_indices_cpu),), 1e9)
        logits.index_put_((active_indices, force_tokens), fill)

评论区精华

验证时机可行性 正确性

hclsys 担心 verify() 可能在模型加载前被调用,导致 get_vocab_size 不可用。作者回应 verify 依赖 ModelConfig,且已用于 logprobs=-1,不会引入新假设。

结论:确认验证在正确时机执行,无需额外守卫。 · 已解决

是否将 ROCm 修复拆分为独立 PR 设计

njhill 建议将非 OOV 的 ROCm 修复拆分离。作者解释这些修复是为了让 CI 通过,部分问题因重构才暴露,需要一起合并。

结论:保留在同一 PR 中,由作者统一管理。 · 已解决

思考预算写入路径的 ROCm 工作原因 性能

作者描述使用 1D index_fill_ 替代 2D index_put_ 以避免 ROCm 上 advanced-indexing 错误。

结论:采用 flattened index_fill_ 作为默认 ROCm 路径,fallback 到循环赋值。 · 已解决

风险与影响

核心风险:logprob_token_ids 验证逻辑新增了 vocab_size 依赖,但已有前置使用(logprobs=-1),应为安全;ROCm 特定修改(shutdown、thinking budget)通过 platform check 隔离,不影响 CUDA;但 shutdown 路径的额外清理可能改变当前进程内的资源释放时序,需关注后续启动稳定性。

影响范围:仅涉及 logprob_token_ids 参数的请求验证和 ROCm 平台下的模型关闭路径。对非 ROCm 用户无行为变化;对 ROCm 用户,无效请求将提前报错而非崩溃,引擎重启更可靠。团队受益于更稳定的 CI。

平台依赖变更 核心验证路径 稳定性修复

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论