执行摘要

本PR修复了Triton W4A16 GEMM内核在BLOCK_K大于量化组大小时，因单个计算瓦片跨越多个scale组却只使用第一个组的scale，导致尾部行数据静默损坏的问题。修复方法是在内核启动器中强制将BLOCK_K限制为不超过group_size。该问题在ROCm平台上使用特定量化模型进行长上下文工具调用时表现为模型行为异常，修复后模型运行正确且效率提升。

功能与动机

问题背景：作者在使用ROCm RDNA3平台运行Qwen3.5-35B-A3B-GPTQ-W4A16-G32模型时，发现模型在超过10K令牌的长上下文工具调用中表现异常：重复调用相同参数的工具、未完成任务、或幻觉指令。

根本原因：Triton W4A16 GEMM内核（由PR #37352引入）在triton_w4a16_gemm函数中，当BLOCK_K（默认32）大于量化group_size（如32）时，单个计算瓦片会跨越多个量化scale组，但内核只加载第一个组的scale应用于整个瓦片，导致尾部行使用错误的scale进行反量化，静默损坏权重数据。

引用PR body关键表述：

"When BLOCK_K exceeds group_size, a single tile spans multiple scale groups, but only the first group's scales are applied to all rows in the tile. This silently corrupts the dequantized weights in the tail rows."

实现拆解

仅修改一个文件：vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py。

关键改动：在triton_w4a16_gemm函数中，在设置BLOCK_M、BLOCK_N、BLOCK_K默认值后，添加条件检查：

if group_size < BLOCK_K:
    BLOCK_K = group_size

逻辑说明：

• 默认BLOCK_K = 32，但量化组大小group_size可能更小（如32）。
• 当group_size < BLOCK_K时，将BLOCK_K限制为group_size，确保每个计算瓦片不超过一个量化组，避免scale错配。
• 这保证了内核加载的scale与瓦片内所有行对应同一个量化组。

评论区精华

Review讨论较少，但有两个关键评论：

1. gemini-code-assist[bot]：

   "This pull request introduces a safety check in the triton_w4a16_gemm function to clamp the BLOCK_K parameter to the group_size. This change prevents potential data corruption that could occur if a processing tile spans multiple quantization groups, ensuring that the correct scales and zeros are applied during dequantization."

2. yewentao256：

   "LGTM, thanks for the work!"

没有争议点，修复被迅速批准。

风险与影响

技术风险：

1. 静默数据损坏风险已修复：原问题导致权重反量化错误，输出不可预测，修复后确保正确性。
2. 性能潜在影响：当group_size较小时（如32），BLOCK_K被限制为较小值，可能降低计算效率，但这是正确性必需的权衡。
3. 回归风险低：仅影响使用triton_w4a16_gemm且BLOCK_K > group_size的场景，其他场景不受影响。
4. 缺少测试覆盖：PR未添加测试用例，但基于问题描述，修复已通过实际模型（Qwen3.5-35B-A3B-GPTQ-W4A16-G32）验证。

影响评估：

• 用户影响：使用Triton W4A16量化内核（特别是ROCm平台）的用户在长上下文推理中将获得正确结果，避免模型行为异常。
• 系统影响：确保量化权重反量化正确性，提升模型输出质量和可靠性。
• 团队影响：揭示了内核实现中一个隐蔽的正确性问题，提醒在量化内核设计中需考虑BLOCK_K与group_size的匹配关系。

关联脉络

与历史PR的关联：

• PR #37352：引入了Triton W4A16 GEMM内核，本PR修复了该内核中的一个bug。PR body中明确提及："with the kernel introduced in PR #37352"。

功能演进方向：

• 近期多个PR涉及量化内核的bugfix和优化（如PR #39717、#39604、#39418、#38707），显示团队在持续完善量化支持，特别是在多平台（ROCm、XPU）上的正确性和性能。
• 本PR是这一趋势的一部分，专注于ROCm平台上量化内核的正确性修复，确保长上下文推理的可靠性。