Skip to content

[CUDA] QMoE GEMV fast path for batch-1 decode #29038

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tianleiwu merged 8 commits into from
Jun 15, 2026
Merged

[CUDA] QMoE GEMV fast path for batch-1 decode #29038

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a447477
QMoE GEMV kernel
tianleiwu Jun 13, 2026
a37ac01
address feedbacks
tianleiwu Jun 13, 2026
c51e7b4
gemv for group wise quantization
tianleiwu Jun 13, 2026
9f70173
support bf16
tianleiwu Jun 13, 2026
b9c3580
support int8 per-column
tianleiwu Jun 14, 2026
a8088b9
support block size 32
tianleiwu Jun 14, 2026
f7b900e
Fix buffer aliasing and Use fp32 accumulation
tianleiwu Jun 15, 2026
79cb66a
lintrunner
tianleiwu Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 139 additions & 4 deletions docs/contrib_ops/cuda/moe_qmoe.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ and have been significantly modified for ONNX Runtime — see
10. [FP8 (W8A16) Details](#10-fp8-w8a16-details)
11. [WFP4AFP8 Details](#11-wfp4afp8-details)
12. [Future / Deferred Modes](#12-future--deferred-modes)
- [12.1 MoE GEMV Optimization Summary](#121-moe-gemv-optimization-summary)
13. [Testing](#13-testing)
14. [Build Configuration](#14-build-configuration)
15. [Limitations & Known Issues](#15-limitations--known-issues)
Expand Down Expand Up @@ -175,16 +176,42 @@ values that require them are rejected at construction time:

## 4. Architecture Dispatch & Kernel Paths

The runner selects between three CUTLASS kernel families at runtime. The choice is
The runner selects between three CUTLASS kernel families and one small-row GEMV
fast path at runtime. The choice is
made by `CutlassMoeFCRunner::supportsTmaWarpSpecialized()` and the dispatch headers
under [onnxruntime/contrib_ops/cuda/llm/moe_gemm/](onnxruntime/contrib_ops/cuda/llm/moe_gemm/).

| Path | CUTLASS class | Used for | SM range |
|------|---------------|----------|----------|
| **MoE GEMV fast path** | `fpA_intB_gemv`-based custom kernel | INT4/INT8 per-column W*A16 and symmetric INT4/INT8 block-wise W*A16 with FP16 or BF16 activations and true decode row counts | SM80+ |
| **Ampere GemmGrouped** | `cutlass::gemm::kernel::GemmGrouped` | INT4/INT8 W*A16, FP8 W8A16 dequant fallback, FP32 | SM75–SM89, plus all mixed-input on SM90/SM120 |
| **TMA Warp-Specialized (mixed-input)** | `CollectiveBuilderMixedInput` | Same-type FP16×FP16 / BF16×BF16, native MXFP4 W4A16 | SM90 (same-type), SM120 (FP4 W4A16) |
| **Block-Scaled Tensor Op** | `OpClassBlockScaledTensorOp` | Native FP8×MXFP4 (`wfp4afp8`) | SM100+ (Blackwell) |

The MoE GEMV fast path is selected before the Ampere grouped GEMM for integer
QMoE when all of the following are true:

- activation/output dtype is FP16 or BF16;
- scales and biases use the same dtype as the activation;
- weights are INT4 or INT8 for per-column scales, or INT4/INT8 for symmetric block-wise scales;
- `block_size <= 0` for per-column INT4/INT8, or `block_size` is 32, 64, or 128 for block-wise INT4/INT8;
- block-wise GEMV is symmetric only; if zero-point compensation is present, dispatch falls back to grouped GEMM;
- `expanded_num_rows = num_tokens * top_k` is in `(0, 8]`;
- `N >= 512` and `K >= 512`;
- if `expanded_num_rows > 4`, the logical MoE intermediate size is at least 512;
- `N` is divisible by the column-interleaved tile width (32 for INT4, 16 for INT8),
and `K` satisfies the kernel step and, for block-wise scales, complete-block alignment.

Asymmetric block-wise quantization, broader row counts, and dimensions
outside the profiled gate stay on grouped GEMM until profile data shows an
end-to-end GEMV win. FP16 and BF16 share the same dispatch gate and custom
kernels; for a given shape, BF16 routes to GEMV exactly where FP16 does and
shows comparable latency.

Set `ORT_DISABLE_MOE_GEMV=1` before process start to force the grouped GEMM
fallback for debugging, benchmarking, or bisecting numerical differences. The
switch is cached on first use.

### 4.1 Per-mode dispatch matrix

| Mode | SM75-89 (Ampere/Ada) | SM90 (Hopper) | SM100 (Blackwell) | SM120 (RTX 5090) |
Expand Down Expand Up @@ -338,6 +365,10 @@ This is the layout transform applied either offline by
`weights_prepacked=0` (see [§5.1](#51-weights-input-2--5--8)).


INT MoE/QMoE CUDA kernels, including the small-row MoE GEMV path, consume the
SM80 `ColumnMajorTileInterleave<64, 4>` layout. Pack INT4/INT8 MoE weights with
the SM80 target layout even when the runtime GPU is Hopper or newer.

1. **Input layout**: `[N, K]` per expert (Out × In), 2 elements per byte for INT4.
2. **Transpose & signed conversion**:
- Unpack `uint4 [0, 15]` → subtract 8 → `int8 [-8, 7]`.
Expand Down Expand Up @@ -424,9 +455,9 @@ Dequantization (symmetric): `W = (W_stored - 128) * scale`.

| Architecture | Activation | Supported `block_size` |
|--------------|-----------|------------------------|
| SM75–89 (Turing/Ampere/Ada) | FP16/BF16 | 64, 128 |
| SM90 (Hopper) | FP16/BF16 | any multiple of 64 |
| SM100/120 (Blackwell) | FP16/BF16 | falls back to Ampere — 64 or 128 |
| SM75–89 (Turing/Ampere/Ada) | FP16/BF16 | 32, 64, 128 |
| SM90 (Hopper) | FP16/BF16 | falls back to Ampere — 32, 64, 128 |
| SM100/120 (Blackwell) | FP16/BF16 | falls back to Ampere — 32, 64, 128 |

For MXFP4, the block size is fixed at **32** by the format.

Expand Down Expand Up @@ -882,6 +913,110 @@ software dequant of the mixed-input path.
The schema reserves the necessary input slots (18–21) so adding these modes
will not change the operator interface.

### 12.1 MoE GEMV Optimization Summary

The MoE GEMV work targets decode-sized integer QMoE workloads where grouped GEMM
launch overhead, prologue overhead, and intermediate traffic dominate the FC
compute. Detailed measurements are recorded in
[qmoe_gemv_experiments.md](qmoe_gemv_experiments.md). This section is the
implementation summary and current backlog.

#### Completed per-column INT4 work

Per-column here means the original INT4 W4A16 path with one scale per output
column: scale tensors are `[E, N]`, and `block_size <= 0`.

| Area | What is implemented | Result |
|------|---------------------|--------|
| Benchmarking | Added `profile_qmoe_gemv.py` and `profile_qmoe_gemv.sh` to run GEMV-enabled and `ORT_DISABLE_MOE_GEMV=1` grouped-GEMM profiles in separate processes. | Stable A/B profiles with the `benchmark` NVTX range; use `parse_nsys.py --pattern '%'` so fallback CUTLASS kernels are visible. |
| Route policy | Dispatch is data-gated to FP16/BF16 integer QMoE decode shapes, `expanded_num_rows <= 8`, `N >= 512`, `K >= 512`, and profiled alignment constraints. | Keeps tiny shapes and unprofiled row counts on grouped GEMM. |
| Row-to-expert lookup | The prologue materializes the local expert id for each permuted row and passes it to the GEMV kernels. | Removes repeated prefix-offset scans inside each N-tile CTA; GPT-OSS and Gemma model-shape kernels improved. |
| FC1 interleaved SwiGLU | The FC1 GEMV path can apply interleaved SwiGLU in the GEMV epilogue for the profiled FP16/BF16 INT4 path. | Removes the separate activation launch and FC1 intermediate traffic; GPT-OSS and Gemma improved end-to-end. |
| One-row finalize | `num_rows == 1`, `top_k <= 4` has a static top-k finalize specialization. | Modest GPT-OSS finalize improvement while preserving the existing FC2 GEMV parallelism. |
| End-to-end GPT-OSS | The final per-column GEMV path was validated in ORT GenAI on GPT-OSS-20B INT4. | About 15% faster than the grouped-GEMM baseline and about 8% faster than the FasterTransformer reference at batch 1. |

#### Completed per-column INT8 work

Per-column INT8 means W8A16 with one symmetric scale per output column: scale
tensors are `[E, N]` and `block_size <= 0`. The per-column path previously
required `block_size == 0` exactly in `is_moe_gemv_supported`, but the QMoE
runtime carries per-column scales as `group_size = -1` (the `QuantParams::Int`
default), so per-column INT4 *and* INT8 silently fell back to grouped GEMM. The
gate now treats any `group_size <= 0` as the per-column case, matching the GEMV
launcher dispatch that already mapped `group_size <= 0` to the `GroupSize == 0`
kernel.

| Area | What is implemented | Result |
|------|---------------------|--------|
| Gate fix | `is_moe_gemv_supported` accepts `group_size <= 0` (per-column) alongside 32/64/128 block-wise group sizes. | Per-column INT4 and INT8 now reach the GEMV kernels, and block-wise INT4/INT8 includes `block_size=32`. |
| INT8 details | The existing `(half, uint8_t)` and `(__nv_bfloat16, uint8_t)` GEMV kernel details cover per-column INT8 with no new instantiation. | FC1 interleaved-SwiGLU and FC2 per-column INT8 GEMV run for FP16 and BF16. |
| Profiling | `int8_per_column_*_1024x4096_e8` and `gpt_oss_20b_*_int8_2880x2880_e32` cases profiled in GEMV and `ORT_DISABLE_MOE_GEMV=1` modes. | Real `moe_gemv_kernel` / `moe_gemv_interleaved_swiglu_kernel` confirmed; about 1.2x–1.4x lower benchmark latency than grouped GEMM with valid output. |

#### Completed block-wise INT4/INT8 work


Block-wise here means `quant_type="int"` with `block_size` 32, 64, or 128 and scales
provided as `[E, N, K / block_size]`. QMoE prepack/runtime transposes those
scales to `[E, K_blocks, N]`, and the GEMV kernels consume that same layout.

| Area | What is implemented | Result |
|------|---------------------|--------|
| INT4 and INT8 details | The GEMV kernel details support `(half, cutlass::uint4b_t)`, `(half, uint8_t)`, and the matching `__nv_bfloat16` weight-type pairs. | Both weight types use the SM80 column-interleaved `fpA_intB` layout on all GPUs, matching the grouped-GEMM path, for FP16 and BF16 activations. |
| Group-size dispatch | GEMV templates now cover `GroupSize == 0`, 32, 64, and 128. | Per-column and block-wise paths share the same kernel structure while preserving complete-block checks. |
| Scale indexing | Block-wise scale loads use `real_offset_k / GroupSize * n + real_offset_n` with a K-loop scale step. | Reuses QMoE's existing `[E, K_blocks, N]` runtime scale layout; no new scale pack format is needed. |
| Symmetric-only gate | Block-wise GEMV runs only when zero-point compensation is absent. | Asymmetric block-wise models stay on grouped GEMM until a zero-point GEMV path is implemented and profiled. |
| Model-shape b64 profile | GPT-OSS-20B and Qwen3.6-35B-A3B were profiled with `block_size=64`. | Both use real GEMV kernels under the current 512 threshold and show about 1.4x lower benchmark latency than grouped GEMM fallback. |

#### Completed BF16 enablement

BF16 activations now share the exact dispatch gate and custom GEMV kernels with
FP16. The runtime gate relaxes from `T == half` to `T == half || T ==
__nv_bfloat16`, and `__nv_bfloat16` template instantiations were added for the
per-column INT4, block-wise INT4/INT8, and interleaved-SwiGLU GEMV kernels.

| Area | What is implemented | Result |
|------|---------------------|--------|
| Gate relaxation | `tryLaunchMoeGemvIntSymmetric` and the interleaved-SwiGLU variant accept `__nv_bfloat16` activations with `ScaleBiasType == T`. | For a given shape, BF16 routes to GEMV exactly where FP16 does. |
| Kernel instantiation | `moe_gemv.cu` adds `__nv_bfloat16` details/instantiations (group sizes 0/32/64/128, INT4/INT8, bias on/off) under `ENABLE_BF16`. | The custom FC1/FC2 GEMV kernels run for BF16; no grouped-GEMM fallback when the FP16 gate would route. |
| Profiling | GPT-OSS-20B, Qwen3.6-35B-A3B, and Gemma model shapes profiled with `block_size=64` for both dtypes. | BF16 matches FP16 routing and latency within noise (about 1.3x–1.5x faster than grouped GEMM); SwiGLU BF16 parity tests pass. |

#### Experiments rejected after profiling

| Experiment | Why it was rejected |
|------------|---------------------|
| Broad GEMV enablement for tiny 128x256 cases | Raw GEMV compute kernels were faster, but end-to-end ORT loop latency was worse than grouped GEMM. |
| Expanded rows 8/16 for 1024x4096 before model-specific tuning | GEMV compute stayed competitive, but total latency regressed for larger expanded-row counts. |
| `CtaN=16, Threads=128` | Fewer N-tile CTAs did not offset lower per-CTA efficiency; GPT and Gemma kernels slowed down. |
| `CtaN=8, Threads=64` | Lower thread count slowed both profiled model-size cases. |
| Map-specialized launch | Avoiding the runtime map/prefix branch was neutral or slightly worse and added code size. |
| Naive FC2 GEMV + finalize fusion | Serialized top-k expert GEMVs inside one CTA; GPT-OSS FC2/finalize kernel time regressed sharply. |
| One-row finalize with 128 threads | Underutilized the 2880-wide GPT-OSS output; the 256-thread static top-k variant was better. |
| Asymmetric block-size-128 fallback stress cases | Existing grouped-GEMM parity tolerance was exceeded in this environment; they were not kept in the GEMV-focused matrix. |

#### Remaining ideas not tried

- A better FC2/finalize design that preserves parallelism across top-k experts
and N tiles, then reduces partial results with bounded numerical change.
- Architecture-specific dispatch thresholds or a small autotuner for SM80, SM89,
SM90, SM100, and SM120 rather than one broad hand-tuned gate.
- Asymmetric block-wise GEMV with zero-point compensation in the kernel, if model
demand and grouped-GEMM baseline data justify the maintenance cost.
- More model-shape block-wise profiling, especially `block_size=128`, INT8, and
end-to-end GenAI runs for models that ship block-wise QMoE weights.
- Native validation on SM100/SM120 for interactions between GEMV routing and the
FP4 / WFP4AFP8 paths.

Keep `ORT_DISABLE_MOE_GEMV=1` available. It is useful for A/B testing, fallback
validation, and bisecting numerical or performance regressions. For quick
profiling, use
[profile_qmoe_gemv.sh](../../../onnxruntime/test/python/transformers/profile_qmoe_gemv.sh):

```bash
onnxruntime/test/python/transformers/profile_qmoe_gemv.sh \
--case gpt_oss_20b_m1_top4_fp16_2880x2880_e32 \
--block-size 64 --warmup 5 --repeat 100
```

---

## 13. Testing
Expand Down
Loading
Loading