A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
[Bugfix][Kernel][ROCm] Fix triton_w4a16 scales mismatch when BLOCK_K > group_size (#39705)
Signed-off-by: JartX <sagformas@epdcenter.es>
JartX committed
Apr 13, 2026 at 20:29 UTC
f414f90601b9ef766017f1fb9398c1bed9bc871a
1 file changed
+8
vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py
+8
@@ -235,6 +235,14 @@ def triton_w4a16_gemm(
235
else:
236
BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
237
238
+ # The kernel loads scales/zeros for a single group per BLOCK_K tile
239
+ # (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail
240
+ # of the tile dequantize with the wrong group's scales, silently
241
+ # corrupting the output. Clamp BLOCK_K to group_size to keep one
242
+ # scale group per tile.
243
+ if group_size < BLOCK_K:
244
+ BLOCK_K = group_size
245
+
246
grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
247
248
triton_w4a16_gemm_kernel[grid](