[Bugfix] Fix fused MoE int32 overflow in stride*offset without perf regression

[haosdent]

Purpose

Fixes #34413

PR #34279 annotated all stride parameters as tl.int64 to fix an int32 overflow crash, but this caused ~60x perf regression on small GPUs (e.g. NVIDIA GB10) due to register pressure. PR #34530 reverted that fix.

This patch prevents the overflow with minimal register impact by casting offs_token to int64 after loading instead of widening all strides. When chunking is disabled and M is large, stride_cm * offs_token (where stride_cm = N = w1.size(1) and offs_token up to M*topk) can exceed int32 max. The cast leverages Triton type promotion (int32 * int64 -> int64) following the existing pattern used for off_experts and offs_bn.

Adds a regression test that disables chunking with M=100000, n=2048, topk=6 (product = 4096 * 600000 = 2.46B > int32 max) and validates correctness against the torch_moe reference.

Overflow safety audit (all stride*offset products verified):

┌───────────────────────────────────────┬──────────────────┬─────────────┬────────────────────┬───────┐
│              Expression               │   Offset type    │ Stride type │       Result       │ Safe? │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ offs_token // top_k * stride_am       │ int64            │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ offs_k * stride_ak                    │ int32            │ int32       │ int32 (tiny: ~128) │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ off_experts * stride_be               │ int64 (existing) │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ offs_bn * stride_bn                   │ int64 (existing) │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ stride_cm * offs_token                │ int64            │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ BLOCK_SIZE_K * stride_ak (inner loop) │ constexpr (~128) │ int32       │ int32 (tiny)       │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ offs_token // top_k * stride_asm      │ int64            │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ off_experts * stride_bse              │ int64 (existing) │ int32       │ int64              │ Yes   │
├───────────────────────────────────────┼──────────────────┼─────────────┼────────────────────┼───────┤
│ offs_bn * stride_bsn                  │ int64 (existing) │ int32       │ int64              │ Yes   │
└───────────────────────────────────────┴──────────────────┴─────────────┴────────────────────┴───────┘

Test Plan

Add test case test_fused_moe_int64_overflow, and run python -m pytest tests/kernels/moe/test_moe.py::test_fused_moe_int64_overflow -v 2>&1

Test Result

New test case pass while it fail on the main branch.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request effectively addresses a significant performance regression in the fused MoE kernels on register-constrained GPUs. The previous approach of casting all stride parameters to tl.int64 was too broad and caused excessive register pressure. The new approach is more targeted: it reverts the strides to their default integer type and instead casts offs_token to tl.int64 just before its use in pointer arithmetic. This correctly prevents overflow in stride-offset products by leveraging Triton's type promotion, while minimizing the impact on register usage. The change is applied consistently across both fused_moe_kernel and fused_moe_kernel_gptq_awq and aligns with existing patterns in the code. The fix is well-reasoned and should restore the expected performance. I approve these changes.

[eugr]

@haosdent - thanks, this fixes the performance regression on Spark.
@johnnynunez - fyi.

[mgoin]

Thanks a lot @haosdent for the analysis, we will revert the original PR for now to simplify #34530

[tlrmchlsmth]

Sorry, went a bit scorched earth with that one - have you tried the repro scrip in #34279 on this by any chance? Hoping this more selective fix still resolves the IMAs

[haosdent]

I run the script as well, this fix works for me

Testing m=1000, n=1024, k=1024, e=8, topk=2
  stride_cm=2048, max_offs_token=2000, max_c_offset=4096000, int32_max=2147483647
  stride_am=1024, max_a_offset=1024000
  SUCCESS - output shape: torch.Size([1000, 1024])

Testing m=40000, n=1024, k=1024, e=8, topk=6
  stride_cm=6144, max_offs_token=240000, max_c_offset=1474560000, int32_max=2147483647
  stride_am=1024, max_a_offset=40960000
  SUCCESS - output shape: torch.Size([40000, 1024])

Testing m=100000, n=2048, k=1024, e=8, topk=6
  stride_cm=12288, max_offs_token=600000, max_c_offset=7372800000, int32_max=2147483647
  stride_am=1024, max_a_offset=102400000
  ** C OFFSET WILL OVERFLOW INT32 **
  SUCCESS - output shape: torch.Size([100000, 1024])

but if switched to main, it would crash

Testing m=1000, n=1024, k=1024, e=8, topk=2
  stride_cm=2048, max_offs_token=2000, max_c_offset=4096000, int32_max=2147483647
  stride_am=1024, max_a_offset=1024000
  SUCCESS - output shape: torch.Size([1000, 1024])

Testing m=40000, n=1024, k=1024, e=8, topk=6
  stride_cm=6144, max_offs_token=240000, max_c_offset=1474560000, int32_max=2147483647
  stride_am=1024, max_a_offset=40960000
  SUCCESS - output shape: torch.Size([40000, 1024])

Testing m=100000, n=2048, k=1024, e=8, topk=6
  stride_cm=12288, max_offs_token=600000, max_c_offset=7372800000, int32_max=2147483647
  stride_am=1024, max_a_offset=102400000
  ** C OFFSET WILL OVERFLOW INT32 **
  SUCCESS - output shape: torch.Size([100000, 1024])
  CRASH: AcceleratorError: CUDA error: an illegal memory access was encountered
Search for `cudaErrorIllegalAddress' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

[haosdent]

@SouthWest7 thank you for this

The test_moe test cases do not cover the overflow scenario.

I have refered the reproduce script and add new test case test_fused_moe_int64_overflow to cover it.

[haosdent]

Hi @mgoin I have updated the PR, may you help to take a look to check if this is still necessary?

[mgoin]

Great work, glad you setup the regression test as well. LGTM but will let @tlrmchlsmth sign off. It would be good to get AMD perf validation as well

[tlrmchlsmth]

Thank you! This has been annoying me for a long time so glad to see it get fixed

[mgehre-amd]

Unfortunately, I still see the perf degradation:

Benchmark: cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit on Strix Halo (gfx1151)

Just d7982da (int64 strides) by itself gave a 5% regression.
Adding 54ba644 on top does not recover the original performance.

Setup: --input-len 512 --max-model-len 4096 --dtype float16 --target-gpu-memory-gb 20 --num-prompts 5

Metric	Baseline (neither d7982da nor 54ba644)	Modified (both d7982da and 54ba644)	Delta
Decode	46.5 tok/s	43.7 tok/s	-6.0%
TPOT	21.48 ms	22.89 ms	+6.5%
TTFT	275 ms	434 ms	+58%

Each configuration was run twice with consistent results.

[haosdent]

and the targeted offs_token cast in 54ba644 does not compensate for it on this workload.

@mgehre-amd Thanks for your test, it has indeed degraded a bit after becoming int64.
But I'm not sure if compromise on this to fix the overflow issue or revert the change. I could not judge it from my side (either of them is ok for me until we find a better approach), so we may need vLLM experts' help to decide this @mgoin @tlrmchlsmth .

[mgehre-amd]

If you revert both d7982da and 54ba644 , do the degradation happen? @mgehre-amd

With both commits reverted, I don't see the regression. Updated my previous message to be clearer.