[Bugfix] Fix gridDim.y overflow for large row counts (#45255)
Signed-off-by: Jason Li <li.jason.cs@gmail.com>
JasonLi314 committed
Jun 19, 2026 at 20:27 UTC
93bad119120d0f9bff707dcbf5af5c029158b969
2 files changed
+71
-12
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+14
-12
@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
301
302
const int sf_k_local = local_group_id % kGroupsPerBlockX;
303
const int row_local = local_group_id / kGroupsPerBlockX;
304
- const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
305
- const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
304
+ // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
305
+ const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
306
+ const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;
307
308
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
309
asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
497
" is not a multiple of 4.");
498
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
499
const int ry = 16 / kx;
499
- const int64_t blocks_x = padded_groups_per_row / kx;
500
- const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
500
+ const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
501
+ const int64_t sf_k_blocks = padded_groups_per_row / kx;
502
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
502
- // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
503
- STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
504
- blocks_y <= static_cast<int64_t>(INT32_MAX),
503
+ // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
504
+ constexpr int64_t kMaxGridDimYZ = 65535;
505
+ STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
506
+ sf_k_blocks <= kMaxGridDimYZ,
507
"per_token_group_quant_8bit_packed grid too large: (",
506
- blocks_x, ", ", blocks_y, ").");
508
+ row_blocks, ", ", sf_k_blocks, ").");
509
510
auto dst_type = output_q.scalar_type();
511
@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
515
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
516
do { \
517
cudaLaunchConfig_t config = {}; \
516
- config.gridDim = dim3(static_cast<unsigned int>(blocks_x), \
517
- static_cast<unsigned int>(blocks_y)); \
518
+ config.gridDim = dim3(static_cast<unsigned int>(row_blocks), \
519
+ static_cast<unsigned int>(sf_k_blocks)); \
520
config.blockDim = dim3(num_threads); \
521
config.dynamicSmemBytes = 0; \
522
config.stream = stream; \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
541
#else
542
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
543
do { \
542
- dim3 grid(static_cast<unsigned int>(blocks_x), \
543
- static_cast<unsigned int>(blocks_y)); \
544
+ dim3 grid(static_cast<unsigned int>(row_blocks), \
545
+ static_cast<unsigned int>(sf_k_blocks)); \
546
dim3 block(num_threads); \
547
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
548
RY> \
tests/kernels/quantization/test_per_token_group_quant.py
+57
@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
345
)
346
347
348
+@pytest.mark.skipif(
349
+ not current_platform.is_cuda_alike(),
350
+ reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
351
+)
352
+def test_per_token_group_quant_fp8_packed_large_mn():
353
+ """Regression test for https://github.com/vllm-project/vllm/issues/45099.
354
+
355
+ Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
356
+ 2^16 - 1, respectively.
357
+ Prior code introduced a bug where it incorrectly assumed grid.x and y both have
358
+ 2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
359
+ launched with a large mn that exceeds grid.y limit (2^16 - 1).
360
+
361
+ This issue doesn't surface often because each forward pass only processes a
362
+ bounded token batch, not the full context.
363
+ Quantizing tensors with more rows than that will fail at launch with
364
+ "CUDA error: invalid argument".
365
+ This is a differential test that compares fp8 output against Triton output
366
+ reference when token size sits just above the gridDim.y 2^16 - 1 limit.
367
+ """
368
+
369
+ device = "cuda"
370
+ group_size = 128
371
+ # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
372
+ # row, so any mn > 65535 overflowed grid.y before the fix.
373
+ num_tokens, hidden_dim = 65537, 2048
374
+ torch.manual_seed(42)
375
+ x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
376
+
377
+ out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
378
+ x,
379
+ group_size=group_size,
380
+ use_ue8m0=True,
381
+ )
382
+
383
+ with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
384
+ ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
385
+ x, group_size, use_ue8m0=True
386
+ )
387
+
388
+ assert torch.equal(out_q, ref_q), "Quantized output mismatch"
389
+
390
+ # Vectorized packed-scale check; the per-element loop used by the smaller
391
+ # tests is too slow at this size. groups_per_row is a multiple of 4 here,
392
+ # so there is no K padding and the packed view lines up.
393
+ mn = num_tokens
394
+ groups_per_row = hidden_dim // group_size
395
+ k_num_packed = (groups_per_row + 3) // 4
396
+ assert groups_per_row % 4 == 0
397
+ ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
398
+ exp = ref_exponents.view(mn, k_num_packed, 4)
399
+ expected = (
400
+ exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
401
+ )
402
+ assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
403
+
404
+
405
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
406
@pytest.mark.parametrize("group_size", [64, 128])
407
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")