[Bugfix] Fix int32 token offset overflow in fused SiLU block quant (#53409)
Signed-off-by: Anni <dranox2222@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
canlahlah committed
Aug 28, 2026 at 13:23 UTC
06cccf8730a199817be58e2de576918fe7a46e1e
2 files changed
+26
-1
csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu
+1
-1
@@ -26,7 +26,7 @@ __global__ void silu_and_mul_per_block_quant_kernel(
26
"group_size must be a power of 2 for correct reduction");
27
28
// Grid: (num_tokens, num_groups)
29
- int const token_idx = blockIdx.x;
29
+ int64_t const token_idx = blockIdx.x;
30
int const group_idx = blockIdx.y;
31
int const tid = threadIdx.x; // tid in [0, group_size)
32
int const num_tokens = gridDim.x;
tests/kernels/core/test_fused_silu_mul_block_quant.py
+25
@@ -7,6 +7,7 @@ import torch.nn.functional as F
7
8
import vllm._custom_ops as ops
9
from tests.kernels.utils import opcheck
10
+from tests.utils import large_gpu_test
11
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
12
per_token_group_quant_fp8,
13
)
@@ -187,3 +188,27 @@ def test_silu_block_quant_edge_cases(
188
assert not torch.isnan(out.float()).any()
189
assert not torch.isnan(scales).any()
190
assert not torch.isinf(scales).any()
191
+
192
+
193
+@large_gpu_test(min_gb=24)
194
+@torch.inference_mode()
195
+def test_silu_block_quant_int32_offset_overflow(default_vllm_config):
196
+ """Token offsets beyond INT32_MAX must not wrap (issue #53390)."""
197
+ hidden_size = 16384
198
+ group_size = 128
199
+ quant_dtype = current_platform.fp8_dtype()
200
+ # Smallest batch whose last token's input offset (token_idx * 2 *
201
+ # hidden_size) exceeds INT32_MAX.
202
+ num_tokens = (2**31 - 1) // (2 * hidden_size) + 2
203
+
204
+ x = torch.randn(num_tokens, hidden_size * 2, dtype=torch.bfloat16, device="cuda")
205
+ out, scales = ops.silu_and_mul_per_block_quant(
206
+ x, group_size, quant_dtype, None, False
207
+ )
208
+ torch.accelerator.synchronize()
209
+
210
+ expected_out, expected_scales = ops.silu_and_mul_per_block_quant(
211
+ x[-1:].contiguous(), group_size, quant_dtype, None, False
212
+ )
213
+ assert torch.equal(scales[-1:], expected_scales)
214
+ assert torch.equal(out[-1:], expected_out)