Fix int32 overflow in Triton grouped-GEMM expert offset (#8261)
## Summary
`group_gemm_triton.py` computed the expert base pointer as `selected *
stride_be` in **int32**. Since `stride_be` is `K*N` **elements**, any
expert index past `2**31 / (K*N)` wraps to a negative offset, so the
kernel reads out of bounds and faults with an illegal memory access.
```python
b_base = b_ptr + selected * stride_be # int32 product -> wraps
```
A 64-expert layer with `K=4096, N=14336` places the last expert at
**3.70e9 elements**, past the int32 limit of **2.15e9**:
```
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered
```
The fix casts the expert index to int64 before the multiply. The failure
is a hard fault rather than silent corruption, so no previously-produced
numerics are suspect.
## Reachability
This is not a synthetic configuration. `prefer_triton_grouped_mm()`
selects the Triton path on **all sm<90 devices**, and on **any device
when `torch._grouped_mm` is unavailable**
(`accelerator/cuda_accelerator.py:281`). The overflow needs only
`(E_local - 1) * K * N >= 2**31` — reached by fine-grained-MoE
configurations at Mixtral-scale FFN dimensions when enough experts land
on one rank (low EP degree or high expert count).
## Root-cause evidence
Bisection matches the `2**31` threshold exactly — the trigger is the
**product**, not the expert count alone:
| K | N | experts | `(E-1)*K*N` | vs int32 max | result |
|---|---|---|---|---|---|
| 4096 | 14336 | 64 | 3.70e9 | over | **FAIL** |
| 4096 | 14336 | 64 (32 rows/expert) | 3.70e9 | over | **FAIL** |
| 4096 | 14336 | 32 | 1.82e9 | under | OK |
| 1024 | 14336 | 64 | 9.24e8 | under | OK |
| 4096 | 4096 | 64 | 1.06e9 | under | OK |
| 4096 | 14336 | 8 | 4.11e8 | under | OK |
Rows-per-expert is irrelevant (row 2), which rules out an M-dimension
indexing issue and isolates the fault to the expert-stride term.
## Testing
**Regression test** added to the existing
`tests/unit/v1/moe/test_group_gemm_triton.py` (no new file). It
fails-before / passes-after:
```
# before the fix
FAILED test_expert_offset_exceeds_int32 - RuntimeError: CUDA error: an illegal memory access was encountered
# after the fix
1 passed
```
The test asserts `(num_experts - 1) * stride_be > 2**31 - 1` up front,
so it fails loudly rather than silently stopping exercising the overflow
if the sizes are ever tuned.
**Numerics** at the previously-faulting shape (`E=64, K=4096, N=14336`),
checked per-expert against a `torch` reference:
```
worst relative error across all 64 experts = 3.31e-03 (bf16, K=4096)
```
**Suite results** on `torch 2.8.0+cu128`, triton 3.4.0, H100 (sm90):
```
71 passed, 0 failed (of 75 collected)
```
No failures, and no regressions attributable to this change.
## Verification caveat
Verified on **H100 (sm90)**. The overflow is arch-independent integer
arithmetic, so the diagnosis and fix carry over — but note that on a
current PyTorch, sm90 does **not** select this kernel in production
(`prefer_triton_grouped_mm()` returns `False` there, confirmed on `torch
2.8.0`). The regression test calls `group_gemm_triton` directly, so it
exercises the overflow regardless. An Ampere confirmation would still be
a welcome extra check.
The regression test allocates ~4.5 GiB of expert weights. That is
inherent — the bug cannot manifest below `2**31` elements — and it skips
automatically when free device memory is short.
## Incidental observation (not addressed here)
While benchmarking this kernel on sm90, the Triton grouped GEMM measured
**slower than a per-expert `mm` loop** (8192 tokens, H=4096, FFN=14336,
torch 2.4.1 where sm90 still selected the Triton path):
| experts | for-loop | Triton grouped | ratio |
|---|---|---|---|
| 2 | 1.46 ms | 2.26 ms | 0.65x |
| 8 | 1.87 ms | 2.21 ms | 0.85x |
| 32 | 2.11 ms | 2.74 ms | 0.77x |
This is consistent with the module's documented sm80/sm86 scope and is
not a defect — on sm90 with a current PyTorch,
`prefer_triton_grouped_mm()` correctly returns `False` and the native
`torch._grouped_mm` is used instead. Flagging it only because on a torch
build without `torch._grouped_mm`, that helper returns `True` regardless
of architecture, which routes sm90 onto the slower path. Worth a
separate look if that combination is considered supported.
Signed-off-by: Zhipeng Wang <zhipengbayern@gmail.com>
DeepSpeed
c5331bc1 - Fix int32 overflow in Triton grouped-GEMM expert offset (#8261)
