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 8, 05:00 UTC
[Bugfix] Fix fused MoE int32 overflow in stride*offset without perf regression (#34507)
Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
haosdent committed
Feb 17, 2026 at 09:58 UTC
b68fd899d1317259d61d8bbaf79dcb2749d17634
2 files changed
+54
-1
tests/kernels/moe/test_moe.py
+49
@@ -396,6 +396,55 @@ def test_fused_moe(
396
)
397
398
399
+def test_fused_moe_int64_overflow(monkeypatch, workspace_init):
400
+ """Regression test for int32 overflow in stride*offset products.
401
+
402
+ When chunking is disabled and M is large, stride_cm * offs_token can
403
+ exceed int32 max. Verifies the offs_token int64 cast (fix for #34413)
404
+ prevents overflow and produces correct results.
405
+
406
+ Reproduces the scenario from PR #34279.
407
+ """
408
+ # ~12 GB GPU memory needed for intermediate caches
409
+ free_mem = torch.cuda.mem_get_info()[0]
410
+ if free_mem < 12 * 1024**3:
411
+ pytest.skip("Insufficient GPU memory for overflow test")
412
+
413
+ set_random_seed(7)
414
+
415
+ m, n, k, e, topk = 100000, 2048, 1024, 8, 6
416
+ dtype = torch.bfloat16
417
+
418
+ # Disable chunking to expose the overflow-prone code path
419
+ monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "10000000")
420
+
421
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
422
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
423
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
424
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
425
+
426
+ # Verify the test exercises the overflow condition:
427
+ # C has shape (M, topk, N) where N = w1.size(1) = 2*n
428
+ # stride_cm = C.stride(1) = N, max offs_token = M * topk
429
+ # Product must exceed int32 max for this test to be meaningful
430
+ N = w1.size(1)
431
+ assert N * m * topk > 2**31 - 1, "Test params don't trigger int32 overflow"
432
+
433
+ fused_moe_fn = functools.partial(fused_moe, renormalize=False)
434
+
435
+ with set_current_vllm_config(vllm_config):
436
+ run_moe_test(
437
+ torch_moe,
438
+ fused_moe_fn,
439
+ a=a,
440
+ w1=w1,
441
+ w2=w2,
442
+ score=score,
443
+ topk=topk,
444
+ global_num_experts=e,
445
+ )
446
+
447
+
448
@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS_SMALL_M)
449
@pytest.mark.parametrize("e", NUM_EXPERTS_LARGE)
450
@pytest.mark.parametrize("topk", TOP_KS_SMALL)
vllm/model_executor/layers/fused_moe/fused_moe.py
+5
-1
@@ -175,7 +175,8 @@ def fused_moe_kernel_gptq_awq(
175
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
176
return
177
offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
178
- offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
178
+ # Cast to int64 to prevent overflow in stride*offset products
179
+ offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64)
180
token_mask = offs_token < num_valid_tokens
181
182
off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
@@ -426,6 +427,9 @@ def fused_moe_kernel(
427
pid_m, # first element = pid_m
428
num_valid_tokens, # remaining elements = constant
429
)
430
+ # Cast to int64 to prevent overflow in stride*offset products
431
+ # (e.g. stride_cm * offs_token can exceed int32 for large token counts)
432
+ offs_token = offs_token.to(tl.int64)
433
434
token_mask = offs_token < num_valid_tokens
435