[Bug] cuDNN FP8 MOE grouped matmul silently computes only the first expert group with backend 9.18–9.20 (SM120)
op: moe
## Summary
`flashinfer.grouped_mm.grouped_mm_fp8` (cuDNN backend) returns **silently wrong results** with cuDNN backend **9.18.0 – 9.20.0**: only the **first expert group** is computed; every group `e >= 1` of the output is left **zero**. `graph.check_support()` passes, no error or warning is raised. Fixed in cuDNN **9.21.0**.
The current version floor admits the broken versions: `_CUDNN_MOE_MIN_VERSION = 91800` (`flashinfer/grouped_mm/cudnn/core.py:48`) is used for `grouped_mm_fp8` (`flashinfer/grouped_mm/core.py:269`), and the test gate `requires_cudnn_moe` (`tests/grouped_mm/conftest.py:28`) is `>= 91800` as well.
**Blast radius:** `torch 2.11.0+cu130` pins `nvidia-cudnn-cu13==9.19.0.56` — i.e. a stock cu13 PyTorch environment gets exactly a broken version, and MoE FP8 grouped GEMM silently produces garbage for every expert but the first.
## Reproduction (RTX PRO 6000 Blackwell, SM120, CUDA 13, torch 2.11.0+cu130)
```python
import torch
from flashinfer.grouped_mm import grouped_mm_fp8
torch.manual_seed(42)
E, tpe, k, n = 4, 64, 256, 256
a = torch.randn(E * tpe, k, device="cuda").clamp(-1, 1).to(torch.float8_e4m3fn)
b = torch.randn(E, n, k, device="cuda").clamp(-1, 1).to(torch.float8_e4m3fn)
m_indptr = (torch.arange(E + 1, device="cuda") * tpe).to(torch.int32)
alpha = torch.tensor([1.0], dtype=torch.float32, device="cuda")
out = grouped_mm_fp8(a, b, m_indptr, alpha=alpha, out_dtype=torch.bfloat16)
for e in range(E):
s, t = e * tpe, (e + 1) * tpe
ref = a[s:t].float() @ b[e].float().T
ok = torch.isclose(out[s:t].float(), ref, atol=0.125, rtol=0.125).float().mean()
zero = (out[s:t] == 0).float().mean()
print(f"expert {e}: match={ok:.3f} zero-frac={zero:.3f}")
```
With cuDNN 9.19.0.56:
```
expert 0: match=1.000 zero-frac=0.004
expert 1: match=0.014 zero-frac=1.000
expert 2: match=0.010 zero-frac=1.000
expert 3: match=0.008 zero-frac=1.000
```
Same pattern for `float8_e5m2`, all shapes tried (`tpe ∈ {32..256}`, `k ∈ {256,512}`, `n ∈ {256,1024}`), with and without `alpha`. `pytest tests/grouped_mm/test_grouped_mm_fp8.py` → **58 failed** (everything with `num_experts > 1`).
## What rules out the wrapper
`grouped_mm_bf16` goes through the **same** `_run_cudnn_moe_grouped_gemm` / `_build_cudnn_moe_grouped_gemm_graph` code path (same `first_token_offset` construction, `flashinfer/grouped_mm/cudnn/core.py:248`) and is correct on every cuDNN version tested — `tests/grouped_mm/test_grouped_mm_bf16.py` = 46 passed on 9.19.0.56. Only the FP8 data types trip it, so the `first_token_offset`/graph semantics on the FlashInfer side are fine; the regression is inside the cuDNN backend's FP8 MOE kernels.
## cuDNN backend bisect (pip `nvidia-cudnn-cu13`, `tests/grouped_mm/test_grouped_mm_fp8.py`)
| cuDNN backend | result |
| --- | --- |
| 9.18.0.77 | 58 failed / 34 passed |
| 9.18.1.3 | 58 failed / 34 passed |
| 9.19.0.56 (torch 2.11 cu13 pin) | 58 failed / 34 passed |
| 9.19.1.2 | 58 failed / 34 passed |
| 9.20.0.48 | 58 failed / 34 passed |
| **9.21.0.82** | **92 passed** |
| 9.22.0.52 | 92 passed |
| 9.23.2.1 | 92 passed |
(The passing subset on broken versions is the `num_experts = 1` / empty-expert cases.)
Bonus datapoint: with 9.23.2.1 the whole `tests/grouped_mm/` directory is **305 passed** on SM120, including the `>= 9.21` block-scale MXFP8/FP4 suites.
## Scope caveat
Verified on **SM120 (RTX PRO 6000 Blackwell, CUDA 13)** only — I don't have SM90/SM100 hardware to check whether 9.18–9.20 mis-execute there too. bf16 vs fp8 discrimination above is from the same SM120 box.
## Proposed fix
Mirror the existing block-scale floor: add `_CUDNN_MOE_FP8_MIN_VERSION = 92100` and use it in `grouped_mm_fp8`'s runtime check (`flashinfer/grouped_mm/core.py:269`), plus a matching `requires_cudnn_moe_fp8` test gate, so the broken combination fails fast with an actionable error instead of silently corrupting MoE outputs. PR incoming.
Found while validating current `main` (`8fc7f079`) on RTX PRO 6000 for the v0.6.14 SM12x release push (#3783). AI-assisted (Claude Code) under human direction; every number above is from live runs on the real card.
1 条评论