linkedin / linkedin/Liger-Kernel
Fused MoE kernels crash on Qwen3.5-397B (int32 offset overflow)
- Dominant language
- Python
- Stars
- 6.6k
- Forks
- 603
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 47
Description
### 🐛 Describe the bug
The fused MoE Triton kernels in `liger_kernel/ops/fused_moe_kernels.py` compute the per-expert weight offset as `expert_idx * stride_w_E` in **int32**. When a model's expert weight tensor exceeds 2^31 elements, this offset overflows to a negative value, producing an out-of-bounds pointer and a CUDA `illegal memory access`.
This only manifests on large MoE models. Smaller ones stay under the limit and work fine, which makes the bug easy to miss.
- **liger-kernel version:** 0.8.0
- **Model that triggers it:** Qwen3.5-397B-A17B (`num_experts=512`, `hidden_size=4096`, `moe_intermediate_size=1024`)
- **Setup:** LoRA SFT, 2 nodes × 8×H200, DeepSpeed ZeRO-3, gradient checkpointing, `use_liger_kernel=True`
Forward pass dies inside `_fused_up_proj_swiglu_kernel`:
```
File ".../liger_kernel/ops/fused_moe.py", line 203, in forward
_fused_up_proj_swiglu_kernel[...](...)
...
torch.AcceleratorError: CUDA error: an illegal memory access was encountered
```
Same root cause also fires in the backward `_moe_bwd_dW2_kernel` / `_moe_bwd_dW1_kernel`.
In each kernel the expert offset is built like this (forward up-proj shown):
```python
gate_up_proj_ptr + expert_idx * stride_w_E + n_idx[:, None] * stride_w_N + k_idx[None, :] * stride_w_K
```
`expert_idx` is loaded from an int32 tensor (`tile_expert`) — or derived via `pid0 // N_M_TILES` — and `stride_w_E` fits in int32, so Triton evaluates `expert_idx * stride_w_E` in int32.
For `gate_up_proj` of shape `(E, 2*I, H)`:
```
stride_w_E = 2 * I * H = 2 * 1024 * 4096 = 8,388,608
max offset = (E-1) * stride_w_E = 511 * 8,388,608 = 4,286,578,688
```
`4,286,578,688 > 2^31-1 (2,147,483,647)` → overflow → negative offset → OOB access once the routed expert index gets high enough (~256+).
Quick check across a few models (max expert offset into `gate_up_proj` vs int32 max):
| Model | E | hidden | moe_inter | max expert offset | overflows int32? |
|---|---|---|---|---|---|
| Qwen3.5-35B-A3B | 256 | 2048 | 512 | 534,773,760 | no |
| Qwen3.5-122B-A10B | 256 | 3072 | 1024 | 1,604,321,280 | no |
| **Qwen3.5-397B-A17B** | **512** | **4096** | **1024** | **4,286,578,688** | **yes** |
### Reproduce
Multi-node LoRA SFT of **Qwen3.5-397B-A17B** via TRL `SFTTrainer` (DeepSpeed ZeRO-3, gradient checkpointing, `use_liger_kernel=True`, bf16). It crashes on the first training step inside the fused MoE kernel. The same recipe works fine on the smaller Qwen3.5-35B / 122B MoE checkpoints — only the 397B one triggers it.
### Versions
Environment Report:
-------------------
Operating System: Linux-6.8.0-71-generic-x86_64-with-glibc2.39
Python version: 3.13.12
Liger Kernel version: 0.8.0
PyTorch version: 2.12.1+cu130
CUDA version: 13.0
HIP(ROCm) version: Not available
Triton version: 3.7.1
Transformers version: 5.13.0
XPU version: XPU Not Available
Contributor guide
Research direction
Start in liger_kernel/ops/fused_moe_kernels.py and trace the offset expressions used by _fused_up_proj_swiglu_kernel, _moe_bwd_dW2_kernel, and _moe_bwd_dW1_kernel. Reproduce with the Qwen3.5-397B-A17B configuration, then verify that expert offsets no longer overflow and that the first training step completes without a CUDA illegal memory access.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100