[CUDA] QMoE GEMV fast path for batch-1 decode (#29038)
## Summary
Adds a symmetric weight-only **MoE GEMV fast path** for single-token
(batch-1) decode of quantized MoE models such as GPT-OSS-20B, replacing
the CUTLASS grouped-GEMM path when the expanded row count is small. The
fast path now covers **INT4 and INT8** weights, **per-column and
block-wise (group size 32/64/128)** scales, and **FP16 and BF16**
activations. On an H200 (SM90) this improves end-to-end GPT-OSS-20B INT4
token-generation throughput by roughly **15% over the CUTLASS
grouped-GEMM baseline** and about **8% over the FasterTransformer kernel
used in ORT 1.26** across prompt lengths 128/1024/2048, while producing
bit-faithful output.
## Motivation
At batch-1 decode each token expands to `top_k` rows (4 for
GPT-OSS-20B), so the MoE FC1/FC2 GEMMs are extremely skinny. The CUTLASS
grouped GEMM is built for throughput at larger M and leaves the decode
path memory-bound and underutilized. A dedicated weight-only INT GEMV
with per-expert dispatch is a better fit for this regime and closes the
gap to (and surpasses) ORT 1.26 for GPT-OSS-20B.
This also adds block_size=32 support as requested in
https://github.com/microsoft/onnxruntime/issues/29035.
## Key Changes
### New MoE GEMV kernel
| File | Change |
|---|---|
| `contrib_ops/cuda/llm/moe_gemm/moe_gemv.h` | Public interface:
`is_moe_gemv_supported` dispatch predicate and the symmetric INT
launchers `launch_moe_gemv_int_symmetric<T, WeightType>` /
`launch_moe_gemv_int_symmetric_interleaved_swiglu<T, WeightType>` (plus
the original INT4 per-channel launchers). |
| `contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu` | Symmetric INT MoE GEMV
(INT4 `uint4b_t` / INT8 `uint8_t`, FP16/BF16). One CTA per expanded row
× N-tile; per-expert weight/scale/bias offsets via a direct
row-to-expert map (prefix-offset scan fallback). Supports per-column
(`group_size <= 0`) and block-wise (`group_size` 64/128) scales. FC1 has
an interleaved SwiGLU-fused epilogue; a static top-k one-row finalize
specialization is used for FC2 routing. |
### Runner wiring and profiler
| File | Change |
|---|---|
| `contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu` / `.h` | Branch FC1 and
FC2 to the GEMV fast path when supported (carrying `group_size` through
`QuantParams`); fall back to grouped GEMM otherwise.
`ORT_DISABLE_MOE_GEMV=1` forces the grouped-GEMM path for A/B testing. |
| `contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc` / `.h`,
`moe_util_kernels.h` | Use M (token count) in the MoE GEMM profiler so
tactic selection reflects the decode shape. |
### Backward-compatible SwiGLU fusion
| File | Change |
|---|---|
| `contrib_ops/cuda/moe/moe.cc`, `moe_quantization.cc` (+ `.h`) | Treat
`swiglu_fusion == 0` with no separate FC3 as interleaved fusion (`==
1`). The published GPT-OSS-20B model (and any model exported by ORT <
1.27) hard-coded the interleaved layout and emits no `swiglu_fusion`
attribute, so it defaults to 0; those weights are pre-fused into FC1 and
must be treated as fusion mode 1. |
### Tests, profiling, and docs
| File | Change |
|---|---|
| `test/python/transformers/profile_qmoe_gemv.py` / `.sh` | Standalone
QMoE GEMV profiling harness (NVTX-ranged, GEMV-vs-grouped-GEMM kernel
comparison). |
| `test/python/transformers/test_qmoe_cuda.py`, `test_moe_cuda.py` |
QMoE GEMV decode-latency/parity coverage across INT4/INT8,
per-column/block-wise, FP16/BF16, plus import tidy-ups. |
| `docs/contrib_ops/cuda/qmoe_gemv_experiments.md` | Full experiment
log: kernel-level sweeps, dispatch-gate tuning, block-wise and BF16
enablement, INT8 per-column root-cause + fix, and the GenAI end-to-end
throughput comparison (CUTLASS baseline / GEMV final / FT baseline). |
| `docs/contrib_ops/cuda/moe_qmoe.md` | QMoE doc updates describing the
GEMV fast path and its dispatch gate. |
## Results
### GPT-OSS-20B INT4 end-to-end (H200/SM90, batch 1)
Token-generation throughput (tps, higher is better):
| Prompt length | CUTLASS baseline (gemm) | GEMV (final) | FT (ORT 1.26)
| GEMV vs cutlass | GEMV vs FT |
|---|---|---|---|---|---|
| 128 | 248.9 | 288.0 | 265.2 | +15.7% | +8.6% |
| 1024 | 237.8 | 272.2 | 252.9 | +14.5% | +7.6% |
| 2048 | 231.3 | 265.0 | 245.6 | +14.6% | +7.9% |
### Decode microbenchmark coverage (H200/SM90, batch 1)
Benchmark-loop latency in milliseconds, lower is better. `Enabled` is
the default GEMV build; `Fallback` sets `ORT_DISABLE_MOE_GEMV=1`
(grouped GEMM). Every case reported `has_invalid_output=false`.
| Case | Quant | DType | Enabled ms | Fallback ms | Speedup |
|---|---|---|---|---|---|
| `int8_per_column_m1_top2_1024x4096_e8` | INT8 per-column | FP16 |
0.0566 | 0.0816 | 1.44x |
| `int8_per_column_m1_top2_1024x4096_e8` | INT8 per-column | BF16 |
0.0578 | 0.0862 | 1.49x |
| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | INT8 per-column | FP16 |
0.0785 | 0.0947 | 1.21x |
| `gpt_oss_20b_m1_top4_int8_2880x2880_e32` | INT8 per-column | BF16 |
0.0785 | 0.0989 | 1.26x |
Block-wise INT4 (`block_size=64`, `1024x4096`, e8) routes to GEMV for
both FP16 and BF16 with FC1/FC2 kernel times within noise across dtypes
(FP16 4.53/6.95 us, BF16 4.62/7.01 us). See `qmoe_gemv_experiments.md`
for the full sweep.
## Correctness
- GEMV-enabled and `ORT_DISABLE_MOE_GEMV=1` (grouped GEMM) produce
identical, correct output for the GPT-OSS-20B sanity prompt; Nsight
traces confirm `moe_gemv_kernel` (FC2) and
`moe_gemv_interleaved_swiglu_kernel` (FC1) run for the decode shapes.
- INT4/INT8 SwiGLU parity cases pass for FP16 and BF16 (max absolute
difference ~1e-3 against the reference). Regression: `pytest -k
"TestSwigluQMoE or TestQMoEIntPrePackSmoke"` → `34 passed, 4 skipped`.
- Test on GPT-OSS-20b with block_size=32, and generated results looks
good.
## Testing Notes
- Build: `bash .env/cuda_130.sh --build` (CUDA 13.0,
`CMAKE_CUDA_ARCHITECTURES=89;90`). Use `--clean_moe` after dispatch-code
edits to avoid stale `moe_kernels.cu.o`.
- Kernel profiling:
`onnxruntime/test/python/transformers/profile_qmoe_gemv.sh`.
- Parity/latency: `pytest
onnxruntime/test/python/transformers/test_qmoe_cuda.py`.
- End-to-end: GenAI `benchmark_e2e.py` on GPT-OSS-20B INT4, batch 1,
prompt lengths 128/1024/2048. Run on an idle GPU — shared-GPU contention
corrupts the decode-latency measurement.
- A/B check: compare default vs `ORT_DISABLE_MOE_GEMV=1` for both
throughput and output parity.
