diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 77709bc436e66..8a255c4c897c4 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -4935,7 +4935,7 @@ This version of the operator has been available since version 1 of the 'com.micr
activation_type : string
Activation function to use. Choose from relu, gelu, silu, swiglu and identity. Default is relu
block_size : int
-
Size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). If provided, both hidden_size and inter_size must be divisible by the block size. Otherwise, there is no blocking and a whole column shares one scaling factor.
+
Size of each quantization block along the K (input feature) dimension. Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). Both hidden_size and inter_size must be divisible by the block size. The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. For integer quantization ('int'), omitting block_size means there is no blocking and a whole column shares one scaling factor.
expert_weight_bits : int
Number of bits used in quantized weights. Supported values are 2, 4, and 8. Default is 4 bits
k : int
diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index dd451cbe44c03..44cd7928c4b26 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -21,6 +21,11 @@ and have been significantly modified for ONNX Runtime — see 7. [Cross-Architecture Packing Compatibility](#7-cross-architecture-packing-compatibility) 8. [SwiGLU Fusion](#8-swiglu-fusion) 9. [FP4 (MXFP4) Details](#9-fp4-mxfp4-details) + - [9.9 Runtime environment variables](#99-runtime-environment-variables) + - [9.10 Interleaved GEMV layout + dtype-conditional accumulation](#910-interleaved-gemv-layout--dtype-conditional-accumulation) +9b. [NVFP4 (W4A16, block-16) Details](#9b-nvfp4-w4a16-block-16-details) + - [9b.2 Fused GEMV decode (group-16)](#9b2-fused-gemv-decode-group-16) + - [9b.3 Fast E2M1 → half/bf16 decode](#9b3-fast-e2m1--halfbf16-decode) 10. [FP8 (W8A16) Details](#10-fp8-w8a16-details) 11. [WFP4AFP8 Details](#11-wfp4afp8-details) 12. [Future / Deferred Modes](#12-future--deferred-modes) @@ -70,8 +75,8 @@ input tokens → router (top-k softmax) → permute by expert | `activation_beta` | float | `0.0` | SwiGLU beta. Default `0.0` (Standard SwiGLU); GPT-OSS uses `1.0`. | | `swiglu_limit` | float | unset (`+inf`) | SwiGLU clamp limit. Unset means no clamp (Standard SwiGLU); GPT-OSS uses `7.0`. | | `expert_weight_bits` (QMoE only) | int | 4 | 4 (INT4/MXFP4) or 8 (INT8/FP8). | -| `block_size` (QMoE only) | int | -1 | Group size for INT4/INT8 group-wise quantization. -1 = per-output-channel. | -| `quant_type` (QMoE only) | string | `"int"` | `"int"`, `"fp4"`, `"fp8"`, `"wfp4afp8"`. See [§3](#3-quantization-modes). | +| `block_size` (QMoE only) | int | -1 | Group size for INT4/INT8 group-wise quantization. -1 = per-output-channel. FP4/WFP4AFP8 normalize an omitted value to 32 and require 32 when present; NVFP4 normalizes an omitted value to 16 and requires 16 when present. | +| `quant_type` (QMoE only) | string | `"int"` | `"int"`, `"fp4"`, `"nvfp4"`, `"fp8"`, `"wfp4afp8"`. See [§3](#3-quantization-modes). | | `weights_prepacked` (QMoE only) | int | -1 | Tri-state, only meaningful when `quant_type="int"`. The prepacked layouts selected by `-1` and `1` are **EP-determined**. `-1` (default): the INT4/INT8 `fc1`/`fc2` initializers are already prepacked in the EP's default layout (e.g. from `pack_weights_for_cuda_mixed_gemm` for the CUDA EP). `1`: already prepacked in an alternate EP-selected layout. `0`: the initializers are raw `[E, N, K/pack]` tensors (as produced by `quantize_matmul_{4,8}bits`) and the kernel runs the CUTLASS layout transform in `PrePack()`. **Note:** the CUDA EP INT4/INT8 MoE GEMM always runs the Ampere (SM80) kernel — even on SM90 — so it consumes the SM80 `fpA_intB` layout on all architectures; `-1` and `1` are therefore equivalent for the CUDA EP today, and `1` is reserved for a possible future Hopper-specific layout. See [§5.1](#51-weights-input-2--5--8). | ### 2.2 Type Constraints @@ -79,8 +84,8 @@ input tokens → router (top-k softmax) → permute by expert | Constraint | Allowed Types | Used By | |------------|---------------|---------| | `T` | `float`, `float16`, `bfloat16` | input, output, biases, router | -| `T1` | `uint8`, `float8e4m3fn` | quantized weights and zero points: INT4/INT8/FP4 weights use `uint8`; FP8 weights use `float8e4m3fn` | -| `T2` | `float`, `float16`, `bfloat16`, `uint8` | INT4/INT8 weight scales use floating-point tensors; MXFP block scales use `uint8` storage | +| `T1` | `uint8`, `float8e4m3fn` | quantized weights and zero points: INT4/INT8/FP4/NVFP4 weights use `uint8`; FP8 weights use `float8e4m3fn` | +| `T2` | `float`, `float16`, `bfloat16`, `uint8`, `float8e4m3fn` | INT4/INT8 weight scales use floating-point tensors; MXFP block scales use `uint8` (`float8e8m0`) storage; **NVFP4 block scales use `float8e4m3fn`** | | `T4` | `float` | per-expert global scales, FP8 activation scales | ### 2.3 Inputs @@ -93,10 +98,10 @@ to the selected `quant_type` are simply omitted (most are `Optional`). | 0 | `input` | T | `(num_tokens, hidden_size)` | all | | 1 | `router_probs` | T | `(num_tokens, num_experts)` | all | | 2 | `fc1_experts_weights` | T1 | `(E, fusion×inter, hidden/pack)` | all | -| 3 | `fc1_scales` | T2 (Opt) | varies — see [§2.4](#24-input-369-interpretation-by-quant_type) | int, fp4, wfp4afp8 | +| 3 | `fc1_scales` | T2 (Opt) | varies — see [§2.4](#24-input-369-interpretation-by-quant_type) | int, fp4, nvfp4, wfp4afp8 | | 4 | `fc1_experts_bias` | T (Opt) | `(E, fusion×inter)` | optional | | 5 | `fc2_experts_weights` | T1 | `(E, hidden, inter/pack)` | all | -| 6 | `fc2_scales` | T2 (Opt) | varies | int, fp4, wfp4afp8 | +| 6 | `fc2_scales` | T2 (Opt) | varies | int, fp4, nvfp4, wfp4afp8 | | 7 | `fc2_experts_bias` | T (Opt) | `(E, hidden)` | optional | | 8 | `fc3_experts_weights` | T1 (Opt) | `(E, inter, hidden/pack)` | optional (SwiGLU split-weight) | | 9 | `fc3_scales` | T2 (Opt) | varies | optional | @@ -105,8 +110,8 @@ to the selected `quant_type` are simply omitted (most are `Optional`). | 12 | `fc2_zero_points` | T1 (Opt) | matches `fc2_scales` | int only | | 13 | `fc3_zero_points` | T1 (Opt) | matches `fc3_scales` | optional, int only | | 14 | `router_weights` | T (Opt) | `(num_tokens, num_experts)` | optional (DeepSeek noaux_tc) | -| 15 | `fc1_global_scale` | T4 (Opt) | `(num_experts,)` | fp4, fp8, wfp4afp8 | -| 16 | `fc2_global_scale` | T4 (Opt) | `(num_experts,)` | fp4, fp8, wfp4afp8 | +| 15 | `fc1_global_scale` | T4 (Opt) | `(num_experts,)` | fp4, nvfp4, fp8, wfp4afp8 | +| 16 | `fc2_global_scale` | T4 (Opt) | `(num_experts,)` | fp4, nvfp4, fp8, wfp4afp8 | | 17 | `fc1_act_scale` | T4 (Opt) | `(1,)` or `(num_experts,)` | wfp4afp8 (Variant A) | | 18 | `fc2_act_scale` | T4 (Opt) | `(1,)` or `(num_experts,)` | wfp4afp8 (Variant A) | | 19 | `fc1_act_block_scale` | T2 (Opt, float8e8m0) | `(E, M_pad, K/32)` | wfp4afp8 (Variant B) | @@ -127,11 +132,12 @@ is used for both (backward compatible). | `"int"` (group-wise) | float / fp16 / bf16 | `(E, N, K/block_size)` | `w_float = w_int × scale (+ zero)` | | `"int"` (per-channel) | float / fp16 / bf16 | `(E, N)` | per-output-channel scale | | `"fp4"` | uint8 (`float_ue8m0_t`) | `(E, N, K/32)` | MXFP4 block scale, group=32 | +| `"nvfp4"` | uint8 (`float8e4m3fn` bytes) | `(E, N, K/16)` | NVFP4 block scale, group=16 (needs `fc*_global_scale`) | | `"fp8"` | — | — | not used; only the per-expert global scale (input 15/16/17) is needed | | `"wfp4afp8"` | uint8 (`float_ue8m0_t`) | `(E, N, K/32)` | MXFP4 block scale, group=32 | Inputs 11/12/13 (`fc*_zero_points`) are valid only for `"int"`. FP8 e4m3 and -FP4 e2m1 are symmetric formats with no zero-point. +FP4 e2m1 (both MXFP4 and NVFP4) are symmetric formats with no zero-point. --- @@ -143,19 +149,21 @@ FP4 e2m1 are symmetric formats with no zero-point. | `"int"` (8-bit) | W8A16 | FP16/BF16 | INT8 group-wise | SM75+ | — | always | | `"fp8"` | W8A16-fp8 | BF16/FP16 | FP8 e4m3 (no packing) | **SM90+** native | dequant→A16 on SM<90 | `ENABLE_FP8` (CUDA ≥ 11.8) | | `"fp4"` | W4A16-MXFP4 | BF16/FP16 | MXFP4 e2m1, group=32 | **SM120+** native | dequant→A16 on SM<120 | `ENABLE_FP4` + `USE_FP4_QMOE` (CUDA ≥ 12.8) | +| `"nvfp4"` | W4A16-NVFP4 | BF16/FP16 | NVFP4 e2m1, group=16, `float8e4m3fn` block scale + per-expert FP32 global scale | dequant→A16 (all SM) + **fused GEMV decode** | `ENABLE_FP4` + `USE_FP4_QMOE` (CUDA ≥ 12.8) | | `"wfp4afp8"` | W4A8-MXFP4×FP8 | FP8 e4m3 (quantized in-runner) | MXFP4 e2m1, group=32 | **SM100+** native | dequant→A16 on SM<100 | `ENABLE_FP4` + `USE_FP4_QMOE` + `ENABLE_FP8` | Selection logic (see [moe_quantization.cc](onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc)): ```cpp if (quant_type_ == "fp4") use_fp4_dequant_fallback_ = (sm_ < 120); +if (quant_type_ == "nvfp4") use_fp4_dequant_fallback_ = true; // no native block-scaled path yet; fused GEMV decode covers small-decode shapes if (quant_type_ == "wfp4afp8") use_wfp4afp8_dequant_fallback_ = (sm_ < 100); if (quant_type_ == "fp8") use_fp8_dequant_fallback_ = (sm_ < 90); ``` `expert_weight_bits` validation: - `int` → 4 or 8 -- `fp4`, `wfp4afp8` → must be 4 +- `fp4`, `nvfp4`, `wfp4afp8` → must be 4 - `fp8` → must be 8 When the build is configured without the corresponding flags, `quant_type` @@ -183,7 +191,7 @@ under [onnxruntime/contrib_ops/cuda/llm/moe_gemm/](onnxruntime/contrib_ops/cuda/ | 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+ | +| **MoE GEMV fast path** | `fpA_intB_gemv`-based custom kernel | INT4/INT8 per-column W*A16 and symmetric INT4/INT8 block-wise W*A16, and **MXFP4 (group-32) / NVFP4 (group-16) W4A16** decode, 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) | @@ -220,6 +228,7 @@ switch is cached on first use. | FP16/BF16 (no quant, MoE op) | Ampere GemmGrouped | TMA WS (same-type) | TMA WS / valid Blackwell spec | TMA WS / Ampere fallback | | FP8 W8A16 native | dequant fallback | TMA WS | TMA WS | SM89 FP8 kernel redirect | | FP4 W4A16 native | dequant fallback | dequant fallback | dequant fallback | TMA WS mixed-input FP4 | +| NVFP4 W4A16 (group-16) | dequant fallback + fused GEMV decode | dequant fallback + fused GEMV decode | dequant fallback + fused GEMV decode | dequant fallback + fused GEMV decode | | WFP4AFP8 native | dequant fallback | dequant fallback | Block-scaled tensor op | Block-scaled tensor op | | FP32 | Ampere GemmGrouped (forced) | same | same | same | @@ -752,6 +761,142 @@ switch (hopper_inputs.fusion) { | `moe_gemm_template_dispatch_tma_ws_mixed_dtype.h` | `FUSION` param throughout; `PackedScalesNum`-based K tile; direct N tile mapping; workspace calc with `Ntile=128` | | `moe_gemm_template_dispatch.h` | FUSION routing in `dispatchToArch`; removed restrictive wfp4a16 config filter | +### 9.9 Runtime environment variables + +The FP4 path is gated by several process-start environment variables (read once in the +QMoE constructor). Defaults give the validated production behavior; the rest are opt-in or +debug switches. + +| Variable | Default | Effect | +|----------|---------|--------| +| `ORT_ENABLE_FP4_GEMV` | on | Fused MXFP4 GEMV decode kernel. Set to `0` to force the dequant-to-dense fallback (debugging/bisecting). Active in the SM<120 fallback regime, and as the decode arm when native CUTLASS prefill is enabled. | +| `ORT_FP4_GEMV_AUTOTUNE` | `0` | Opt-in per-shape autotune of the GEMV CtaN/Threads tiling. Enabling it synchronizes the first uncached inference for each shape. | +| `ORT_FP4_GEMV_AUTOTUNE_LOG` | `0` | Set to `1` to log the chosen GEMV configs per shape. | +| `ORT_FP4_GEMV_INTERLEAVED` | `0` | **Experimental, opt-in.** Routes the MXFP4 decode GEMV through the `ColumnMajorInterleaved` weight layout (`kInterleave=4`, `kStepK=32`) with dtype-conditional accumulation. fp16 gets faster decode; bf16 stays accuracy-safe. Default off keeps the shipping `ColumnMajor` path byte-for-byte unchanged. See [§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation). | +| `ORT_FP4_GEMV_INTERLEAVED_HALFACC` | `0` | **Override.** When `ORT_FP4_GEMV_INTERLEAVED=1`, forces 16-bit accumulation for *both* fp16 and bf16, overriding the dtype-conditional policy; regresses bf16 accuracy, so it is off by default. | +| `ORT_FP4_SM80_GEMM` | `1` | Routes SM80/Ampere FP4 prefill through the fused-dequant grouped GEMM. Set to `0` to force dense fallback for debugging or comparison. Decode still routes through fused MXFP4 GEMV when supported. | +| `ORT_ENABLE_FP4_CUTLASS_GEMM` | `0` | Opt-in native SM90 WFP4A16 CUTLASS GEMM (fast prefill). Requires FP16, SM90, and aligned shapes (`hidden`/`inter` divisible by 256). Must be combined with `ORT_ENABLE_FP4_CUTLASS_UNSAFE=1`. | +| `ORT_ENABLE_FP4_CUTLASS_UNSAFE` | `0` | Confirms use of the experimental native SM90 path. Without it, a request to enable native GEMM logs a warning and falls back to dequant/GEMV. | +| `ORT_FP4_PREFILL_MIN_TOKENS` | `64` | When native CUTLASS is enabled, the per-node decode threshold. Tokens with `M >= threshold` (prefill) route to native CUTLASS; `M < threshold` (decode) route to the fused GEMV. Both weight/scale layouts are pre-packed so one node serves both regimes. | + +When native CUTLASS is enabled, weights and scales are dual-prepacked (native layout plus +raw e8m0 block scales for GEMV), so `ORT_FP4_PREFILL_MIN_TOKENS` selects per call between the +prefill and decode kernels with no extra conversion. + +### 9.10 Interleaved GEMV layout + dtype-conditional accumulation + +**Opt-in (`ORT_FP4_GEMV_INTERLEAVED=1`, default off).** The shipping MXFP4 decode GEMV uses a +non-interleaved `ColumnMajor` weight layout (`kInterleave=1`, `kStepK=8`). This experimental path +instead mirrors the INT4 GEMV's `ColumnMajorInterleaved` layout (`kInterleave=4`, `kStepK=32`, 4× +fewer K-loop trips) and pairs it with a **dtype-conditional accumulator** to balance bf16 accuracy +against occupancy: + +```cpp +// onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +template +using Fp4GemvAccT = std::conditional_t::value, half, float>; +``` + +| Activation | Accumulator | Rationale | +|-----------|-------------|-----------| +| **fp16** | fp16 (`half`) | fp16's 10-bit mantissa tolerates 16-bit accumulation over the longer `kStepK=32` chains; the cheaper accumulator keeps register use low so the interleaved layout's K-trip savings translate into a real speedup. | +| **bf16** | fp32 (`float`) | bf16's 7-bit mantissa loses too much precision under 16-bit accumulation (fails tolerance), so it must accumulate in fp32 — at the cost of the extra registers that erase the speedup. | + +The interleaved weights are produced by the `gemv_interleaved` branch of +[`PrePackRepackFP4Weights`](onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc), which routes the +raw e2m1 codes per-expert through the CUTLASS fpA_intB SM80 `W4_A16` preprocessor with +`apply_bias_interleave=false` (the integer-only `+8`/pair-interleave step would corrupt the +floating-point e2m1 codes; the layout-only steps apply unchanged). Block scales are unchanged +(`kStepK=32` equals the MXFP4 block size). The shape gate requires `n % (CtaN*4)==0` and +`k % 64==0`; `CtaN`/`Threads` are pinned (`4`/`128`) so the kernel always matches the prepacked +weights. + +The interleaved layout speeds up fp16 decode; bf16 gets no speedup (the fp32 register cost cancels +it) but stays accurate. The path is opt-in so fp16 deployments can take the win without affecting +bf16 or the shipping default. + +--- + +## 9b. NVFP4 (W4A16, block-16) Details + +`quant_type="nvfp4"` is a second FP4 weight-only mode, distinct from MXFP4 (`"fp4"`): + +| | MXFP4 (`"fp4"`) | **NVFP4 (`"nvfp4"`)** | +|---|---|---| +| Weight codes | E2M1, 2 codes/byte (`uint8`) | E2M1, 2 codes/byte (`uint8`) — same | +| Block size | 32 | **16** | +| Block scale | `float8e8m0` (E8M0, pow-of-two), stored `uint8` | **`float8e4m3fn`** (E4M3) | +| Global scale | 1.0 (implicit) | **per-expert FP32** `fc*_global_scale` (`weight_scale_2`) | +| Reconstruct | `w = e2m1(code) · e8m0(block_scale)` | `w = e2m1(code) · e4m3(block_scale[n, k/16]) · global_scale[e]` | + +NVFP4 is the format emitted by NVIDIA Model-Optimizer (e.g. `nvidia/Qwen3.6-35B-A3B-NVFP4`). + +### 9b.1 Dispatch + +There is no native block-scaled CUTLASS path for NVFP4 today (that is Blackwell-only and not yet +wired), so NVFP4 **always** uses the dequant-to-A16 fallback ([§4.3](#43-dequant-to-a16-fallback)) +for prefill / GEMV-unsupported shapes, and the **fused GEMV decode fast path** +([§4](#4-architecture-dispatch--kernel-paths)) for small-decode shapes. `enable_fp4_gemv_` is on by +default for NVFP4 (opt-out `ORT_ENABLE_FP4_GEMV=0`); `enable_fp4_sm80_gemm_` stays off (the SM80 +grouped-GEMM FP4 prefill path is MXFP4-only). + +### 9b.2 Fused GEMV decode (group-16) + +The MXFP4 decode GEMV ([§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation)) is +block-size generic: the kernel indexes scales as `scales[e][k/group_size][n]`, and `GroupSize` is a +compile-time template (`static_assert((CtaK/kInterleave) % GroupSize == 0)`). NVFP4 reuses it with +`group_size = 16`: + +- `is_moe_gemv_fp4_supported` accepts `group_size ∈ {16, 32}`; the dispatch instantiates the + `GroupSize=16` cases in `dispatch_moe_gemv_group_size` / + `dispatch_moe_gemv_interleaved_swiglu_group_size` ([moe_gemv_device.cuh](onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh)). +- NVFP4 uses **only** the non-interleaved `ColumnMajor` layout; the opt-in interleaved path + ([§9.10](#910-interleaved-gemv-layout--dtype-conditional-accumulation)) is MXFP4-only because its + `kStepK=32` tile is tied to the block-32 scale layout. +- `QMoECombineNvfp4ScalesForGemv` ([qmoe_kernels.cu](onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu)) + decodes the `float8e4m3fn` block scales, folds in the per-expert FP32 global scale, and rewrites + `[E, n, k/16] → [E, k/16, n]` in the activation dtype (`TypeA`) that the GEMV expects. +- The decode gate ([moe_quantization.cc](onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc)) fires + when `expanded = num_tokens·top_k ∈ (0, 8]`, SwiGLU is fused, and both FC1 + (`n=2·inter`, `k=hidden`) and FC2 (`n=hidden`, `k=inter`) satisfy `n,k ≥ 512` and group-16 block + alignment. For Qwen3.6-35B-A3B (`hidden=2048`, `inter=512`, `E=256`, `top_k=8`) both GEMMs qualify. + +### 9b.3 Fast E2M1 → half/bf16 decode + +Profiling the NVFP4 decode GEMV on H200 (SM90) at the Qwen shapes showed the FC1/FC2 kernels are +**ALU-pipeline bound** (ncu: ALU ≈ 79%, DRAM ≈ 7%), i.e. the per-element FP4 dequantization — not +memory or tiling — dominates. The original `Fp4I2FConverter::decode` used a `float` lookup table, +a sign branch, and a per-element `float → half/bf16` conversion. It is replaced by a **branchless +`prmt.b32` byte-select** that builds the 16-bit float bit pattern directly (E2M1 has only eight +magnitudes, so the half/bf16 encodings are packed into two 32-bit constants and selected by the +low three code bits, with the sign bit shifted into place): + +```cpp +// onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h — Fp4I2FConverter::decode (bit-identical to the LUT) +uint32_t sel = code & 0x7u, sign = uint32_t(code & 0x8u) << 12; // sign -> bit 15 +// half: low byte always 0 -> one prmt; bf16: two prmt (high + low byte). +``` + +This is numerically bit-identical to the previous LUT. Measured decode speedup (H200, `tokens=1`, +`top_k=8`, `hidden=2048`, `inter=512`, `E=256`, bf16, autotuned tiling): + +| kernel | before | after | Δ | +|--------|-------:|------:|---| +| FC1 GEMV (SwiGLU-fused, `n=1024 k=2048`) | 20.6 µs | **13.25 µs** | **−36%** | +| FC2 GEMV (`n=2048 k=512`) | 10.7 µs | **7.55 µs** | **−29%** | + +ncu confirms the mechanism: FC1 compute (SM) throughput dropped 68.8% → 55.9% and DRAM rose +7.3% → 10.2% (a more balanced kernel). Full-node latency (incl. expand / sort / finalize / +softmax-topk overhead) fell 66.8 µs → 56.1 µs, and end-to-end Qwen3.6 decode reaches ~131 tok/s on +H200. + +### 9b.4 Testing & benchmarking + +- Parity: [test_qmoe_nvfp4_cuda.py](onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py) + (`-k gemv` compares `ORT_ENABLE_FP4_GEMV=1` vs `=0` vs a torch reference). +- Kernel microbench: `bench_qmoe_nvfp4_gemv.py` builds the QMoE node at the Qwen decode shape and + times GEMV vs the dequant fallback (with `ORT_FP4_GEMV_AUTOTUNE_LOG=1` for the chosen tiling). + --- ## 10. FP8 (W8A16) Details diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp index b713c97747fb2..2d79718e0e986 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp @@ -379,7 +379,7 @@ struct MixedGroupedGemmInputUtils { constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa; asm volatile( "{\n" - " lop3 .b32 %0, %2, %4, %5, %6;\n" + " lop3.b32 %0, %2, %4, %5, %6;\n" " xor .b32 %1, %3, %5; \n" "}\n" : "=r"(scale_pos_[0]), "=r"(scale_pos_[1]) diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h index c0656ac784830..51cffac23f48d 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h @@ -107,6 +107,23 @@ struct LayoutDetailsB +struct LayoutDetailsB= 75>::type> { + static constexpr int ThreadblockK = 128 * 8 / cutlass::sizeof_bits::value; + + private: + static constexpr int ElementsPerCacheLine = 128 * 8 / sizeof_bits::value; + static constexpr int ColumnsInterleaved = ElementsPerCacheLine / ThreadblockK; + + public: + using Layout = layout::ColumnMajorTileInterleave; + static constexpr int ElementsPerAccess = 128 / cutlass::sizeof_bits::value; + using Operator = cutlass::arch::OpMultiplyAddDequantizeInterleavedBToA; +}; + } // namespace kernel } // namespace gemm } // namespace cutlass diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h index ecfebd7a258d9..7ab652ab66863 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h @@ -180,7 +180,8 @@ struct MoeFCGemm { int64_t const* total_tokens_including_expert, int64_t gemm_n, int64_t gemm_k, GemmCoord* host_problem_sizes = nullptr) : problem_count(problem_count), threadblock_count(threadblock_count), group_size(group_size), output_op(output_op), ptr_A(const_cast(ptr_A)), ptr_B(const_cast(ptr_B)), weight_scales(const_cast(weight_scales)), weight_zeros(const_cast(weight_zeros)), ptr_C(const_cast(ptr_C)), C_is_broadcast{C_is_broadcast}, ptr_D(ptr_D), total_tokens_including_expert(total_tokens_including_expert), gemm_n(gemm_n), gemm_k(gemm_k), host_problem_sizes(nullptr) { - if (platform::is_same::value || platform::is_same::value) { + if (platform::is_same::value || platform::is_same::value || + platform::is_same::value) { assert(weight_scales); } this->gather_A_indices = nullptr; @@ -283,7 +284,8 @@ struct MoeFCGemm { } static Status can_implement(Arguments const& args) { - if constexpr (platform::is_same::value || platform::is_same::value) { + if constexpr (platform::is_same::value || platform::is_same::value || + platform::is_same::value) { if (args.weight_scales == nullptr) { CUTLASS_TRACE_HOST("MoeFCGemm::can_implement() - weight scales are required for uint8_t and uint4b_t"); return Status::kInvalid; diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h index 8d73329ed7713..86f38300d4b27 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h @@ -118,8 +118,9 @@ struct DqMma::value, "Mma multistage must dequantize after ldsm"); - static_assert(platform::is_same::value || platform::is_same::value, - "Element B must be uint8 or uint4"); + static_assert(platform::is_same::value || platform::is_same::value || + platform::is_same::value, + "Element B must be uint8, uint4 or float_e2m1"); static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits::value * kAlignmentA) == 128) ? cutlass::arch::CacheOperation::Global @@ -218,8 +219,9 @@ struct DqMma::value, "Mma multistage must dequantize after ldsm"); - static_assert(platform::is_same::value || platform::is_same::value, - "Element B must be uint8 or uint4"); + static_assert(platform::is_same::value || platform::is_same::value || + platform::is_same::value, + "Element B must be uint8, uint4 or float_e2m1"); static cutlass::arch::CacheOperation::Kind const CacheOpA = ((sizeof_bits::value * kAlignmentA) == 128) ? cutlass::arch::CacheOperation::Global diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h index ae0cee20d3575..14b59fd783331 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h @@ -116,8 +116,9 @@ struct DqMma::value || platform::is_same::value, "Element A must be fp16 or bf16"); - static_assert(platform::is_same::value || platform::is_same::value, - "Element B must be uint8 or uint4"); + static_assert(platform::is_same::value || platform::is_same::value || + platform::is_same::value, + "Element B must be uint8, uint4 or float_e2m1"); using OperatorInfo = arch::DetagOperator; using Operator = typename OperatorInfo::Operator; @@ -201,8 +202,9 @@ struct DqMma::value || platform::is_same::value, "Element A must be fp16 or bf16"); - static_assert(platform::is_same::value || platform::is_same::value, - "Element B must be uint8 or uint4"); + static_assert(platform::is_same::value || platform::is_same::value || + platform::is_same::value, + "Element B must be uint8, uint4 or float_e2m1"); using OperatorInfo = arch::DetagOperator; using Operator = typename OperatorInfo::Operator; diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h index dfe99c271f547..8f95212619437 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h @@ -221,6 +221,91 @@ struct DefaultMma +struct DefaultMma { + private: + static constexpr int kAlignmentScale = 128 / sizeof_bits::value; + + using Mma = DqMma; + + public: + using MmaCore = typename Mma::MmaCore; + using IteratorA = typename Mma::IteratorA; + using IteratorB = typename Mma::IteratorB; + using ThreadblockMma = typename Mma::ThreadblockMma; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Specialization for row-major output (OperatorClass TensorOp), fp16 activation & MXFP4 (e2m1) weight, +/// mma multistage (stage>=3) +template < + /// Layout type for A matrix operand + typename LayoutA, + /// Access granularity of A matrix in units of elements + int kAlignmentA, + /// Layout type for B matrix operand + typename LayoutB, + /// Access granularity of B matrix in units of elements + int kAlignmentB, + /// Element type for internal accumulation + typename ElementAccumulator, + /// Tag indicating architecture to tune for + typename ArchTag, + /// Threadblock-level tile size (concept: GemmShape) + typename ThreadblockShape, + /// Warp-level tile size (concept: GemmShape) + typename WarpShape, + /// Instruction-level tile size (concept: GemmShape) + typename InstructionShape, + /// Operation performed by GEMM + typename Operator, + /// + int kStages, + /// Shared memory clear option + SharedMemoryClearOption SharedMemoryClear> +struct DefaultMma { + private: + static constexpr int kAlignmentScale = 128 / sizeof_bits::value; + + using Mma = DqMma; + + public: + using MmaCore = typename Mma::MmaCore; + using IteratorA = typename Mma::IteratorA; + using IteratorB = typename Mma::IteratorB; + using ThreadblockMma = typename Mma::ThreadblockMma; +}; + #ifdef ENABLE_FP8 //////////////////////////////////////////////////////////////////////////////// /// Specialization for row-major output (OperatorClass TensorOp), fp8 activation & int4 weight, mma multistage diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h index cb5ce0f72b362..ba40045dfdb9b 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h @@ -331,6 +331,52 @@ struct DefaultMma +struct DefaultMma { + private: + static constexpr int kAlignmentScale = 128 / sizeof_bits::value; + + using Mma = DqMma; + + public: + using MmaCore = typename Mma::MmaCore; + using IteratorA = typename Mma::IteratorA; + using IteratorB = typename Mma::IteratorB; + using ThreadblockMma = typename Mma::ThreadblockMma; +}; + +//////////////////////////////////////////////////////////////////////////////// +/// Specialization for row-major output (OperatorClass TensorOp), bf16 activation & MXFP4 (e2m1) weight, +/// mma multistage (stage>=3) +template < + typename LayoutA, int kAlignmentA, typename LayoutB, int kAlignmentB, typename ElementAccumulator, + typename ArchTag, typename ThreadblockShape, typename WarpShape, typename InstructionShape, typename Operator, + int kStages, SharedMemoryClearOption SharedMemoryClear> +struct DefaultMma { + private: + static constexpr int kAlignmentScale = 128 / sizeof_bits::value; + + using Mma = DqMma; + + public: + using MmaCore = typename Mma::MmaCore; + using IteratorA = typename Mma::IteratorA; + using IteratorB = typename Mma::IteratorB; + using ThreadblockMma = typename Mma::ThreadblockMma; +}; + } // namespace threadblock } // namespace gemm } // namespace cutlass diff --git a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h index 86c45a865954e..788b7dc701ced 100644 --- a/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h +++ b/onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h @@ -394,6 +394,187 @@ struct FastInterleavedAndBiasedNumericArrayConverter { ///////////////////////////////////////////////////////////////////////////////////////////////// +// FP4 (e2m1) -> half/bf16 fast converters. +// +// These mirror the uint4b_t specializations above so the SM80 fused-dequant grouped GEMM +// (DqMmaMultistage) can target FP4 weights, exactly as it does INT4 weights today. Two +// differences from the integer path: +// 1. No zero-point bias. INT4 packs with a +8 bias that the converter subtracts. e2m1 codes +// are raw 4-bit floats, so the matching weight pre-pack interleaves WITHOUT adding a bias. +// 2. Conversion is a bit-field expand, not an integer-to-float magic-add. e2m1 = [s e1 e0 m0] +// with exponent bias 1; the 8 magnitudes are {0, .5, 1, 1.5, 2, 3, 4, 6}. +// +// The de-interleave order matches add_bias_and_interleave_int4s_inplace_kernel's permutation +// [e0,e2,e4,e6,e1,e3,e5,e7] (sans bias), so the same weight-prepack interleave is reused. +template <> +struct FastInterleavedAndBiasedNumericArrayConverter { + using result_type = Array; + using source_type = Array; + + // Convert a single 4-bit e2m1 code (0..15) to the raw 16-bit IEEE half bit pattern. + CUTLASS_HOST_DEVICE + static uint32_t e2m1_to_half_bits(uint32_t v) { + uint32_t sign = (v & 0x8u) << 12; // e2m1 sign bit -> half bit 15 + uint32_t e = (v >> 1) & 0x3u; // 2-bit exponent (bias 1) + uint32_t m = v & 0x1u; // 1-bit mantissa + // Normal (e != 0): half exponent = e - 1 + 15 = e + 14; mantissa bit -> half mantissa MSB. + // e == 0: subnormal 0.5 (m == 1) -> 0x3800, or zero (m == 0) -> 0x0000. + uint32_t mag = (e != 0u) ? (((14u + e) << 10) | (m ? 0x200u : 0u)) + : (m ? 0x3800u : 0u); + return sign | mag; + } + + CUTLASS_DEVICE + static result_type convert(source_type const& source) { + result_type result; + uint16_t* r = reinterpret_cast(&result); + uint32_t const packed = reinterpret_cast(source); + + uint32_t d[8]; + CUTLASS_PRAGMA_UNROLL + for (int k = 0; k < 8; ++k) { + d[k] = (packed >> (4 * k)) & 0xFu; + } + + // Invert the [e0,e2,e4,e6,e1,e3,e5,e7] interleave so result holds logical order e0..e7. + r[0] = static_cast(e2m1_to_half_bits(d[0])); + r[1] = static_cast(e2m1_to_half_bits(d[4])); + r[2] = static_cast(e2m1_to_half_bits(d[1])); + r[3] = static_cast(e2m1_to_half_bits(d[5])); + r[4] = static_cast(e2m1_to_half_bits(d[2])); + r[5] = static_cast(e2m1_to_half_bits(d[6])); + r[6] = static_cast(e2m1_to_half_bits(d[3])); + r[7] = static_cast(e2m1_to_half_bits(d[7])); + return result; + } + + CUTLASS_DEVICE + result_type operator()(source_type const& s) { + return convert(s); + } +}; + +template +struct FastInterleavedAndBiasedNumericArrayConverter { + static constexpr int VEC_WIDTH = 8; + static_assert(!(N % VEC_WIDTH), "N must be multiple of 8."); + + using result_type = Array; + using source_type = Array; + + CUTLASS_DEVICE + static result_type convert(source_type const& source) { + using scalar_result_type = typename result_type::Element; + using scalar_source_type = typename source_type::Element; + FastInterleavedAndBiasedNumericArrayConverter + convert_vector_; + + result_type result; + using vec_result = Array; + using vec_source = Array; + + vec_result* result_ptr = reinterpret_cast(&result); + vec_source const* source_ptr = reinterpret_cast(&source); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / VEC_WIDTH; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + CUTLASS_DEVICE + result_type operator()(source_type const& s) { + return convert(s); + } +}; + +template <> +struct FastInterleavedAndBiasedNumericArrayConverter { + using result_type = Array; + using source_type = Array; + + // Convert a single 4-bit e2m1 code (0..15) to the raw 16-bit bfloat16 bit pattern. + CUTLASS_HOST_DEVICE + static uint32_t e2m1_to_bf16_bits(uint32_t v) { + uint32_t sign = (v & 0x8u) << 12; // e2m1 sign bit -> bf16 bit 15 + uint32_t e = (v >> 1) & 0x3u; // 2-bit exponent (bias 1) + uint32_t m = v & 0x1u; // 1-bit mantissa + // Normal (e != 0): bf16 exponent = e - 1 + 127 = e + 126; mantissa bit -> bf16 mantissa MSB. + // e == 0: subnormal 0.5 (m == 1) -> 0x3F00, or zero (m == 0) -> 0x0000. + uint32_t mag = (e != 0u) ? (((126u + e) << 7) | (m ? 0x40u : 0u)) + : (m ? 0x3F00u : 0u); + return sign | mag; + } + + CUTLASS_DEVICE + static result_type convert(source_type const& source) { + result_type result; + uint16_t* r = reinterpret_cast(&result); + uint32_t const packed = reinterpret_cast(source); + + uint32_t d[8]; + CUTLASS_PRAGMA_UNROLL + for (int k = 0; k < 8; ++k) { + d[k] = (packed >> (4 * k)) & 0xFu; + } + + r[0] = static_cast(e2m1_to_bf16_bits(d[0])); + r[1] = static_cast(e2m1_to_bf16_bits(d[4])); + r[2] = static_cast(e2m1_to_bf16_bits(d[1])); + r[3] = static_cast(e2m1_to_bf16_bits(d[5])); + r[4] = static_cast(e2m1_to_bf16_bits(d[2])); + r[5] = static_cast(e2m1_to_bf16_bits(d[6])); + r[6] = static_cast(e2m1_to_bf16_bits(d[3])); + r[7] = static_cast(e2m1_to_bf16_bits(d[7])); + return result; + } + + CUTLASS_DEVICE + result_type operator()(source_type const& s) { + return convert(s); + } +}; + +template +struct FastInterleavedAndBiasedNumericArrayConverter { + static constexpr int VEC_WIDTH = 8; + static_assert(!(N % VEC_WIDTH), "N must be multiple of 8."); + + using result_type = Array; + using source_type = Array; + + CUTLASS_DEVICE + static result_type convert(source_type const& source) { + using scalar_result_type = typename result_type::Element; + using scalar_source_type = typename source_type::Element; + FastInterleavedAndBiasedNumericArrayConverter + convert_vector_; + + result_type result; + using vec_result = Array; + using vec_source = Array; + + vec_result* result_ptr = reinterpret_cast(&result); + vec_source const* source_ptr = reinterpret_cast(&source); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < N / VEC_WIDTH; ++i) { + result_ptr[i] = convert_vector_(source_ptr[i]); + } + + return result; + } + + CUTLASS_DEVICE + result_type operator()(source_type const& s) { + return convert(s); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + } // namespace cutlass ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h index c3b734816cf84..b1af64b1e16a4 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h @@ -33,6 +33,17 @@ enum class QuantType { int get_arch_for_mixed_gemm_weight_preprocess(int arch); +// ``apply_bias_interleave`` controls the final integer-only "+bias + pair-interleave" step +// (step 4). It MUST be left true for genuine signed-integer weights (INT4/INT8). For MXFP4 +// (e2m1) codes the nibbles are floating-point, so the integer +8 bias would corrupt them; +// FP4 callers pass false to skip step 4 while keeping the layout-only steps 1-3 (row-permute, +// subbyte-transpose, column-interleave), which apply to e2m1 unchanged. +// +// ``interleave_without_bias`` (only consulted when ``apply_bias_interleave`` is false) applies +// step 4's [e0,e2,e4,e6,e1,e3,e5,e7] nibble pair-interleave WITHOUT the +8 bias. This is the +// layout the SM80 MoE grouped GEMM's e2m1 dequant converter expects (it inverts that +// permutation). The fused MXFP4 GEMV decode kernel uses a different layout and leaves this +// false. void preprocess_weights_for_mixed_gemm_cuda(cudaStream_t stream, int arch, int8_t* preprocessed_quantized_weight, @@ -40,7 +51,9 @@ void preprocess_weights_for_mixed_gemm_cuda(cudaStream_t stream, int32_t* d_permutation_map, std::vector const& shape, QuantType quant_type, - bool synchronize = true); + bool synchronize = true, + bool apply_bias_interleave = true, + bool interleave_without_bias = false); } // namespace weight_only } // namespace kernels diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu index 7e83bdda72eab..e5823c88d3d7f 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu @@ -482,6 +482,49 @@ __global__ void add_bias_and_interleave_int4s_inplace_kernel(uint32_t* tensor, s } } +// Interleave-only variant for MXFP4 (e2m1) weights: applies the SAME +// [e0,e2,e4,e6,e1,e3,e5,e7] nibble pair-interleave as +// add_bias_and_interleave_int4s_inplace_kernel, but writes the raw 4-bit code unchanged (no +// signed +8 bias). The e2m1 dequant converter +// (FastInterleavedAndBiasedNumericArrayConverter) inverts exactly this +// permutation, so the grouped GEMM requires the interleave while the floating-point nibbles must +// stay untouched. +__global__ void interleave_int4s_inplace_kernel(uint32_t* tensor, size_t num_elts) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t register_idx = static_cast(idx); + if (register_idx < num_elts / 8) { + uint32_t current_register = tensor[register_idx]; + uint32_t transformed_register = 0; + for (int i = 0; i < 8; ++i) { + uint8_t raw_nibble = (current_register >> (i * 4)) & 0x0F; + int dest_idx = ((i % 2) == 0) ? (i / 2) : ((i - 1) / 2 + 4); + transformed_register |= (static_cast(raw_nibble & 0x0F) << (dest_idx * 4)); + } + tensor[register_idx] = transformed_register; + } +} + +void interleave_without_bias_quantized_tensor_inplace_cuda( + int8_t* tensor, + size_t num_elts, + QuantType quant_type, + cudaStream_t stream) { + ORT_ENFORCE(quant_type == QuantType::W4_A16 || quant_type == QuantType::W4_AFP8, + "Interleave-without-bias is only supported for 4-bit (e2m1) weights."); + if (num_elts == 0) { + return; + } + ORT_ENFORCE(num_elts >= 8 && num_elts % 8 == 0, + "Interleave-without-bias requires the number of 4-bit elements to be a non-zero multiple of 8, got ", + num_elts, "."); + const int threads_per_block = 256; + const int num_registers = SafeInt(num_elts) / 8; + const int num_blocks = (num_registers + threads_per_block - 1) / threads_per_block; + interleave_int4s_inplace_kernel<<>>( + reinterpret_cast(tensor), + num_elts); +} + /** * @brief Launches the CUDA kernel for in-place bias addition and interleaving. * @@ -541,7 +584,9 @@ void preprocess_weights_for_mixed_gemm_cuda(cudaStream_t stream, int32_t* d_permutation_map, std::vector const& shape, QuantType quant_type, - bool synchronize) { + bool synchronize, + bool apply_bias_interleave, + bool interleave_without_bias) { LayoutDetails details = getLayoutDetailsForTransform(quant_type, arch); ORT_ENFORCE(shape.size() == 2 || shape.size() == 3, "Shape must be 2-D or 3-D"); @@ -578,11 +623,24 @@ void preprocess_weights_for_mixed_gemm_cuda(cudaStream_t stream, std::swap(src_buf, dst_buf); } - add_bias_and_interleave_quantized_tensor_inplace_cuda( - src_buf, - num_elts, - quant_type, - stream); + // Step 4 (integer-only): add the +bias and pair-interleave. Skipped for MXFP4 (e2m1) + // float codes, which would be corrupted by the integer bias; the preceding layout-only + // steps already produced the ColumnMajorInterleaved nibble order the GEMV consumes. + if (apply_bias_interleave) { + add_bias_and_interleave_quantized_tensor_inplace_cuda( + src_buf, + num_elts, + quant_type, + stream); + } else if (interleave_without_bias) { + // MXFP4 grouped-GEMM layout: apply step 4's [e0,e2,e4,e6,e1,e3,e5,e7] nibble interleave + // (which the e2m1 dequant converter inverts) WITHOUT the integer +8 bias. + interleave_without_bias_quantized_tensor_inplace_cuda( + src_buf, + num_elts, + quant_type, + stream); + } if (preprocessed_quantized_weight != src_buf) { const size_t num_bytes = num_elts * static_cast(get_weight_quant_bits(quant_type)) / static_cast(8); diff --git a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h index 4fa64ef329c57..2504f12b57776 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h @@ -94,6 +94,15 @@ struct Int4DetailsW { static constexpr int kElemBits = 4; }; +struct Fp4DetailsW { + static constexpr int kElemBits = 4; +}; + +template +struct IsFp4Weight : std::false_type {}; +template <> +struct IsFp4Weight : std::true_type {}; + template struct ColumnMajor { using DetailsA = TypeDetailsA; @@ -223,12 +232,66 @@ struct I2FConverter { } }; +template +struct Fp4I2FConverter { + static_assert(std::is_same_v || std::is_same_v); + + // Branchless E2M1 (FP4) -> half/bf16 decode. E2M1 has only eight magnitudes + // {0, 0.5, 1, 1.5, 2, 3, 4, 6}, so the 16-bit float bit pattern is built directly with a + // single prmt.b32 byte-select from packed magnitude constants (two for bf16, whose low byte + // is not always zero) plus a shifted sign bit. This replaces the per-element float LUT + // lookup + sign branch + float->AType conversion, which profiling showed to be the dominant + // ALU-pipeline cost of the small-decode FP4 GEMV (ncu: ALU ~79% of a compute-bound kernel). + // The magnitude bytes below are the exact half/bf16 encodings of the eight FP4 values, so the + // result is bit-identical to the previous LUT path. + __device__ __forceinline__ static AType decode(uint8_t code) { +#if defined(__CUDA_ARCH__) + uint32_t const sel = code & 0x7u; + uint32_t const sign = static_cast(code & 0x8u) << 12; // FP4 sign bit -> bit 15 + if constexpr (std::is_same_v) { + // half high byte per magnitude (low byte is always 0): + // codes 0..3 -> {0x00, 0x38, 0x3C, 0x3E}, codes 4..7 -> {0x40, 0x42, 0x44, 0x46}. + uint32_t hb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3E3C3800u), "r"(0x46444240u), "r"(sel)); + return __ushort_as_half(static_cast(((hb & 0xFFu) << 8) | sign)); + } else { + // bf16 high byte {0x00,0x3F,0x3F,0x3F, 0x40,0x40,0x40,0x40} and + // low byte {0x00,0x00,0x80,0xC0, 0x00,0x40,0x80,0xC0}. + uint32_t hb, lb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3F3F3F00u), "r"(0x40404040u), "r"(sel)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lb) : "r"(0xC0800000u), "r"(0xC0804000u), "r"(sel)); + return __ushort_as_bfloat16(static_cast(((hb & 0xFFu) << 8) | (lb & 0xFFu) | sign)); + } +#else + constexpr float kValues[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f}; + float v = kValues[code & 0x7]; + return static_cast((code & 0x8) ? -v : v); +#endif + } + + template + __device__ __forceinline__ static void convert(void* src, void* dst) { + static_assert(N % 2 == 0); + uint8_t const* s = reinterpret_cast(src); + AType* d = reinterpret_cast(dst); +#pragma unroll + for (int i = 0; i < N; i += 2) { + uint8_t byte = s[i >> 1]; + d[i] = decode(static_cast(byte & 0x0F)); + d[i + 1] = decode(static_cast((byte >> 4) & 0x0F)); + } + } +}; + template struct ConverterWrapper { using TypeDetailsA = typename Details::TypeDetailsA; using TypeDetailsW = typename Details::TypeDetailsW; static constexpr bool kUseInterleavedConverter = Details::kUseInterleavedConverter; - using Converter = I2FConverter; + using Converter = std::conditional_t< + IsFp4Weight::value, + Fp4I2FConverter, + I2FConverter>; }; template diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl index ba23174d3c203..70badcbb41087 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl @@ -14,10 +14,10 @@ * limitations under the License. */ -#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#ifdef __GNUC__ // Check if the compiler is GCC or Clang #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wstrict-aliasing" -#endif // __GNUC__ +#endif // __GNUC__ #include "cutlass/epilogue/collective/default_epilogue.hpp" #include "cutlass/epilogue/thread/linear_combination.h" @@ -50,9 +50,9 @@ #include "contrib_ops/cuda/llm/cutlass_extensions/gemm/collective/collective_builder_mixed_input.hpp" #include "contrib_ops/cuda/llm/cutlass_extensions/gemm_configs.h" -#ifdef __GNUC__ // Check if the compiler is GCC or Clang +#ifdef __GNUC__ // Check if the compiler is GCC or Clang #pragma GCC diagnostic pop -#endif // __GNUC__ +#endif // __GNUC__ #include "core/common/common.h" #include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" @@ -61,12 +61,9 @@ #include "contrib_ops/cuda/llm/cutlass_type_conversion.h" #include "contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.h" -namespace onnxruntime::llm -{ -namespace kernels -{ -namespace cutlass_kernels -{ +namespace onnxruntime::llm { +namespace kernels { +namespace cutlass_kernels { namespace tk = onnxruntime::llm::common; namespace tkc = onnxruntime::llm::cutlass_extensions; @@ -75,238 +72,225 @@ using EpilogueFusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion; using namespace cute; template + typename CTAShape, typename ClusterShape, typename MainloopScheduleType, typename EpilogueScheduleType, + cutlass::WeightOnlyQuantOp QuantOp> void sm90_generic_mixed_moe_gemm_kernelLauncher(GroupedGemmInput inputs, - TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) -{ - ORT_LLM_LOG_ENTRY(); - - ///////////////////////////////////////////////////////////////////////////////////////////////// - /// GEMM kernel configurations - ///////////////////////////////////////////////////////////////////////////////////////////////// - static_assert(FUSION == EpilogueFusion::NONE || FUSION == EpilogueFusion::FINALIZE, - "Unimplemented fusion provided to TMA WS Mixed MoE gemm launcher"); - constexpr static bool IsFinalizeFusion = FUSION == EpilogueFusion::FINALIZE; - - ///////////////////////////////////////////////////////////////////////////////////////////////// - /// GEMM kernel configurations - ///////////////////////////////////////////////////////////////////////////////////////////////// - - // A matrix configuration - using ElementA = typename CudaToCutlassTypeAdapter::type; - using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand - constexpr int AlignmentA - = 128 / cutlass::sizeof_bits::value; // Alignment of A matrix in units of elements (up to 16 bytes) - - // B matrix configuration - using ElementB_ = typename CudaToCutlassTypeAdapter::type; - using ElementB = std::conditional_t, cutlass::int4b_t, ElementB_>; - using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand - constexpr int AlignmentB - = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of - // elements (up to 16 bytes) - - // This example manually swaps and transposes, so keep transpose of input layouts - using LayoutA_Transpose = typename cutlass::layout::LayoutTranspose::type; - using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose::type; - - // Need to pass a pointer type to make the 3rd dimension of Stride be _0 - using StrideA = cute::remove_pointer_t>; - using StrideB = cute::remove_pointer_t>; - - // Scale configuration - constexpr bool use_wfp4a16 = std::is_same_v; - constexpr int group_size = use_wfp4a16 ? cutlass::gemm::collective::detail::mxfp4_group_size - : cutlass::gemm::collective::detail::int4_group_size; - constexpr int PackedScalesNum = get<2>(CTAShape{}) / group_size; - using ElementScale = std::conditional_t; - using ElementScalePacked = cutlass::Array; - using LayoutScale = cutlass::layout::RowMajor; - - // C/D matrix configuration - using ElementC = typename CudaToCutlassTypeAdapter::type; - using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands - constexpr int AlignmentC - = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of - // elements (up to 16 bytes) - - // D matrix configuration - using ElementD = ElementC; - using LayoutD = LayoutC; - constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; - - using ElementFinalOutput = ElementC; - using ElementBias = ElementFinalOutput; - using ElementRouterScales = float; - - // Core kernel configurations - using ElementAccumulator = float; // Element type for internal accumulation - using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature - using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag - using TileShape = CTAShape; // Threadblock-level tile size - using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size - using KernelSchedule - = std::conditional_t, - cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong, - cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative>; - using EpilogueSchedule - = std::conditional_t, - cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong, - cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>; // Epilogue to launch - - using StrideC = TmaWarpSpecializedGroupedGemmInput::StrideC; - - // Default epilogue (NONE fusion) - using CollectiveEpilogueDefault = typename cutlass::epilogue::collective::CollectiveBuilder::type*, - AlignmentC, ElementD, typename cutlass::layout::LayoutTranspose::type*, AlignmentD, - EpilogueSchedule>::CollectiveOp; - - // Fused finalize epilogue (FINALIZE fusion) - using CollectiveEpilogueFinalize = - typename cutlass::epilogue::collective::EpilogueMoeFusedFinalizeBuilder< - ArchTag, TileShape, - ElementC, StrideC*, - ElementFinalOutput, - TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideFinalOutput, - ElementAccumulator, - ElementAccumulator, - ElementBias, TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideBias, - ElementRouterScales, - TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideRouterScales - >::CollectiveOp; - - using CollectiveEpilogue = std::conditional_t; - - // =========================================================== MIXED INPUT WITH SCALES - // =========================================================================== The Scale information must get paired - // with the operand that will be scaled. In this example, B is scaled so we make a tuple of B's information and the - // scale information. - using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilderMixedInput, LayoutB_Transpose*, AlignmentB, ElementA, LayoutA_Transpose*, - AlignmentA, ElementAccumulator, TileShape, ClusterShape, - cutlass::gemm::collective::StageCountAutoCarveout( - sizeof(typename CollectiveEpilogue::SharedStorage))>, - KernelSchedule>::CollectiveOp; - - using GemmKernel = cutlass::gemm::kernel::GemmUniversal>, - CollectiveMainloop, CollectiveEpilogue>; - - using GemmGrouped = cutlass::gemm::device::GemmUniversalAdapter; - using StrideD = typename GemmKernel::InternalStrideD; - using StrideS = typename CollectiveMainloop::StrideScale; - - GemmGrouped gemm; - using Args = typename GemmGrouped::Arguments; - - cutlass::KernelHardwareInfo hw_info; - hw_info.device_id = 0; - hw_info.sm_count = sm_count_; - - using EpilogueArguments = typename CollectiveEpilogue::Arguments; - using EpilogueScalars = decltype(EpilogueArguments{}.thread); - - auto make_epilogue_scalars = [&]() -> EpilogueScalars { - if constexpr (IsFinalizeFusion) { - return EpilogueScalars{ElementAccumulator(1.f), ElementAccumulator(0.f)}; - } else { - EpilogueScalars scalars; - scalars.alpha = use_wfp4a16 ? 1 : 0; - scalars.beta = 0; - scalars.alpha_ptr = nullptr; - scalars.beta_ptr = nullptr; - scalars.alpha_ptr_array = use_wfp4a16 ? nullptr : inputs.alpha_scales; - scalars.beta_ptr_array = nullptr; - // One alpha and beta per each group - scalars.dAlpha = {cute::_0{}, cute::_0{}, use_wfp4a16 ? 0 : 1}; - scalars.dBeta = {cute::_0{}, cute::_0{}, use_wfp4a16 ? 0 : 1}; - return scalars; - } - }; - - auto make_epilogue_args = [&]() -> EpilogueArguments { - auto scalars = make_epilogue_scalars(); - if constexpr (IsFinalizeFusion) { - auto epi_params = hopper_inputs.fused_finalize_epilogue; - return EpilogueArguments{ - scalars, - nullptr, hopper_inputs.stride_c, // C params - reinterpret_cast(epi_params.ptr_final_output), - epi_params.stride_final_output, // D (output) params - reinterpret_cast(epi_params.ptr_bias), - epi_params.stride_bias, // Bias params - epi_params.ptr_router_scales, epi_params.stride_router_scales, // Router scales - epi_params.ptr_expert_first_token_offset, // Offset of this expert's token in the router scales - epi_params.ptr_source_token_index, // Index of the source token to sum into - epi_params.num_rows_in_final_output // Number of tokens in the output buffer - }; - } else { - return EpilogueArguments{scalars, - reinterpret_cast(hopper_inputs.ptr_c), hopper_inputs.stride_c, - reinterpret_cast(hopper_inputs.default_epilogue.ptr_d), - hopper_inputs.default_epilogue.stride_d}; - } - }; - - if (workspace_size != nullptr) - { - const Args args{cutlass::gemm::GemmUniversalMode::kGrouped, - {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, - {reinterpret_cast(hopper_inputs.ptr_b), hopper_inputs.stride_b, - reinterpret_cast(hopper_inputs.ptr_a), hopper_inputs.stride_a, - reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), - hopper_inputs.int4_groupwise_params.stride_s_a, group_size}, - make_epilogue_args(), - hw_info}; - *workspace_size = gemm.get_workspace_size(args); - return; + TmaWarpSpecializedGroupedGemmInput hopper_inputs, int sm_count_, size_t* workspace_size) { + ORT_LLM_LOG_ENTRY(); + + ///////////////////////////////////////////////////////////////////////////////////////////////// + /// GEMM kernel configurations + ///////////////////////////////////////////////////////////////////////////////////////////////// + static_assert(FUSION == EpilogueFusion::NONE || FUSION == EpilogueFusion::FINALIZE, + "Unimplemented fusion provided to TMA WS Mixed MoE gemm launcher"); + constexpr static bool IsFinalizeFusion = FUSION == EpilogueFusion::FINALIZE; + + ///////////////////////////////////////////////////////////////////////////////////////////////// + /// GEMM kernel configurations + ///////////////////////////////////////////////////////////////////////////////////////////////// + + // A matrix configuration + using ElementA = typename CudaToCutlassTypeAdapter::type; + using LayoutA = cutlass::layout::RowMajor; // Layout type for A matrix operand + constexpr int AlignmentA = 128 / cutlass::sizeof_bits::value; // Alignment of A matrix in units of elements (up to 16 bytes) + + // B matrix configuration + using ElementB_ = typename CudaToCutlassTypeAdapter::type; + using ElementB = std::conditional_t, cutlass::int4b_t, ElementB_>; + using LayoutB = cutlass::layout::ColumnMajor; // Layout type for B matrix operand + constexpr int AlignmentB = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of B matrix in units of + // elements (up to 16 bytes) + + // This example manually swaps and transposes, so keep transpose of input layouts + using LayoutA_Transpose = typename cutlass::layout::LayoutTranspose::type; + using LayoutB_Transpose = typename cutlass::layout::LayoutTranspose::type; + + // Need to pass a pointer type to make the 3rd dimension of Stride be _0 + using StrideA = cute::remove_pointer_t>; + using StrideB = cute::remove_pointer_t>; + + // Scale configuration + constexpr bool use_wfp4a16 = std::is_same_v; + constexpr int group_size = use_wfp4a16 ? cutlass::gemm::collective::detail::mxfp4_group_size + : cutlass::gemm::collective::detail::int4_group_size; + constexpr int PackedScalesNum = get<2>(CTAShape{}) / group_size; + using ElementScale = std::conditional_t; + using ElementScalePacked = cutlass::Array; + using LayoutScale = cutlass::layout::RowMajor; + + // C/D matrix configuration + using ElementC = typename CudaToCutlassTypeAdapter::type; + using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands + constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; // Memory access granularity/alignment of C matrix in units of + // elements (up to 16 bytes) + + // D matrix configuration + using ElementD = ElementC; + using LayoutD = LayoutC; + constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; + + using ElementFinalOutput = ElementC; + using ElementBias = ElementFinalOutput; + using ElementRouterScales = float; + + // Core kernel configurations + using ElementAccumulator = float; // Element type for internal accumulation + using ArchTag = cutlass::arch::Sm90; // Tag indicating the minimum SM that supports the intended feature + using OperatorClass = cutlass::arch::OpClassTensorOp; // Operator class tag + using TileShape = CTAShape; // Threadblock-level tile size + using StageCountType = cutlass::gemm::collective::StageCountAuto; // Stage count maximized based on the tile size + using KernelSchedule = std::conditional_t, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong, + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperative>; + using EpilogueSchedule = std::conditional_t, + cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong, + cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative>; // Epilogue to launch + + using StrideC = TmaWarpSpecializedGroupedGemmInput::StrideC; + + // Default epilogue (NONE fusion) + using CollectiveEpilogueDefault = typename cutlass::epilogue::collective::CollectiveBuilder::type*, + AlignmentC, ElementD, typename cutlass::layout::LayoutTranspose::type*, AlignmentD, + EpilogueSchedule>::CollectiveOp; + + // Fused finalize epilogue (FINALIZE fusion) + using CollectiveEpilogueFinalize = + typename cutlass::epilogue::collective::EpilogueMoeFusedFinalizeBuilder< + ArchTag, TileShape, + ElementC, StrideC*, + ElementFinalOutput, + TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideFinalOutput, + ElementAccumulator, + ElementAccumulator, + ElementBias, TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideBias, + ElementRouterScales, + TmaWarpSpecializedGroupedGemmInput::FusedFinalizeEpilogue::StrideRouterScales>::CollectiveOp; + + using CollectiveEpilogue = std::conditional_t; + + // =========================================================== MIXED INPUT WITH SCALES + // =========================================================================== The Scale information must get paired + // with the operand that will be scaled. In this example, B is scaled so we make a tuple of B's information and the + // scale information. + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilderMixedInput, LayoutB_Transpose*, AlignmentB, ElementA, LayoutA_Transpose*, + AlignmentA, ElementAccumulator, TileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + KernelSchedule>::CollectiveOp; + + using GemmKernel = cutlass::gemm::kernel::GemmUniversal>, + CollectiveMainloop, CollectiveEpilogue>; + + using GemmGrouped = cutlass::gemm::device::GemmUniversalAdapter; + using StrideD = typename GemmKernel::InternalStrideD; + using StrideS = typename CollectiveMainloop::StrideScale; + + GemmGrouped gemm; + using Args = typename GemmGrouped::Arguments; + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = 0; + hw_info.sm_count = sm_count_; + + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueScalars = decltype(EpilogueArguments{}.thread); + + auto make_epilogue_scalars = [&]() -> EpilogueScalars { + if constexpr (IsFinalizeFusion) { + return EpilogueScalars{ElementAccumulator(1.f), ElementAccumulator(0.f)}; + } else { + auto const* alpha_ptr_array = use_wfp4a16 ? hopper_inputs.alpha_scale_ptr_array : inputs.alpha_scales; + EpilogueScalars scalars; + scalars.alpha = alpha_ptr_array ? 0 : 1; + scalars.beta = 0; + scalars.alpha_ptr = nullptr; + scalars.beta_ptr = nullptr; + scalars.alpha_ptr_array = alpha_ptr_array; + scalars.beta_ptr_array = nullptr; + // One alpha and beta per each group + scalars.dAlpha = {cute::_0{}, cute::_0{}, alpha_ptr_array ? 1 : 0}; + scalars.dBeta = {cute::_0{}, cute::_0{}, 0}; + return scalars; } - - auto arguments = Args{cutlass::gemm::GemmUniversalMode::kGrouped, - {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, - {reinterpret_cast(hopper_inputs.ptr_b), hopper_inputs.stride_b, - reinterpret_cast(hopper_inputs.ptr_a), hopper_inputs.stride_a, - reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), - hopper_inputs.int4_groupwise_params.stride_s_a, group_size}, - make_epilogue_args(), - hw_info}; - - size_t const required_workspace = gemm.get_workspace_size(arguments); - ORT_ENFORCE(required_workspace <= hopper_inputs.gemm_workspace_size, - "[Mixed dtype WS grouped GEMM] given workspace size insufficient, ", hopper_inputs.gemm_workspace_size, - " < ", required_workspace); - - auto can_implement = gemm.can_implement(arguments); - if (can_implement != cutlass::Status::kSuccess) - { - std::string err_msg = "mixed dtype WS grouped cutlass kernel will fail for params. Error: " - + std::string(cutlassGetStatusString(can_implement)); - std::cout << err_msg << std::endl; - ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); - } - - auto init_status = gemm.initialize(arguments, hopper_inputs.gemm_workspace, inputs.stream); - if (init_status != cutlass::Status::kSuccess) - { - std::string err_msg = "Failed to initialize cutlass mixed dtype WS grouped gemm. Error: " - + std::string(cutlassGetStatusString(init_status)); - ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); - } - - auto run_status = gemm.run(inputs.stream); - if (run_status != cutlass::Status::kSuccess) - { - std::string err_msg = "Failed to run cutlass mixed dtype WS grouped gemm. Error: " - + std::string(cutlassGetStatusString(run_status)); - ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); + }; + + auto make_epilogue_args = [&]() -> EpilogueArguments { + auto scalars = make_epilogue_scalars(); + if constexpr (IsFinalizeFusion) { + auto epi_params = hopper_inputs.fused_finalize_epilogue; + return EpilogueArguments{ + scalars, + nullptr, hopper_inputs.stride_c, // C params + reinterpret_cast(epi_params.ptr_final_output), + epi_params.stride_final_output, // D (output) params + reinterpret_cast(epi_params.ptr_bias), + epi_params.stride_bias, // Bias params + epi_params.ptr_router_scales, epi_params.stride_router_scales, // Router scales + epi_params.ptr_expert_first_token_offset, // Offset of this expert's token in the router scales + epi_params.ptr_source_token_index, // Index of the source token to sum into + epi_params.num_rows_in_final_output // Number of tokens in the output buffer + }; + } else { + return EpilogueArguments{scalars, + reinterpret_cast(hopper_inputs.ptr_c), hopper_inputs.stride_c, + reinterpret_cast(hopper_inputs.default_epilogue.ptr_d), + hopper_inputs.default_epilogue.stride_d}; } + }; + + if (workspace_size != nullptr) { + const Args args{cutlass::gemm::GemmUniversalMode::kGrouped, + {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, + {reinterpret_cast(hopper_inputs.ptr_b), hopper_inputs.stride_b, + reinterpret_cast(hopper_inputs.ptr_a), hopper_inputs.stride_a, + reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), + hopper_inputs.int4_groupwise_params.stride_s_a, group_size}, + make_epilogue_args(), + hw_info}; + *workspace_size = gemm.get_workspace_size(args); return; + } + + auto arguments = Args{cutlass::gemm::GemmUniversalMode::kGrouped, + {inputs.num_experts, hopper_inputs.int4_groupwise_params.shape.problem_shapes, nullptr}, + {reinterpret_cast(hopper_inputs.ptr_b), hopper_inputs.stride_b, + reinterpret_cast(hopper_inputs.ptr_a), hopper_inputs.stride_a, + reinterpret_cast(hopper_inputs.int4_groupwise_params.ptr_s_a), + hopper_inputs.int4_groupwise_params.stride_s_a, group_size}, + make_epilogue_args(), + hw_info}; + + size_t const required_workspace = gemm.get_workspace_size(arguments); + ORT_ENFORCE(required_workspace <= hopper_inputs.gemm_workspace_size, + "[Mixed dtype WS grouped GEMM] given workspace size insufficient, ", hopper_inputs.gemm_workspace_size, + " < ", required_workspace); + + auto can_implement = gemm.can_implement(arguments); + if (can_implement != cutlass::Status::kSuccess) { + std::string err_msg = "mixed dtype WS grouped cutlass kernel will fail for params. Error: " + std::string(cutlassGetStatusString(can_implement)); + std::cout << err_msg << std::endl; + ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); + } + + auto init_status = gemm.initialize(arguments, hopper_inputs.gemm_workspace, inputs.stream); + if (init_status != cutlass::Status::kSuccess) { + std::string err_msg = "Failed to initialize cutlass mixed dtype WS grouped gemm. Error: " + std::string(cutlassGetStatusString(init_status)); + ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); + } + + auto run_status = gemm.run(inputs.stream); + if (run_status != cutlass::Status::kSuccess) { + std::string err_msg = "Failed to run cutlass mixed dtype WS grouped gemm. Error: " + std::string(cutlassGetStatusString(run_status)); + ORT_THROW("[Mixed dtype WS grouped GEMM] " + err_msg); + } + return; } -} // namespace cutlass_kernels -} // namespace kernels -} // namespace onnxruntime::llm +} // namespace cutlass_kernels +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h index 2a17f00191f8f..b6577dfcfff6a 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h @@ -201,8 +201,10 @@ struct TmaWarpSpecializedGroupedGemmInput { FpXBlockScalingType fpX_block_scaling_type = FpXBlockScalingType::NONE; struct INT4GroupwiseParams { - constexpr static int group_size = 128; // Unused, hard-coded to 128 + constexpr static int int4_group_size = 128; + constexpr static int wfp4a16_group_size = 32; bool enabled = false; + bool use_wfp4a16 = false; using SFA = __nv_bfloat16; using SFB = __nv_bfloat16; // Unused using ProblemShapeInt = cutlass::gemm::GroupProblemShape>; @@ -299,9 +301,16 @@ class MoeGemmRunner { TmaWarpSpecializedGroupedGemmInput hopper_inputs); std::vector getConfigs() const; - static std::vector getConfigs(int sm); - static std::vector getTmaWarpSpecializedConfigs(int sm); - static std::vector getAmpereConfigs(int sm); + static std::vector getConfigs(int sm, bool use_sm80_fp4 = false); + static std::vector getTmaWarpSpecializedConfigs(int sm, bool use_sm80_fp4 = false); + static std::vector getAmpereConfigs(int sm, bool use_sm80_fp4 = false); + + // Route wfp4a16 (MXFP4 weights + FP16/BF16 activations) through the SM80 fused-dequant grouped + // GEMM instead of the SM90 TMA WS path. The decision is made once in the QMoE op constructor + // (which reads ORT_FP4_SM80_GEMM / ORT_ENABLE_FP4_CUTLASS_GEMM at the right time) and pushed in + // here, so config selection at inference time does NOT depend on the live environment. + void setUseSm80Fp4(bool v) { use_sm80_fp4_ = v; } + [[nodiscard]] bool useSm80Fp4() const { return use_sm80_fp4_; } [[nodiscard]] bool isTmaWarpSpecialized(cutlass_extensions::CutlassGemmConfig gemm_config) const; [[nodiscard]] bool supportsTmaWarpSpecialized() const; @@ -325,6 +334,9 @@ class MoeGemmRunner { private: int sm_{}; int multi_processor_count_{}; + // wfp4a16 only: when true, config selection offers the SM80 fused-dequant Ampere grouped GEMM + // and suppresses the SM90 TMA WS configs. Set by the QMoE op after construction. + bool use_sm80_fp4_ = false; mutable int num_experts_ = 0; mutable size_t gemm_workspace_size_ = 0; size_t calcMaxWorkspaceSize(int num_experts) const; diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h index 7c63fb948c045..15fd5042fdc38 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h @@ -127,7 +127,11 @@ struct genericMoeGemmKernelLauncher { "Specialized for half, float"); #endif - static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value || cutlass::platform::is_same::value); + static_assert(cutlass::platform::is_same::value || cutlass::platform::is_same::value || cutlass::platform::is_same::value +#if defined(ENABLE_FP4) + || cutlass::platform::is_same::value +#endif + ); static_assert(arch::kMinComputeCapability < 90, "Sm90+ architecture should use specialized kernels"); @@ -371,7 +375,7 @@ static void dispatch(GroupedGemmInput MoeGemmRunner::getConfigs() const { ORT_LLM_LOG_ENTRY(); - return getConfigs(sm_); + return getConfigs(sm_, use_sm80_fp4_); +} + +// Whether wfp4a16 should use the SM80 fused-dequant grouped GEMM (vs the SM90 TMA WS path). +// The ``use_sm80_fp4`` flag is the decision captured by the QMoE op constructor and pushed into +// the runner via setUseSm80Fp4(); we only add the hard architectural guard that the SM80 path is +// for Ampere through pre-Blackwell. No environment is read here, so inference-time config selection +// is independent of the live environment. +inline bool moeUseSm80Fp4(int sm, bool use_sm80_fp4) { + return use_sm80_fp4 && sm >= 80 && sm < 120; } template std::vector MoeGemmRunner::getConfigs( - int sm) { + int sm, bool use_sm80_fp4) { ORT_LLM_LOG_ENTRY(); - std::vector candidate_configs = getTmaWarpSpecializedConfigs(sm); - std::vector ampere_configs = getAmpereConfigs(sm); + std::vector candidate_configs = getTmaWarpSpecializedConfigs(sm, use_sm80_fp4); + std::vector ampere_configs = getAmpereConfigs(sm, use_sm80_fp4); std::copy(ampere_configs.begin(), ampere_configs.end(), std::back_inserter(candidate_configs)); return candidate_configs; } template std::vector -MoeGemmRunner::getAmpereConfigs(int sm) { +MoeGemmRunner::getAmpereConfigs(int sm, bool use_sm80_fp4) { ORT_LLM_LOG_ENTRY(); using onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; static constexpr auto weight_only_flag = std::is_same::value ? CutlassGemmConfig::NONE : CutlassGemmConfig::WEIGHT_ONLY; @@ -630,6 +643,13 @@ MoeGemmRunner::getAmpereConfigs(int sm if (!kernels::cutlass_kernels::isValidAmpereMOESpecialisation() || (use_w4afp8 && sm != 89)) { return {}; } + // wfp4a16 is Ampere-valid (for the SM80 fused-dequant path) but only offer SM80 configs when the + // SM80 FP4 path is enabled (default-on for sm < 120); otherwise the default SM90 TMA WS path is used. + if constexpr (use_wfp4a16) { + if (!moeUseSm80Fp4(sm, use_sm80_fp4)) { + return {}; + } + } std::vector ampere_configs = kernels::cutlass_kernels::get_candidate_configs(sm, max_split_k, config_type_param); return ampere_configs; @@ -637,7 +657,7 @@ MoeGemmRunner::getAmpereConfigs(int sm template std::vector -MoeGemmRunner::getTmaWarpSpecializedConfigs(int sm) { +MoeGemmRunner::getTmaWarpSpecializedConfigs(int sm, bool use_sm80_fp4) { ORT_LLM_LOG_ENTRY(); using onnxruntime::llm::cutlass_extensions::CutlassGemmConfig; static constexpr auto weight_only_flag = std::is_same::value ? CutlassGemmConfig::NONE : CutlassGemmConfig::WEIGHT_ONLY; @@ -652,6 +672,14 @@ MoeGemmRunner::getTmaWarpSpecializedCo auto config_type_param = static_cast(weight_only_flag | simt_only_flag | grouped_gemm_flag | enable_blackwell | enable_hopper | fp8_only_flag | fp4_only_flag); ORT_ENFORCE(!(enable_blackwell && enable_hopper), "Blackwell and hopper flags are mutually exclusive"); + // When the SM80 FP4 path is enabled, wfp4a16 uses the Ampere fused-dequant grouped GEMM only; + // do not offer any TMA WS configs. + if constexpr (use_wfp4a16) { + if (moeUseSm80Fp4(sm, use_sm80_fp4)) { + return {}; + } + } + if (!isTmaWarpSpecializedGroupedGemmCompiledForSm(config_sm)) { if constexpr (use_w4afp8 || use_wfp4a16 || use_wfp4afp8 || use_wfp8a16 || use_fp4) { @@ -698,7 +726,10 @@ bool MoeGemmRunner::supportsTmaWarpSpe } if constexpr (use_wfp4a16) { - return sm_ >= 120 && kernels::cutlass_kernels::isValidHopperMOESpecialisation(); + if (moeUseSm80Fp4(sm_, use_sm80_fp4_)) { + return false; // SM80 fused-dequant grouped GEMM path; not TMA warp specialized. + } + return sm_ >= 90 && kernels::cutlass_kernels::isValidHopperMOESpecialisation(); } else { return (sm_ == 90 && kernels::cutlass_kernels::isValidHopperMOESpecialisation()) || (sm_ >= 100 && sm_ < 120 && kernels::cutlass_kernels::isValidBlackwellMOESpecialisation()) || @@ -773,7 +804,12 @@ void MoeGemmRunner::dispatchToArch( dispatchMoeGemmToCutlass( inputs, multi_processor_count_); } else if constexpr (use_wfp4a16) { - ORT_THROW("wfp4a16 (FP4 weights with FP16/BF16 activations) requires SM120+"); + if constexpr (std::is_same_v || std::is_same_v) { + dispatchMoeGemmToCutlass( + inputs, multi_processor_count_); + } else { + ORT_THROW("wfp4a16 Ampere MoE GEMM currently supports FP16/BF16 activations only"); + } } else if constexpr (use_wfp4afp8) { ORT_THROW("wfp4afp8 (FP4 weights with FP8 activations) requires SM100+"); } else if constexpr (use_wfp8a16) { @@ -855,53 +891,29 @@ void MoeGemmRunner::dispatchToArch( #endif #if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) - // Hopper W4A16 (FP4 weights + FP16/BF16 activations) WS grouped GEMM + // Hopper W4A16 (FP4 weights + FP16/BF16 activations) WS grouped GEMM. + // When the selected config is a non-TMA (SM80) config — used by the SM80 fused-dequant + // grouped GEMM path (ORT_FP4_SM80_GEMM) — fall through to the Ampere fallback below. if constexpr (use_wfp4a16) { + if (inputs.gemm_config.is_tma_warp_specialized) { #ifdef ORT_QUICK_BUILD - // Quick build only instantiates FP16+FP4 kernels; BF16+FP4 is not available. - if constexpr (!std::is_same_v) { - ORT_THROW("BF16+FP4 MoE GEMM is not available under ORT_QUICK_BUILD. Use FP16 activations instead."); - } else { + // Quick build only instantiates FP16+FP4 kernels; BF16+FP4 is not available. + if constexpr (!std::is_same_v) { + ORT_THROW("BF16+FP4 MoE GEMM is not available under ORT_QUICK_BUILD. Use FP16 activations instead."); + } else { #endif - ORT_ENFORCE(inputs.gemm_config.is_tma_warp_specialized, - "wfp4a16 is only supported for TMA warp specialization"); - // Select fusion and K tile based on runtime information - auto select_fusion = [&]() { - switch (hopper_inputs.fusion) { - case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE: - if (inputs.k % 256 == 0) { - sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( - inputs, hopper_inputs, multi_processor_count_, nullptr); - } else { - sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( - inputs, hopper_inputs, multi_processor_count_, nullptr); - } - break; - case TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE: - default: - if (inputs.k % 256 == 0) { - sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( - inputs, hopper_inputs, multi_processor_count_, nullptr); - } else { - sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( - inputs, hopper_inputs, multi_processor_count_, nullptr); - } - break; - } - }; - select_fusion(); - return; + ORT_ENFORCE(hopper_inputs.fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE, + "wfp4a16 does not support fused finalize"); + sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass( + inputs, hopper_inputs, multi_processor_count_, nullptr); + return; #ifdef ORT_QUICK_BUILD - } + } #endif + } + // Non-TMA (SM80) config: fall through to the Ampere fused-dequant grouped GEMM below. } #endif diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh new file mode 100644 index 0000000000000..2d4b400381c36 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh @@ -0,0 +1,384 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Shared device-side machinery for the symmetric weight-only MoE GEMV fast path. +// Contains fpA_intB_gemv kernels and host launch/dispatch helper templates used by +// the MXFP4 public launchers. + +#pragma once + +#include +#include + +#include "core/common/common.h" +#include "contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h" +#include "contrib_ops/cuda/llm/moe_gemm/common.h" + +namespace onnxruntime::llm { +namespace kernels { +namespace fpA_intB_gemv { + +template +__global__ void moe_gemv_kernel(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, + int num_experts, + int64_t weight_expert_stride, int64_t scale_expert_stride, int n, int k) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)) + using AccessTypeA = typename Details::AccessTypeA; + using AccessTypeW = typename Details::AccessTypeW; + + static constexpr bool Mandatory = true; + static constexpr int CtaM = 1; + static constexpr int StepK = Details::kStepK; + static constexpr int CtaK = StepK * Threads; + static_assert(CtaN % 2 == 0); + if constexpr (GroupSize != 0) { + static_assert((CtaK / Details::kInterleave) % GroupSize == 0); + } + + const int row = blockIdx.x; + + int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; +#pragma unroll 1 + for (int e = 0; e < num_experts && permuted_row_to_expert == nullptr; ++e) { + if (row >= static_cast(expert_first_token_offset[e + 1])) { + expert = e + 1; + continue; + } + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + weight += expert * weight_expert_stride; + scales += static_cast(expert) * scale_expert_stride; + if constexpr (EnableBias) { + bias += static_cast(expert) * n; + } + + const int origin_k = k, interleaved_k = k * Details::kInterleave; + + const int tile_id_m = row, tile_id_n = blockIdx.y, tid = threadIdx.x; + const int offset_m = tile_id_m * CtaM, interleaved_offset_n = tile_id_n * CtaN; + const int real_offset_n = interleaved_offset_n * Details::kInterleave + + ((tid * StepK / Details::LayoutDetails::kTileSize) % Details::kInterleave); + const int real_offset_k = + (tid * StepK / (Details::kInterleave * Details::LayoutDetails::kTileSize)) * Details::LayoutDetails::kTileSize + + ((tid * StepK) % Details::LayoutDetails::kTileSize); + + GMemIterator act_iterator( + act, offset_m * origin_k + real_offset_k, CtaK / Details::kInterleave, origin_k); + GMemIterator weight_iterator( + weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, + CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); + GMemIterator scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + + out += offset_m * n + tile_id_n * CtaN * Details::kInterleave; + if constexpr (EnableBias) { + bias += tile_id_n * CtaN * Details::kInterleave; + } + + AccT tile_acc[CtaM * CtaN]; + fill(tile_acc, static_cast(0.f)); + + TypeA vec_scale[CtaN]; + if constexpr (GroupSize == 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, 0, i); + } + } + + for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { + TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; + uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + if constexpr (GroupSize != 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, iter, i); + } + } +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + weight_iterator.load(tile_w_quantized, iter, i); + dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); + pack_to_vec2(tile_w_pack2, tile_w, i); + } +#pragma unroll + for (int i = 0; i < CtaM; ++i) { + act_iterator.load(tile_a, iter, i); + mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + } + } + epilogue(out, n, tile_acc, bias, 1.0f); +#endif +} + +template +__device__ __forceinline__ void swiglu_epilogue(void* out, void* tile_acc, void* bias, + cutlass_kernels::ActivationParams activation_params) { + static constexpr int Interleave = Details::kInterleave; + static constexpr int ThreadsPerInterleavedTile = Details::kThreadsPerInterleavedTile; + static constexpr int WarpSize = Details::kWarpSize; + static constexpr int WarpNum = Threads / WarpSize; + static constexpr int RawCols = CtaN * Interleave; + static_assert(CtaM == 1); + static_assert(RawCols % 2 == 0); + static_assert(Threads % WarpSize == 0); + + __shared__ float shmem[CtaM * CtaN * Interleave * WarpNum]; + int tid = threadIdx.x; + int warp_id = tid / WarpSize, lane_id = tid % WarpSize; +#pragma unroll + for (int n = 0; n < CtaN; ++n) { + float v = static_cast(reinterpret_cast(tile_acc)[n]); + v = warp_reduce_sum(v); + if (lane_id < Interleave * ThreadsPerInterleavedTile && lane_id % ThreadsPerInterleavedTile == 0) { + shmem[warp_id * RawCols + n * Interleave + lane_id / ThreadsPerInterleavedTile] = v; + } + } + __syncthreads(); + +#pragma unroll + for (int pair = tid; pair < RawCols / 2; pair += Threads) { + const int gate_idx = pair * 2; + const int linear_idx = gate_idx + 1; + float gate = 0.f; + float linear = 0.f; +#pragma unroll + for (int warp = 0; warp < WarpNum; ++warp) { + gate += shmem[warp * RawCols + gate_idx]; + linear += shmem[warp * RawCols + linear_idx]; + } + if constexpr (EnableBias) { + gate += static_cast(reinterpret_cast(bias)[gate_idx]); + linear += static_cast(reinterpret_cast(bias)[linear_idx]); + } + if (isfinite(activation_params.limit)) { + gate = fminf(gate, activation_params.limit); + linear = fminf(fmaxf(linear, -activation_params.limit), activation_params.limit); + } + linear += activation_params.beta; + const float sigmoid = 1.0f / (1.0f + expf(-activation_params.alpha * gate)); + reinterpret_cast(out)[pair] = static_cast(gate * sigmoid * linear); + } +} + +template +__global__ void moe_gemv_interleaved_swiglu_kernel( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, + int64_t weight_expert_stride, int64_t scale_expert_stride, int inter_size, int k, + cutlass_kernels::ActivationParams activation_params) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750)) + using AccessTypeA = typename Details::AccessTypeA; + using AccessTypeW = typename Details::AccessTypeW; + + static constexpr bool Mandatory = true; + static constexpr int CtaM = 1; + static constexpr int StepK = Details::kStepK; + static constexpr int CtaK = StepK * Threads; + static_assert(CtaN % 2 == 0); + if constexpr (GroupSize != 0) { + static_assert((CtaK / Details::kInterleave) % GroupSize == 0); + } + + const int row = blockIdx.x; + + int expert = permuted_row_to_expert != nullptr ? permuted_row_to_expert[row] : 0; +#pragma unroll 1 + for (int e = 0; e < num_experts && permuted_row_to_expert == nullptr; ++e) { + if (row >= static_cast(expert_first_token_offset[e + 1])) { + expert = e + 1; + continue; + } + break; + } + if (expert < 0 || expert >= num_experts) { + return; + } + + const float* alpha = activation_params.swiglu_alpha; + const float* beta = activation_params.swiglu_beta; + const float* limit = activation_params.swiglu_limit; + activation_params.alpha = alpha ? alpha[expert] : activation_params.alpha; + activation_params.beta = beta ? beta[expert] : activation_params.beta; + activation_params.limit = limit ? limit[expert] : activation_params.limit; + + const int n = inter_size * 2; + weight += expert * weight_expert_stride; + scales += static_cast(expert) * scale_expert_stride; + if constexpr (EnableBias) { + bias += static_cast(expert) * n; + } + + const int origin_k = k, interleaved_k = k * Details::kInterleave; + + const int tile_id_m = row, tile_id_n = blockIdx.y, tid = threadIdx.x; + const int offset_m = tile_id_m * CtaM, interleaved_offset_n = tile_id_n * CtaN; + const int real_offset_n = interleaved_offset_n * Details::kInterleave + + ((tid * StepK / Details::LayoutDetails::kTileSize) % Details::kInterleave); + const int real_offset_k = + (tid * StepK / (Details::kInterleave * Details::LayoutDetails::kTileSize)) * Details::LayoutDetails::kTileSize + + ((tid * StepK) % Details::LayoutDetails::kTileSize); + + GMemIterator act_iterator( + act, offset_m * origin_k + real_offset_k, CtaK / Details::kInterleave, origin_k); + GMemIterator weight_iterator( + weight, (interleaved_offset_n * interleaved_k + tid * StepK) / Details::kElemsPerByteW, + CtaK / Details::kElemsPerByteW, interleaved_k / Details::kElemsPerByteW); + GMemIterator scales_iterator( + scales, + (GroupSize != 0 ? real_offset_k / GroupSize * n : 0) + real_offset_n, + (GroupSize != 0 ? CtaK / Details::kInterleave / GroupSize * n : 0), Details::kInterleave); + + out += offset_m * inter_size + tile_id_n * CtaN * Details::kInterleave / 2; + if constexpr (EnableBias) { + bias += tile_id_n * CtaN * Details::kInterleave; + } + + AccT tile_acc[CtaM * CtaN]; + fill(tile_acc, static_cast(0.f)); + + TypeA vec_scale[CtaN]; + if constexpr (GroupSize == 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, 0, i); + } + } + + for (int idx_k = tid * StepK, iter = 0; idx_k < interleaved_k; idx_k += CtaK, ++iter) { + TypeA tile_a[StepK], tile_w[StepK], tile_w_pack2[CtaN * StepK]; + uint8_t tile_w_quantized[StepK / Details::kElemsPerByteW]; + if constexpr (GroupSize != 0) { +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + scales_iterator.load(vec_scale + i, iter, i); + } + } +#pragma unroll + for (int i = 0; i < CtaN; ++i) { + weight_iterator.load(tile_w_quantized, iter, i); + dequantize(tile_w, tile_w_quantized, vec_scale + i, nullptr, 1.0f); + pack_to_vec2(tile_w_pack2, tile_w, i); + } +#pragma unroll + for (int i = 0; i < CtaM; ++i) { + act_iterator.load(tile_a, iter, i); + mma(tile_acc + i * CtaN, tile_w_pack2, tile_a); + } + } + swiglu_epilogue(out, tile_acc, bias, activation_params); +#endif +} + +template +static void launch_moe_gemv(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, + int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, + cudaStream_t stream) { + const int64_t weight_expert_stride = n * k / Details::kElemsPerByteW; + const int64_t scale_expert_stride = GroupSize == 0 ? n : ((k + GroupSize - 1) / GroupSize) * n; + dim3 grid(static_cast(expanded_num_rows), static_cast(n / (CtaN * Details::kInterleave))); + dim3 block(Threads); + if (bias != nullptr) { + moe_gemv_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k)); + } else { + moe_gemv_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(n), static_cast(k)); + } +} + +template +static void launch_moe_gemv_interleaved_swiglu( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + const int64_t n = inter_size * 2; + const int64_t weight_expert_stride = n * k / Details::kElemsPerByteW; + const int64_t scale_expert_stride = GroupSize == 0 ? n : ((k + GroupSize - 1) / GroupSize) * n; + dim3 grid(static_cast(expanded_num_rows), static_cast(n / (CtaN * Details::kInterleave))); + dim3 block(Threads); + if (bias != nullptr) { + moe_gemv_interleaved_swiglu_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(inter_size), static_cast(k), activation_params); + } else { + moe_gemv_interleaved_swiglu_kernel<<>>( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + weight_expert_stride, scale_expert_stride, static_cast(inter_size), static_cast(k), activation_params); + } +} + +template +static void dispatch_moe_gemv_group_size(TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, + const int* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t n, int64_t k, + int group_size, cudaStream_t stream) { + if (group_size <= 0) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 16) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 32) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 64) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else if (group_size == 128) { + launch_moe_gemv(act, weight, scales, bias, out, expert_first_token_offset, + permuted_row_to_expert, num_experts, expanded_num_rows, n, k, stream); + } else { + ORT_THROW("unsupported MoE GEMV group_size: ", group_size); + } +} + +template +static void dispatch_moe_gemv_interleaved_swiglu_group_size( + TypeA* act, uint8_t* weight, TypeA* scales, TypeA* bias, TypeA* out, + const int64_t* expert_first_token_offset, const int* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int group_size, + cutlass_kernels::ActivationParams activation_params, cudaStream_t stream) { + if (group_size <= 0) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 16) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 32) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 64) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else if (group_size == 128) { + launch_moe_gemv_interleaved_swiglu( + act, weight, scales, bias, out, expert_first_token_offset, permuted_row_to_expert, num_experts, + expanded_num_rows, inter_size, k, activation_params, stream); + } else { + ORT_THROW("unsupported MoE GEMV group_size: ", group_size); + } +} + +} // namespace fpA_intB_gemv +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu new file mode 100644 index 0000000000000..dcafe5b0cf9e3 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu @@ -0,0 +1,267 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" + +#include +#include + +// Shared INT/FP4 profiled-shape thresholds. +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv.h" +// Shared device-side kernels + launch/dispatch helpers (fpA_intB_gemv namespace). +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh" +#include "core/platform/env_var_utils.h" + +namespace onnxruntime::llm { +namespace kernels { +namespace moe_gemv { + +namespace fiv = onnxruntime::llm::kernels::fpA_intB_gemv; + +// Tiling knobs swept by the FP4 GEMV autotuner. These mirror the INT path's default m-tile +// config (kDefaultCtaN/kDefaultThreads) plus two alternative tilings (kCtaN16/kThreads64). +static constexpr int kDefaultCtaN = 8; +static constexpr int kDefaultThreads = 128; +static constexpr int kCtaN16 = 16; +static constexpr int kThreads64 = 64; + +int CtaNForConfig(MoeGemvConfig config) { + return config == MoeGemvConfig::kCtaN16 ? kCtaN16 : kDefaultCtaN; +} + +// Opt-in interleaved MXFP4 GEMV path (env ORT_FP4_GEMV_INTERLEAVED=1). See moe_gemv_fp4.h. +// Parsed once via ORT's environment helper for consistent parsing/thread-safety. Off by +// default so the shipped single-pass ColumnMajor path stays byte-for-byte unchanged. +bool Fp4MoeGemvUseInterleaved() { + static bool const enabled = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_INTERLEAVED", 0) == 1; + return enabled; +} + +// Override for the interleaved path (env ORT_FP4_GEMV_INTERLEAVED_HALFACC=1). When set, the +// interleaved GEMV forces 16-bit (AccT=TypeA) accumulation for BOTH fp16 and bf16, overriding the +// default dtype-conditional Fp4GemvAccT policy (fp16->fp16 accum, bf16->fp32 accum). Forcing +// 16-bit on bf16 regresses bf16 accuracy, so this override is off by default. +bool Fp4MoeGemvInterleavedHalfAccum() { + static bool const enabled = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_INTERLEAVED_HALFACC", 0) == 1; + return enabled; +} + +// Interleaved-path tiling: a smaller CtaN than the default (8) to recover the occupancy the +// interleaved layout + fp32 accumulation cost. CtaN must be even (kernel static_assert) and a +// divisor that keeps n % (CtaN*kInterleave) == 0 for the target shapes. kInterleave = 4 here. +static constexpr int kInterleavedCtaN = 4; +static constexpr int kInterleavedThreads = 128; + +// --- MXFP4 (e2m1) GEMV: non-interleaved ColumnMajor layout (kInterleave = 1) --- +// Weights are the QMoERepackFP4ColToRow output ([experts, n, k/2] row-major, two e2m1 +// codes per byte, even-K in the low nibble). Block scales are the +// LaunchQMoECombineFp4ScalesForGemv output ([experts, k/32, n] in TypeA, already folded with +// the per-expert global scale). The kernel decodes e2m1 in-register via Fp4I2FConverter. +template +struct Fp4ADetails; +template <> +struct Fp4ADetails { + using Type = fiv::FP16DetailsA; +}; +#ifdef ENABLE_BF16 +template <> +struct Fp4ADetails<__nv_bfloat16> { + using Type = fiv::BF16DetailsA; +}; +#endif + +// TileSizeK is unused by the ColumnMajor (kInterleave = 1) indexing/reduction beyond the +// shmem-write lane gating (only lane 0 of each warp writes), so 64 matches the INT convention. +static constexpr int kTileSizeKFp4 = 64; +template +using Fp4KernelDetails = + fiv::KernelDetails::Type, fiv::Fp4DetailsW, fiv::ColumnMajor, false, kTileSizeKFp4>; + +// Interleaved Details. ColumnMajorInterleaved with TileSizeK = 64 and kElemBits = 4 +// gives kInterleave = 128*8/(64*4) = 4, kStepK = 128/4 = 32, kThreadsPerInterleavedTile = +// 64/32 = 2. The linear Fp4I2FConverter is reused (UseInterleavedConverter = false): the +// preprocessor's layout-only steps 1-3 produce exactly the nibble order the linear converter +// expects. Weights for this layout are produced by the interleaved branch of +// QMoE::PrePackRepackFP4Weights (CUTLASS fpA_intB W4_A16 preprocessor, apply_bias_interleave=false). +template +using Fp4KernelDetailsInterleaved = + fiv::KernelDetails::Type, fiv::Fp4DetailsW, fiv::ColumnMajorInterleaved, false, + kTileSizeKFp4>; + +// Interleaved-path accumulation policy (dtype-conditional). fp16 has a 10-bit mantissa, so 16-bit +// (half) accumulation over the interleaved kStepK=32 chains stays within tolerance and keeps +// register use low. bf16 has only 7 mantissa bits, so 16-bit accumulation loses too much precision +// (bf16 fails tolerance) and must accumulate in fp32. The ORT_FP4_GEMV_INTERLEAVED_HALFACC override +// forces 16-bit accum for BOTH dtypes. +template +using Fp4GemvAccT = std::conditional_t::value, half, float>; + +// MXFP4 GEMV shape support. Mirrors is_moe_gemv_supported but for the non-interleaved +// ColumnMajor layout: kInterleave = 1, so n need only be divisible by the CtaN tile width +// selected by `config`, and the per-thread step is StepK = 128 / activation_bits = 8 +// (not 128 / weight_bits). +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, + MoeGemvConfig config) { + if (sm < 80) { + return false; + } + if (group_size != 16 && group_size != 32) { // 32 = MXFP4 block size, 16 = NVFP4 block size + return false; + } + if (k % group_size != 0) { + return false; + } + if (expanded_num_rows <= 0 || expanded_num_rows > kMaxProfiledExpandedRows) { + return false; + } + if (n < kMinProfiledProblemDim || k < kMinProfiledProblemDim) { + return false; + } + if (expanded_num_rows > kMaxProfiledExpandedRowsForSmallProblemDim && + (n < kMinProfiledProblemDimForExpandedRowsAbove4 || k < kMinProfiledProblemDimForExpandedRowsAbove4)) { + return false; + } + if (Fp4MoeGemvUseInterleaved()) { + // Interleaved path: ColumnMajorInterleaved (kInterleave = 4, kStepK = 32), fixed CtaN = + // kInterleavedCtaN. Each block covers CtaN*kInterleave columns, and a complete interleaved + // K-tile is kStepK*kThreadsPerInterleavedTile = 32*2 = 64 wide, so require + // n % (CtaN*4) == 0 and k % 64 == 0. (gpt-oss-20b fc1 n=5760/k=2880 and fc2 n=2880/k=2880 + // both satisfy this.) `config` is ignored in this mode: CtaN/Threads are pinned to keep the + // prepacked weight layout and the kernel dispatch in agreement. + // The interleaved kStepK=32 tile is tied to the MXFP4 block-32 scale layout, so it only supports + // group_size == 32; NVFP4 (block 16) must use the non-interleaved ColumnMajor path below. + if (group_size != 32) { + return false; + } + if (n % (kInterleavedCtaN * 4) != 0) { + return false; + } + if (k % 64 != 0) { + return false; + } + return true; + } + if (n % CtaNForConfig(config) != 0) { // kInterleave = 1 + return false; + } + // StepK = 128 / activation_bits = 8; is_moe_gemv_fp4_supported requires k % group_size == 0 with + // group_size >= 16 (16 for NVFP4, 32 for MXFP4), so k % 8 == 0 always holds. + if (k % (128 / 16) != 0) { + return false; + } + return true; +} + +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size) { + return is_moe_gemv_fp4_supported(sm, expanded_num_rows, n, k, group_size, MoeGemvConfig::kDefault); +} + +template +void launch_moe_gemv_fp4_symmetric(T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, + int num_experts, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, + int sm, MoeGemvConfig config, cudaStream_t stream) { + ORT_UNUSED_PARAMETER(sm); + // Interleaved path (opt-in): ColumnMajorInterleaved layout + dtype-conditional accumulation + + // smaller CtaN. The prepacked fc2 weights are in the interleaved layout, so the kernel must match. + // CtaN/Threads are pinned (config ignored) so weights and kernel always agree. AccT follows the + // Fp4GemvAccT policy (fp16->fp16 accum, bf16->fp32 accum); HALFACC forces 16-bit for both. + if (Fp4MoeGemvUseInterleaved()) { + using DetailsI = Fp4KernelDetailsInterleaved; + if (Fp4MoeGemvInterleavedHalfAccum()) { // override: force 16-bit accum for all dtypes + fiv::dispatch_moe_gemv_group_size( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, group_size, stream); + } else { + fiv::dispatch_moe_gemv_group_size>( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, group_size, stream); + } + return; + } + using Details = Fp4KernelDetails; + // AccT follows the Fp4GemvAccT policy (fp16->fp16 accum, bf16->fp32 accum): bf16 has only 7 + // mantissa bits, so 16-bit accumulation over K loses too much precision and fails tolerance + // (e.g. NVFP4 block-16 decode at k=512). CtaN/Threads remain pure parallelization/tiling knobs + // and the accumulation dtype is identical for every config, so this sweep stays bit-exact. + auto launch = [&](auto cta_n, auto threads) { + fiv::dispatch_moe_gemv_group_size>( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, n, k, group_size, stream); + }; + if (config == MoeGemvConfig::kCtaN16) { + launch([] { return kCtaN16; }, [] { return kDefaultThreads; }); + } else if (config == MoeGemvConfig::kThreads64) { + launch([] { return kDefaultCtaN; }, [] { return kThreads64; }); + } else { + launch([] { return kDefaultCtaN; }, [] { return kDefaultThreads; }); + } +} + +template +void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, + int64_t expanded_num_rows, int64_t inter_size, int64_t k, int group_size, int sm, + cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, cudaStream_t stream) { + ORT_UNUSED_PARAMETER(sm); + // Interleaved path (opt-in): ColumnMajorInterleaved layout + dtype-conditional accumulation + + // smaller CtaN, fusing SwiGLU. Takes precedence over the split-K path so the kernel matches the + // interleaved prepacked fc1 weights. CtaN/Threads pinned (config ignored). AccT follows the + // Fp4GemvAccT policy: fp16->fp16 accum since fp16's mantissa tolerates it; bf16->fp32 accum + // since 16-bit accum fails bf16 tolerance. HALFACC forces 16-bit for both dtypes. + if (Fp4MoeGemvUseInterleaved()) { + using DetailsI = Fp4KernelDetailsInterleaved; + if (Fp4MoeGemvInterleavedHalfAccum()) { // override: force 16-bit accum for all dtypes + fiv::dispatch_moe_gemv_interleaved_swiglu_group_size( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size, k, group_size, + activation_params, stream); + } else { + fiv::dispatch_moe_gemv_interleaved_swiglu_group_size>( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size, k, group_size, + activation_params, stream); + } + return; + } + using Details = Fp4KernelDetails; + // AccT follows the Fp4GemvAccT policy (fp16->fp16, bf16->fp32); see launch_moe_gemv_fp4_symmetric. + // The CtaN/Threads sweep stays bit-exact across configs since the accumulation dtype is fixed. + auto launch = [&](auto cta_n, auto threads) { + fiv::dispatch_moe_gemv_interleaved_swiglu_group_size>( + const_cast(act), const_cast(weight), const_cast(scales), const_cast(bias), out, + expert_first_token_offset, permuted_row_to_expert, num_experts, expanded_num_rows, inter_size, k, group_size, + activation_params, stream); + }; + if (config == MoeGemvConfig::kCtaN16) { + launch([] { return kCtaN16; }, [] { return kDefaultThreads; }); + } else if (config == MoeGemvConfig::kThreads64) { + launch([] { return kDefaultCtaN; }, [] { return kThreads64; }); + } else { + launch([] { return kDefaultCtaN; }, [] { return kDefaultThreads; }); + } +} + +template void launch_moe_gemv_fp4_symmetric( + half const*, uint8_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, MoeGemvConfig, cudaStream_t); +template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( + half const*, uint8_t const*, half const*, half const*, half*, int64_t const*, int const*, int, + int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, MoeGemvConfig, cudaStream_t); +#ifdef ENABLE_BF16 +template void launch_moe_gemv_fp4_symmetric<__nv_bfloat16>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, cudaStream_t); +template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu<__nv_bfloat16>( + __nv_bfloat16 const*, uint8_t const*, __nv_bfloat16 const*, __nv_bfloat16 const*, __nv_bfloat16*, + int64_t const*, int const*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, + MoeGemvConfig, cudaStream_t); +#endif + +} // namespace moe_gemv +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h new file mode 100644 index 0000000000000..fcbdd35811ff5 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// MXFP4 (e2m1) batched GEMV fast path for symmetric weight-only MoE at small expanded +// row counts (e.g. batch-1 decode with top_k experts). Companion to the INT path declared +// in moe_gemv.h; both share the device-side machinery in moe_gemv_device.cuh. + +#pragma once + +#include +#include + +#include "contrib_ops/cuda/llm/moe_gemm/common.h" + +namespace onnxruntime::llm { +namespace kernels { +namespace moe_gemv { + +// Tiling/parallelization knob selected by the FP4 GEMV autotuner. CtaN/Threads are pure +// tiling knobs (numerically bit-exact), so the sweep only picks the fastest. +enum class MoeGemvConfig { + kDefault, + kCtaN16, + kThreads64, +}; + +// True when the opt-in interleaved MXFP4 GEMV path is enabled (env ORT_FP4_GEMV_INTERLEAVED=1). +// It combines three changes over the default path: (a) the INT4-style ColumnMajorInterleaved FP4 +// weight layout (kInterleave=4, kStepK=32) for 4x fewer K-trips, (b) dtype-conditional accumulation +// (fp32 for bf16) to keep bf16 accuracy across the longer K-chains, and (c) a smaller CtaN to +// recover the occupancy the interleave + fp32-accum cost. Default OFF; when off the shipping +// single-pass ColumnMajor path is byte-for-byte unchanged. Both PrePack (weight layout) and the +// compute dispatch query this so the prepacked weights and the kernel always agree. +bool Fp4MoeGemvUseInterleaved(); + +// FP4 GEMV shape support for the non-interleaved ColumnMajor layout (kInterleave = 1). Shared by +// both MXFP4 (group_size == 32) and NVFP4 (group_size == 16). Requires sm >= 80, n divisible by +// the kernel tile width (kCtaN) selected by `config`, and the profiled small-decode row/dim +// bounds. (The opt-in interleaved layout is MXFP4-only; see is_moe_gemv_fp4_supported in the .cu.) +// See launch_moe_gemv_fp4_symmetric. +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size); +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, + MoeGemvConfig config); + +// Launches the MXFP4 (e2m1) MoE GEMV in the non-interleaved ColumnMajor layout. +// act: [expanded_num_rows, k] permuted activations (row-major), T = half/bf16 +// weight: [num_experts, n, k/2] e2m1 codes packed two per byte (even-K low nibble) +// == LaunchQMoERepackFP4ColToRow output +// scales: [num_experts, k/group_size, n] TypeA block scales already folded with the per-expert +// global scale == LaunchQMoECombineFp4ScalesForGemv (MXFP4, group_size 32) or +// LaunchQMoECombineNvfp4ScalesForGemv (NVFP4, group_size 16) output +// bias: [num_experts, n] (T) or nullptr +// out: [expanded_num_rows, n] (row-major) +// group_size is the FP4 block size (32 for MXFP4, 16 for NVFP4). +template +void launch_moe_gemv_fp4_symmetric( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, cudaStream_t stream); + +// Launches the MXFP4 MoE GEMV and fuses interleaved SwiGLU activation. +// weight/scales/bias use raw FC1 output width n = 2 * inter_size +// out is post-activation [expanded_num_rows, inter_size] +template +void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( + T const* act, uint8_t const* weight, T const* scales, T const* bias, T* out, + int64_t const* expert_first_token_offset, int const* permuted_row_to_expert, int num_experts, int64_t expanded_num_rows, + int64_t inter_size, int64_t k, int group_size, int sm, cutlass_kernels::ActivationParams activation_params, + MoeGemvConfig config, cudaStream_t stream); + +} // namespace moe_gemv +} // namespace kernels +} // namespace onnxruntime::llm diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index a225bff8f51c6..fc56411bc4303 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu @@ -850,12 +850,35 @@ __device__ void computeTmaWarpSpecializedInputPointers(TmaWarpSpecializedGrouped assert(groupwise_scale_group_size > 0); assert(mxfp4_weight_scale || w4a8_weight_scale); if (mxfp4_weight_scale) { - constexpr int scale_cols_alignment = 4; - auto const scale_rows = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( - gemm_n, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX); - auto const scale_cols = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( - gemm_k / groupwise_scale_group_size, scale_cols_alignment); - auto const scale_offset = expert * scale_rows * scale_cols; + constexpr bool use_wfp4a16_scale_layout = +#if defined(ENABLE_FP4) && defined(ENABLE_BF16) + std::is_same_v && + (std::is_same_v || std::is_same_v); +#elif defined(ENABLE_FP4) + std::is_same_v && std::is_same_v; +#else + false; +#endif + int64_t scale_offset = 0; + if constexpr (use_wfp4a16_scale_layout) { + // The WFP4A16 scale buffer is packed in groups of 8 k-blocks (PackedScalesNum = + // CTA_K(256) / group_size(32)). PrePackFp4ScalesForTmaWs zero-pads each expert's + // k-block count up to a multiple of 8 so the GEMM's last partial CTA-K-tile can read + // a full packed group. Stride experts by the padded count to match the packed buffer. + // (For 8-aligned k-block counts this is a no-op, preserving existing 256-aligned shapes.) + constexpr int64_t kPackedScalesPerKTile = 8; + int64_t const k_blocks = gemm_k / groupwise_scale_group_size; + int64_t const k_blocks_padded = + ((k_blocks + kPackedScalesPerKTile - 1) / kPackedScalesPerKTile) * kPackedScalesPerKTile; + scale_offset = expert * (gemm_n * k_blocks_padded); + } else { + constexpr int scale_cols_alignment = 4; + auto const scale_rows = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( + gemm_n, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX); + auto const scale_cols = TmaWarpSpecializedGroupedGemmInput::alignToSfDim( + gemm_k / groupwise_scale_group_size, scale_cols_alignment); + scale_offset = expert * scale_rows * scale_cols; + } layout_info.int4_groupwise_params.ptr_s_a[out_idx] = reinterpret_cast( safe_inc_ptr(mxfp4_weight_scale, scale_offset)); @@ -2392,12 +2415,19 @@ void CutlassMoeFCRunner::value) == 0, + // 64-element (32-byte) K alignment for the FP4 weight. The original guard required + // 128 elements (64 bytes), but the WFP4A16 SM90 grouped GEMM only needs the weight K to + // satisfy the FP4 block-scale vector (32) and TMA (16-byte) alignment; the offline scale + // prepack (PrePackFp4ScalesForTmaWs) zero-pads each expert's k-block count up to a multiple + // of the packed K-tile (8) so the last partial CTA-K-tile can read a full packed group. + // Relaxing 128->64 admits gpt-oss-20b (hidden=inter=2880, 64-aligned but not 128-aligned); + // validated correct (matches dequant fallback) and clean under compute-sanitizer memcheck. + ORT_ENFORCE(hidden_size % (32 * 8 / sizeof_bits::value) == 0, "Hidden size %d does not meet minimum alignment requirements for MXFP8_MXFP4 MOE GEMM %d", - (int)hidden_size, (int)(64 * 8 / sizeof_bits::value)); - ORT_ENFORCE(inter_size % (64 * 8 / sizeof_bits::value) == 0, + (int)hidden_size, (int)(32 * 8 / sizeof_bits::value)); + ORT_ENFORCE(inter_size % (32 * 8 / sizeof_bits::value) == 0, "Inter size %d does not meet minimum alignment requirements for MXFP8_MXFP4 MOE GEMM %d", (int)inter_size, - (int)(64 * 8 / sizeof_bits::value)); + (int)(32 * 8 / sizeof_bits::value)); } else { // Require at least 128 bits of alignment for MOE GEMM ORT_ENFORCE(hidden_size % (128 / sizeof_bits::value) == 0, @@ -2555,12 +2585,18 @@ CutlassMoeFCRunner: layout_info2.stride_c = nullptr; auto alpha_scale_flat1 = use_fp4 ? quant_params.fp4.fc1.global_scale + : use_wfp4a16 ? (quant_params.mxfp8_mxfp4.fc1.global_scale + ? quant_params.mxfp8_mxfp4.fc1.global_scale + : quant_params.fp8_mxfp4.fc1.global_scale) : use_wfp4afp8 ? (quant_params.fp8_mxfp4.fc1.global_scale ? quant_params.fp8_mxfp4.fc1.global_scale : quant_params.mxfp8_mxfp4.fc1.global_scale) : use_fp8 ? fp8_dequant1 : nullptr; auto alpha_scale_flat2 = use_fp4 ? quant_params.fp4.fc2.global_scale + : use_wfp4a16 ? (quant_params.mxfp8_mxfp4.fc2.global_scale + ? quant_params.mxfp8_mxfp4.fc2.global_scale + : quant_params.fp8_mxfp4.fc2.global_scale) : use_wfp4afp8 ? (quant_params.fp8_mxfp4.fc2.global_scale ? quant_params.fp8_mxfp4.fc2.global_scale : quant_params.mxfp8_mxfp4.fc2.global_scale) @@ -2573,6 +2609,8 @@ CutlassMoeFCRunner: layout_info1.int4_groupwise_params.enabled = use_w4afp8 || use_wfp4a16 || quant_params.groupwise.group_size > 0; layout_info2.int4_groupwise_params.enabled = use_w4afp8 || use_wfp4a16 || quant_params.groupwise.group_size > 0; + layout_info1.int4_groupwise_params.use_wfp4a16 = use_wfp4a16; + layout_info2.int4_groupwise_params.use_wfp4a16 = use_wfp4a16; layout_info1.fpX_block_scaling_type = getScalingType(); layout_info2.fpX_block_scaling_type = getScalingType(); @@ -2637,7 +2675,8 @@ CutlassMoeFCRunner: gemm2_tma_ws_input.fusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE; bool apply_bias = parallelism_config.tp_rank == 0; - bool using_hopper_fused_finalize = !use_deterministic_hopper_reduce_ && gemm2_config_->sm_version == 90 && !use_w4afp8; + bool using_hopper_fused_finalize = !use_deterministic_hopper_reduce_ && gemm2_config_->sm_version == 90 && + !use_w4afp8 && !use_wfp4a16; if (using_hopper_fused_finalize) { gemm2_tma_ws_input.fusion = TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE; gemm2_tma_ws_input.setFinalizeFusionParams(final_output, permuted_token_final_scales_, @@ -2804,6 +2843,8 @@ std::map> GemmProfilerBackend::getProfile // nvllm still uses int64 because torch doesn't have fp4 yet. // bool is_fp4_act_quant = mDType == nvinfer::DataType::kFP4 || mDType == nvinfer::DataType::kINT64; bool is_fp4_w_quant = mWType == nvinfer::DataType::kFP4 || mWType == nvinfer::DataType::kINT64; + bool is_wfp4a16_groupwise_quant = is_fp4_w_quant && !is_fp8_act_quant && mSM < 90; + int const effective_group_size = is_wfp4a16_groupwise_quant && mGroupSize <= 0 ? 32 : mGroupSize; bool is_w4afp8_quant = is_int_groupwise_w_quant && is_fp8_act_quant; // bool is_wfp4afp8_quant = is_fp4_w_quant && is_fp8_act_quant; @@ -2828,16 +2869,23 @@ std::map> GemmProfilerBackend::getProfile quant_4_size = quant_2_size; } + if (is_wfp4a16_groupwise_quant) { + quant_1_size = fc1_out_size * num_experts_per_node * dtype_bytes * hidden_size / effective_group_size; + quant_2_size = hidden_size * num_experts_per_node * dtype_bytes * inter_size / effective_group_size; + quant_3_size = 0; + quant_4_size = 0; + } + // FP4 sizes - quant_1_size = is_fp4_w_quant ? sizeof(float) : quant_1_size; - quant_2_size = is_fp4_w_quant ? getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) * sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) - : quant_2_size; - quant_3_size = is_fp4_w_quant ? num_experts_per_node * sizeof(float) : quant_3_size; - quant_4_size = is_fp4_w_quant ? sizeof(float) : quant_4_size; - size_t quant_5_size = is_fp4_w_quant + quant_1_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? sizeof(float) : quant_1_size; + quant_2_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? getOffsetWeightSF(num_experts_per_node, inter_size, hidden_size, mScalingType) * sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) + : quant_2_size; + quant_3_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? num_experts_per_node * sizeof(float) : quant_3_size; + quant_4_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? sizeof(float) : quant_4_size; + size_t quant_5_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? getOffsetWeightSF(num_experts_per_node, hidden_size, inter_size, mScalingType) * sizeof(TmaWarpSpecializedGroupedGemmInput::ElementSF) : 0; - size_t quant_6_size = is_fp4_w_quant ? num_experts_per_node * sizeof(float) : 0; + size_t quant_6_size = is_fp4_w_quant && !is_wfp4a16_groupwise_quant ? num_experts_per_node * sizeof(float) : 0; size_t tma_ws_input_workspace_size = 0; if (is_tma_ws_input) { @@ -3006,6 +3054,9 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr static_cast(quant_3), static_cast(quant_4), static_cast(quant_5), static_cast(quant_6)); + } else if ((mWType == nvinfer::DataType::kFP4 || mWType == nvinfer::DataType::kINT64) && mSM < 90) { + ORT_ENFORCE(quant_1 && quant_2); + mQuantParams = QuantParams::GroupWise(mGroupSize > 0 ? mGroupSize : 32, quant_1, quant_2); } else if (mWType == nvinfer::DataType::kFP4 || mWType == nvinfer::DataType::kINT64) { // W4A16: FP4 weights with FP16/BF16 activations (no activation quantization) ORT_ENFORCE(quant_2 && quant_3 && quant_5 && quant_6); diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h index a447f373bc48e..49020fd13c12c 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h @@ -260,6 +260,12 @@ class CutlassMoeFCRunnerInterface { std::optional gemm2_config) = 0; virtual std::vector getTactics() = 0; + // wfp4a16 only: route prefill through the SM80 fused-dequant grouped GEMM (vs the SM90 TMA WS + // path). The QMoE op decides this at construction time (reading the relevant env vars then) and + // pushes it in here, so inference-time config selection does not depend on the live environment. + // Default no-op for runners that do not implement the SM80 FP4 path. + virtual void setUseSm80Fp4(bool /*use_sm80_fp4*/) {} + virtual void runMoe(void const* input_activations, void const* input_sf, int const* token_selected_experts, float const* token_final_scales, void const* fc1_expert_weights, void const* fc1_expert_biases, ActivationType fc1_activation_type, void const* fc2_expert_weights, void const* fc2_expert_biases, @@ -399,6 +405,18 @@ class CutlassMoeFCRunner : public CutlassMoeFCRunnerInterface { return moe_gemm_runner_.getConfigs(); } + // Push the QMoE op's SM80-FP4 decision into the inner runner and re-pick a valid default tactic + // from the (now filtered) config list, so even if profiling later finds nothing, the preserved + // default matches the selected GEMM family (SM80 Ampere vs SM90 TMA WS). + void setUseSm80Fp4(bool use_sm80_fp4) override { + moe_gemm_runner_.setUseSm80Fp4(use_sm80_fp4); + auto tactics = moe_gemm_runner_.getConfigs(); + if (!tactics.empty()) { + gemm1_config_ = tactics[0]; + gemm2_config_ = tactics[0]; + } + } + static std::vector getTactics(int sm) { using RunnerType = decltype(moe_gemm_runner_); return RunnerType::getConfigs(sm); diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_tma_warp_specialized_traits.h b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_tma_warp_specialized_traits.h index e8e07c78e8139..07bfc3134dc08 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_tma_warp_specialized_traits.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_tma_warp_specialized_traits.h @@ -109,12 +109,14 @@ constexpr bool isValidAmpereMOESpecialisation() { #if defined(ENABLE_FP8) && defined(ENABLE_FP4) // W8A16-FP8 (FP8 weights with non-FP8 activations) is SM90-only, not valid for Ampere. constexpr bool is_wfp8a16 = std::is_same_v && !std::is_same_v && !std::is_same_v; - return !std::is_same_v && !std::is_same_v && !is_wfp8a16; + // wfp4a16 (e2m1 WEIGHT + fp16/bf16 activation) runs the SM80 fused-dequant grouped GEMM, like INT4. + // Only the e2m1 ACTIVATION (true NVFP4) is excluded from the Ampere path. + return !std::is_same_v && !is_wfp8a16; #elif defined(ENABLE_FP8) constexpr bool is_wfp8a16 = std::is_same_v && !std::is_same_v && !std::is_same_v; return !is_wfp8a16; #elif defined(ENABLE_FP4) - return !std::is_same_v && !std::is_same_v; + return !std::is_same_v; #else return true; // Default to true #endif diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index e6ce308553c66..29d733603cdb2 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -18,6 +18,8 @@ #include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h" #include "contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h" #include "contrib_ops/cuda/utils/dump_cuda_tensor.h" #include "contrib_ops/cpu/utils/debug_macros.h" @@ -39,6 +41,91 @@ void LogQMoESwigluFusionRemapOnce() { "for backward compatibility."; }); } + +// Per-shape autotune of the fused FP4 GEMV CtaN/Threads tiling. It is opt-in because profiling +// synchronizes the inference stream; set ORT_FP4_GEMV_AUTOTUNE=1 to enable it. +// ORT_FP4_GEMV_AUTOTUNE_LOG=1 logs the chosen configs per shape. +bool Fp4GemvAutotuneEnabled() { + return onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_AUTOTUNE", 0) == 1; +} + +bool Fp4GemvAutotuneLogEnabled() { + return onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_AUTOTUNE_LOG", 0) == 1; +} + +const char* QMoEGemvConfigName(onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig config) { + using onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig; + if (config == MoeGemvConfig::kCtaN16) { + return "ctan16"; + } + if (config == MoeGemvConfig::kThreads64) { + return "threads64"; + } + return "default"; +} + +bool TryGetStaticDim(const ONNX_NAMESPACE::TensorShapeProto* shape, int dim_index, int64_t& value) { + if (shape == nullptr || dim_index >= shape->dim_size() || !shape->dim(dim_index).has_dim_value()) { + return false; + } + value = shape->dim(dim_index).dim_value(); + return value > 0; +} + +bool StaticFp4CutlassShapeSupported(const OpKernelInfo& op_kernel_info, bool is_swiglu) { +#ifdef BUILD_CUDA_EP_AS_PLUGIN + ORT_UNUSED_PARAMETER(op_kernel_info); + ORT_UNUSED_PARAMETER(is_swiglu); + return false; +#else + const auto& input_defs = op_kernel_info.node().InputDefs(); + if (input_defs.size() <= 5 || input_defs[2] == nullptr || input_defs[5] == nullptr) { + return false; + } + + const auto* fc1_shape = input_defs[2]->Shape(); + const auto* fc2_shape = input_defs[5]->Shape(); + if (fc1_shape == nullptr || fc2_shape == nullptr || fc1_shape->dim_size() != 3 || fc2_shape->dim_size() != 3) { + return false; + } + + int64_t fc1_k = 0; + int64_t fc1_packed_n = 0; + int64_t fc2_k = 0; + int64_t fc2_packed_n = 0; + if (!TryGetStaticDim(fc1_shape, 1, fc1_k) || !TryGetStaticDim(fc1_shape, 2, fc1_packed_n) || + !TryGetStaticDim(fc2_shape, 1, fc2_k) || !TryGetStaticDim(fc2_shape, 2, fc2_packed_n)) { + return false; + } + + const int64_t hidden_size = fc1_k; + const int64_t hidden_size_from_fc2 = fc2_packed_n * 2; + const int64_t fc1_n = fc1_packed_n * 2; + if (hidden_size != hidden_size_from_fc2 || (is_swiglu && fc1_n % 2 != 0)) { + return false; + } + + const int64_t inter_size = fc2_k; + const int64_t inter_size_from_fc1 = is_swiglu ? fc1_n / 2 : fc1_n; + if (inter_size != inter_size_from_fc1) { + return false; + } + + // The SM90 WFP4A16 block-scaled grouped GEMM packs block scales in groups of 8 k-blocks + // (PackedScalesNum = CTA_K(256) / group_size(32)); the GEMM reads ceil(K/256) of those + // packed groups. PrePackFp4ScalesForTmaWs zero-pads each expert's k-block count up to a + // multiple of 8 (and the runner strides experts by the padded count), so K no longer has + // to be 256-aligned: the partial last group's tail blocks are zero and are only read by + // the GEMM's last CTA-K-tile, where the matching A/B K elements are TMA-zeroed. The + // binding requirements are now the MXFP4 block-scale vector size (k % 32 == 0 for an + // integral k-block count) and the weight/activation TMA alignment. Gate on 64 as a + // conservative margin above 32; this admits gpt-oss-20b (hidden=inter=2880) while leaving + // headroom for any untested N-dim/CTA tiling alignment, and can be lowered to 32 once + // broader shapes are validated. + constexpr int64_t kMxfp4CutlassAlignment = 64; + return hidden_size % kMxfp4CutlassAlignment == 0 && inter_size % kMxfp4CutlassAlignment == 0; +#endif +} } // namespace namespace onnxruntime { @@ -58,7 +145,8 @@ namespace cuda { .TypeConstraint("T1", {DataTypeImpl::GetTensorType(), \ DataTypeImpl::GetTensorType()}) \ .TypeConstraint("T2", {DataTypeImpl::GetTensorType(), \ - DataTypeImpl::GetTensorType()}) \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()}) \ .TypeConstraint("T4", DataTypeImpl::GetTensorType()), \ QMoE); @@ -72,8 +160,22 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE block_size_ = op_kernel_info.GetAttrOrDefault("block_size", -1); this->quant_type_ = op_kernel_info.GetAttrOrDefault("quant_type", "int"); - ORT_ENFORCE(quant_type_ == "int" || quant_type_ == "fp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8", - "quant_type must be 'int', 'fp4', 'fp8', or 'wfp4afp8', but got '", quant_type_, "'"); + ORT_ENFORCE(quant_type_ == "int" || quant_type_ == "fp4" || quant_type_ == "nvfp4" || + quant_type_ == "fp8" || quant_type_ == "wfp4afp8", + "quant_type must be 'int', 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8', but got '", quant_type_, "'"); + if (quant_type_ == "nvfp4") { + constexpr int64_t kNvfp4BlockSize = 16; + ORT_ENFORCE(block_size_ == -1 || block_size_ == kNvfp4BlockSize, + "QMoE quant_type='nvfp4' requires block_size=", kNvfp4BlockSize, + ", but got ", block_size_); + block_size_ = kNvfp4BlockSize; + } else if (quant_type_ == "fp4" || quant_type_ == "wfp4afp8") { + constexpr int64_t kMxfp4BlockSize = 32; + ORT_ENFORCE(block_size_ == -1 || block_size_ == kMxfp4BlockSize, + "QMoE quant_type='", quant_type_, "' requires block_size=", kMxfp4BlockSize, + ", but got ", block_size_); + block_size_ = kMxfp4BlockSize; + } // ``weights_prepacked`` is an optional tri-state attribute (default -1) that // declares the layout of the int4/int8 fc1/fc2 weight initializers. The // concrete prepacked layouts selected by -1 and 1 are determined by the @@ -99,6 +201,7 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE weights_prepacked_ = (weights_prepacked_mode != 0); #if !defined(ENABLE_FP4) || !defined(USE_FP4_QMOE) ORT_ENFORCE(quant_type_ != "fp4", "QMoE quant_type='fp4' requires USE_FP4_QMOE with CUDA 12.8 or newer."); + ORT_ENFORCE(quant_type_ != "nvfp4", "QMoE quant_type='nvfp4' requires USE_FP4_QMOE with CUDA 12.8 or newer."); ORT_ENFORCE(quant_type_ != "wfp4afp8", "QMoE quant_type='wfp4afp8' requires USE_FP4_QMOE with CUDA 12.8 or newer."); #endif @@ -118,13 +221,101 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE #endif is_fp16_ = is_fp16; - if (quant_type_ == "fp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8") { + if (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8") { if (quant_type_ == "fp4") { ORT_ENFORCE(expert_weight_bits_ == 4, "FP4 quantization requires expert_weight_bits=4"); #if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) use_fp4_dequant_fallback_ = sm_ < 120; + const bool requested_fp4_cutlass_gemm = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_ENABLE_FP4_CUTLASS_GEMM", 0) == 1; + const bool allow_unsafe_fp4_cutlass_gemm = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_ENABLE_FP4_CUTLASS_UNSAFE", 0) == 1; + const bool fp4_cutlass_shape_supported = StaticFp4CutlassShapeSupported( + op_kernel_info, activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Swiglu); + enable_fp4_cutlass_gemm_ = is_fp16_ && sm_ == 90 && requested_fp4_cutlass_gemm && + allow_unsafe_fp4_cutlass_gemm && fp4_cutlass_shape_supported; + if (requested_fp4_cutlass_gemm && is_fp16_ && sm_ == 90 && !allow_unsafe_fp4_cutlass_gemm) { + LOGS_DEFAULT(WARNING) << "QMoE native FP4 CUTLASS GEMM was requested, but it is disabled by default " + "while its SM90 MXFP4 conversion/layout path is still being validated; " + "falling back to the dequant/GEMV path. Set ORT_ENABLE_FP4_CUTLASS_UNSAFE=1 " + "only for debugging the experimental native path."; + } else if (requested_fp4_cutlass_gemm && is_fp16_ && sm_ == 90 && !fp4_cutlass_shape_supported) { + LOGS_DEFAULT(WARNING) << "QMoE native FP4 CUTLASS GEMM was requested, but the static FP4 weight " + "shape does not meet the current WFP4A16 alignment requirements; falling " + "back to the dequant/GEMV path."; + } + if (enable_fp4_cutlass_gemm_) { + use_fp4_dequant_fallback_ = false; + } + // Fused MXFP4 GEMV (W4A16) decode path for the SM<120 fallback regime. This is the + // default: on real decode shapes it is ~18x faster than re-dequantizing all experts to + // dense BF16/FP16 every token, and it is validated bit-exact against the fallback. Set + // ORT_ENABLE_FP4_GEMV=0 to force the dequant fallback (e.g. for debugging). Prefill and + // any unsupported shape still fall through to the dequant path at dispatch time. + if (use_fp4_dequant_fallback_) { + const char* v = std::getenv("ORT_ENABLE_FP4_GEMV"); + enable_fp4_gemv_ = (v == nullptr || v[0] == '\0' || v[0] != '0'); + } else if (enable_fp4_cutlass_gemm_) { + // Native CUTLASS WFP4A16 scales well with M but is underfilled for decode (small M): + // route prefill (M >= ORT_FP4_PREFILL_MIN_TOKENS) through native, and small-decode + // shapes through the fused GEMV. Both weight/scale layouts are pre-packed alongside + // each other so a single QMoE node can serve both regimes. + enable_fp4_gemv_ = true; + fp4_prefill_min_tokens_ = onnxruntime::ParseEnvironmentVariableWithDefault( + "ORT_FP4_PREFILL_MIN_TOKENS", 64); + // Above this average per-expert token count, prefill routes to the dense A16 fallback + // runner instead of the native FP4 grouped GEMM (measured crossover ~128 tokens/expert + // on H200). 0 disables the upper bound. See ComputeInternal dispatch. + fp4_native_max_tokens_per_expert_ = onnxruntime::ParseEnvironmentVariableWithDefault( + "ORT_FP4_NATIVE_MAX_TOKENS_PER_EXPERT", 128); + } + // SM80 FP4 grouped GEMM (port of the INT4 fused-dequant Ampere path to MXFP4). + // Only meaningful on Ampere through pre-Blackwell in the dequant-fallback regime + // (80 <= sm_ < 120, e.g. H200), where the native SM90 TMA FP4 path is the slow prefill path. + // This SM80 grouped GEMM is several times faster at the gpt-oss-20b prefill regime, so it is enabled by DEFAULT for FP16/BF16; + // set ORT_FP4_SM80_GEMM=0 to fall back to the dequant path. + // If the user EXPLICITLY + // requested the native CUTLASS GEMM (ORT_ENABLE_FP4_CUTLASS_GEMM=1) we honor that intent + // and do not take the SM80 path — this keeps the kernel-side moeUseSm80Fp4() (which reads + // the same two env vars) in lock-step with this decision in every regime, including the + // native-requested-but-shape-unsupported fallback (which then uses the dequant path). + // When enabled we force the GEMV prepack (which also produces the SM80 CUTLASS-interleaved + // e2m1 weights + activation-dtype group scales) and later override the runner to the FP4 runner so + // prefill can dispatch to the SM80 DqMma grouped GEMM (see moeUseSm80Fp4 in the kernels). + enable_fp4_sm80_gemm_ = + use_fp4_dequant_fallback_ && sm_ >= 80 && sm_ < 120 && !requested_fp4_cutlass_gemm && + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_SM80_GEMM", 1) == 1; + if (enable_fp4_sm80_gemm_) { + enable_fp4_gemv_ = true; + } + // Capture the fused MXFP4 GEMV autotune knobs once here (at op-construction time) instead of + // re-reading the environment on every inference call, so a session's autotune behavior cannot + // change underneath it if the process mutates the environment after construction (mirrors the + // ORT_FP4_SM80_GEMM constructor-plumbed decision below). + enable_fp4_gemv_autotune_ = Fp4GemvAutotuneEnabled(); + enable_fp4_gemv_autotune_log_ = Fp4GemvAutotuneLogEnabled(); #else use_fp4_dequant_fallback_ = true; +#endif + } else if (quant_type_ == "nvfp4") { + ORT_ENFORCE(expert_weight_bits_ == 4, "NVFP4 quantization requires expert_weight_bits=4"); + // Native block-scaled CUTLASS GEMM for NVFP4 is Blackwell-only. On all currently + // supported GPUs (including SM90/H200) NVFP4 always uses the dequant-to-A16 fallback: + // dequantize E2M1 weights (E4M3 block scales, block size 16, per-expert global scale) + // to FP16/BF16 and run the dense A16 MoE runner. + use_fp4_dequant_fallback_ = true; +#if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) + // Fused FP4 GEMV (W4A16) decode fast path, shared with MXFP4 but with block size 16 and + // Float8E4M3FN block scales. Default-on (opt-out via ORT_ENABLE_FP4_GEMV=0): small-decode + // shapes route through the fused GEMV instead of re-dequantizing every expert to dense + // BF16/FP16 each token, and any unsupported shape (prefill / large batch) falls through to + // the dequant fallback. NVFP4 has no native CUTLASS/SM80 path, so those stay disabled. + { + const char* v = std::getenv("ORT_ENABLE_FP4_GEMV"); + enable_fp4_gemv_ = (v == nullptr || v[0] == '\0' || v[0] != '0'); + } + enable_fp4_gemv_autotune_ = Fp4GemvAutotuneEnabled(); + enable_fp4_gemv_autotune_log_ = Fp4GemvAutotuneLogEnabled(); #endif } else if (quant_type_ == "wfp4afp8") { ORT_ENFORCE(expert_weight_bits_ == 4, "WFP4AFP8 (W4A8) quantization requires expert_weight_bits=4"); @@ -148,7 +339,7 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE use_fp8_dequant_fallback_ = true; } } - if (quant_type_ == "fp4" && !use_fp4_dequant_fallback_) { + if (quant_type_ == "fp4" && (!use_fp4_dequant_fallback_ || enable_fp4_sm80_gemm_)) { #if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) if (is_fp16) { m_moe_runner = std::make_unique>( @@ -157,6 +348,22 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE m_moe_runner = std::make_unique>( sm_, activation_type_, normalize_routing_weights_, use_sparse_mixer_); } + // Dense A16 fallback runner for the large-per-expert-M prefill regime, where dequantizing + // the MXFP4 weights to FP16/BF16 and running the dense grouped GEMM beats the native FP4 + // grouped GEMM. Constructed alongside the FP4 runner so a single QMoE node can route per + // call (see fp4_native_max_tokens_per_expert_). Not used by the SM80 grouped-GEMM path + // (enable_fp4_sm80_gemm_), which routes prefill through the FP4 runner directly. + if (is_fp16) { + m_fp4_dense_fallback_runner_ = std::make_unique>( + sm_, activation_type_, normalize_routing_weights_, use_sparse_mixer_); + } else { + m_fp4_dense_fallback_runner_ = std::make_unique>( + sm_, activation_type_, normalize_routing_weights_, use_sparse_mixer_); + } + // Capture the SM80-FP4 routing decision (made above from the environment at op-construction + // time) into the runner, so inference-time config/tactic selection does not re-read the + // environment (which may have changed since the session was created, e.g. in unit tests). + m_moe_runner->setUseSm80Fp4(enable_fp4_sm80_gemm_); #endif } else if (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { #if defined(ENABLE_FP4) && defined(USE_FP4_QMOE) && defined(ENABLE_FP8) && defined(USE_FP8_QMOE) @@ -185,7 +392,7 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE } #endif } else { - // FP4/WFP4AFP8 dequant fallback or FP8 dequant fallback: use A16 runner + // FP4/NVFP4/WFP4AFP8 dequant fallback or FP8 dequant fallback: use A16 runner if (is_fp16) { m_moe_runner = std::make_unique>( sm_, activation_type_, normalize_routing_weights_, use_sparse_mixer_); @@ -227,13 +434,19 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE Status QMoE::ComputeInternal(OpKernelContext* context) const { const bool is_fp4 = (quant_type_ == "fp4"); + const bool is_nvfp4 = (quant_type_ == "nvfp4"); + // NVFP4 shares the E2M1 4-bit weight format, [E,K,N/2] weight layout, and the per-expert + // global-scale (15/16) + block-scale (3/6) inputs with MXFP4 "fp4"; it differs only in the + // block-scale decode (Float8E4M3FN vs Float8E8M0) and block size (16 vs 32). Treat both as an + // fp4-family type wherever weight-scale handling is shared. + const bool is_fp4_family = is_fp4 || is_nvfp4; const bool is_fp8 = (quant_type_ == "fp8"); const bool is_wfp4afp8 = (quant_type_ == "wfp4afp8"); const bool is_int = (quant_type_ == "int"); - // Modes that consume MXFP4 weight block scales (inputs 3/6) and per-expert global weight scales. - const bool uses_fp4_weight_scales = is_fp4 || is_wfp4afp8; + // Modes that consume FP4 weight block scales (inputs 3/6) and per-expert global weight scales. + const bool uses_fp4_weight_scales = is_fp4_family || is_wfp4afp8; // Modes that consume per-expert FP-format global weight scales (inputs 15/16). - const bool uses_global_weight_scales = is_fp4 || is_fp8 || is_wfp4afp8; + const bool uses_global_weight_scales = is_fp4_family || is_fp8 || is_wfp4afp8; const Tensor* input = context->Input(0); const Tensor* router_probs = context->Input(1); // When PrePack consumed the int4/int8 expert-weight initializers @@ -304,10 +517,13 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } // Unified FP4 inputs: block scales in fc*_scales (3/6), global scales in 15/16. - const Tensor* fp4_fc1_block_scales = (uses_fp4_weight_scales && !packed_fp4_fc1_block_scales_) ? context->Input(3) : nullptr; - const Tensor* fp4_fc2_block_scales = (uses_fp4_weight_scales && !packed_fp4_fc2_block_scales_) ? context->Input(6) : nullptr; - const Tensor* fc1_global_scale = (uses_global_weight_scales && !packed_fc1_global_scale_) ? context->Input(15) : nullptr; - const Tensor* fc2_global_scale = (uses_global_weight_scales && !packed_fc2_global_scale_) ? context->Input(16) : nullptr; + // FP4 scales are copied to GPU during PrePack but intentionally remain live inputs. The raw + // tensors provide the dtype and shape metadata required to validate packed buffers before any + // dequant or GEMV kernel indexes them. + const Tensor* fp4_fc1_block_scales = uses_fp4_weight_scales ? context->Input(3) : nullptr; + const Tensor* fp4_fc2_block_scales = uses_fp4_weight_scales ? context->Input(6) : nullptr; + const Tensor* fc1_global_scale = uses_global_weight_scales ? context->Input(15) : nullptr; + const Tensor* fc2_global_scale = uses_global_weight_scales ? context->Input(16) : nullptr; // W4A8 (WFP4AFP8) optional Variant A activation scales (per-tensor or per-expert FP8 global act scale). const Tensor* fc1_act_scale = (is_wfp4afp8 && !packed_fc1_act_scale_) ? context->Input(17) : nullptr; @@ -362,12 +578,26 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } if (uses_fp4_weight_scales) { - constexpr int64_t fp4_block_size = 32; + // MXFP4 ("fp4"/"wfp4afp8") uses block size 32 with Float8E8M0 block scales; NVFP4 ("nvfp4") + // uses block size 16 with Float8E4M3FN block scales. Both are consumed as raw uint8 bytes. + const int64_t fp4_block_size = is_nvfp4 ? 16 : 32; + const char* fp4_mode_label = is_nvfp4 ? "quant_type='nvfp4'" : "quant_type='fp4'/'wfp4afp8'"; + const char* fp4_scale_label = is_nvfp4 ? "NVFP4 block scales" : "MXFP4 block scales"; + ORT_RETURN_IF_NOT(moe_params.hidden_size % fp4_block_size == 0, + "QMoE ", fp4_mode_label, " requires hidden_size to be a multiple of ", + fp4_block_size, " for ", fp4_scale_label, ", got hidden_size=", moe_params.hidden_size, "."); + ORT_RETURN_IF_NOT(moe_params.inter_size % fp4_block_size == 0, + "QMoE ", fp4_mode_label, " requires inter_size to be a multiple of ", + fp4_block_size, " for ", fp4_scale_label, ", got inter_size=", moe_params.inter_size, "."); const int64_t fc1_out_size = is_fused_swiglu ? moe_params.inter_size * 2 : moe_params.inter_size; - auto check_fp4_block_scale = [](const Tensor* tensor, const char* name, int64_t num_experts, - int64_t n, int64_t k) -> Status { - ORT_RETURN_IF_NOT(tensor != nullptr, "QMoE quant_type='fp4'/'wfp4afp8' requires ", name, "."); - ORT_RETURN_IF_NOT(tensor->IsDataType(), name, " must be a float8e8m0 MXFP block-scale tensor."); + auto check_fp4_block_scale = [is_nvfp4](const Tensor* tensor, const char* name, int64_t num_experts, + int64_t n, int64_t k) -> Status { + ORT_RETURN_IF_NOT(tensor != nullptr, "QMoE quant_type='fp4'/'nvfp4'/'wfp4afp8' requires ", name, "."); + if (is_nvfp4) { + ORT_RETURN_IF_NOT(tensor->IsDataType(), name, " must be a float8e4m3fn NVFP4 block-scale tensor."); + } else { + ORT_RETURN_IF_NOT(tensor->IsDataType(), name, " must be a float8e8m0 MXFP block-scale tensor."); + } const auto& dims = tensor->Shape().GetDims(); ORT_RETURN_IF_NOT(dims.size() == 3 && dims[0] == num_experts && dims[1] == n && dims[2] == k, name, " must have shape (", num_experts, ", ", n, ", ", k, "), got ", tensor->Shape().ToString(), "."); @@ -451,6 +681,36 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { bool use_awq = (fc1_zeros != nullptr) || (packed_fc1_bias_ != nullptr); onnxruntime::llm::kernels::cutlass_kernels::MOEParallelismConfig parallelism_config{}; + // Per-call native-vs-dense routing for the FP4 prefill regime. The native FP4 grouped GEMM wins + // when each expert processes few tokens, but the dense A16 grouped GEMM (MXFP4 weights + // dequantized to FP16/BF16) is faster once the average per-expert token count is large + // (compute-bound; measured crossover ~128 tokens/expert on H200). When native FP4 is enabled and + // the average exceeds fp4_native_max_tokens_per_expert_, route this call through the dense + // fallback runner instead. avg_tokens_per_expert = num_rows * top_k / num_experts. + const bool fp4_native_available = is_fp4 && !use_fp4_dequant_fallback_ && m_fp4_dense_fallback_runner_ != nullptr; + const int64_t avg_tokens_per_expert = + moe_params.num_experts > 0 + ? (static_cast(moe_params.num_rows) * static_cast(k_)) / + static_cast(moe_params.num_experts) + : static_cast(moe_params.num_rows) * static_cast(k_); + const bool route_native_fp4 = + fp4_native_available && + (fp4_native_max_tokens_per_expert_ <= 0 || avg_tokens_per_expert <= fp4_native_max_tokens_per_expert_); + // SM80 FP4 grouped-GEMM prefill path. Active only when ORT_FP4_SM80_GEMM=1 and the + // GEMV prepack produced the SM80 CUTLASS-interleaved e2m1 weights + activation-dtype group scales. Decode + // shapes are still served by the fused GEMV (which returns early below); everything that + // falls through to the runner here (prefill / GEMV-unsupported shapes) runs on the FP4 + // runner via the Ampere/SM80 DqMma grouped GEMM with QuantParams::GroupWise(32, ...). + const bool fp4_sm80_prefill = + is_fp4 && enable_fp4_sm80_gemm_ && + gemv_fp4_fc1_weights_ != nullptr && gemv_fp4_fc2_weights_ != nullptr && + gemv_fp4_fc1_scales_ != nullptr && gemv_fp4_fc2_scales_ != nullptr; + // The dense fallback runner is selected only when native FP4 is available but this call exceeds + // the per-expert threshold. In every other configuration (native chosen, or native unavailable so + // m_moe_runner is itself the dense A16 runner) we use m_moe_runner. + onnxruntime::llm::kernels::cutlass_kernels::CutlassMoeFCRunnerInterface* active_runner = + (fp4_native_available && !route_native_fp4) ? m_fp4_dense_fallback_runner_.get() : m_moe_runner.get(); + // Profile and capture the best tactics under the profiler mutex, then release the mutex so // that scratch allocation, weight dequantization, scale prepping, softmax, and other // CPU-bound work can proceed concurrently across QMoE inferences. The mutex is reacquired @@ -469,15 +729,18 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // skip profiling and reuse a config cached from an earlier non-capturing run, falling back to // the default tactic when nothing is cached. cudaStream_t compute_stream = Stream(context); + // compute_stream == nullptr is the legacy default stream (stream 0), which can itself be + // under CUDA-graph capture; cudaStreamIsCapturing handles it, so query it directly rather + // than assuming a null stream is never capturing. const bool stream_is_capturing = onnxruntime::llm::common::isCapturing(compute_stream); // Use profiler with proper weight type for quantized weights if (onnxruntime::llm::common::getEnvForceDeterministicMOE()) { - auto tactics = m_moe_runner->getTactics(); + auto tactics = active_runner->getTactics(); if (!tactics.empty()) { config1 = tactics[0]; config2 = tactics[0]; - m_moe_runner->setTactic(config1, config2); + active_runner->setTactic(config1, config2); } } else { AllocatorPtr allocator; @@ -485,17 +748,28 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { mGemmProfiler.setAllocator(std::move(allocator)); mGemmProfiler.setProfilerParams(static_cast(moe_params.num_experts), static_cast(k_), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size), - static_cast(block_size_), activation_type_, + fp4_sm80_prefill ? int64_t{32} : static_cast(block_size_), activation_type_, false, true, parallelism_config, sm_); onnxruntime::llm::nvinfer::DataType dtype = is_fp16_ ? onnxruntime::llm::nvinfer::DataType::kHALF : onnxruntime::llm::nvinfer::DataType::kBF16; if (is_wfp4afp8 && !use_wfp4afp8_dequant_fallback_) { dtype = onnxruntime::llm::nvinfer::DataType::kFP8; } - // Weight type: FP4 for MXFP4, INT4 for 4-bit integer, INT8 for 8-bit integer + // Weight type: FP4 for MXFP4, INT4 for 4-bit integer, INT8 for 8-bit integer. When the + // native FP4 path routes to the dense fallback for this call, profile the dense (A16) tactic. onnxruntime::llm::nvinfer::DataType wtype; - if (is_fp4) { - wtype = use_fp4_dequant_fallback_ ? dtype : onnxruntime::llm::nvinfer::DataType::kFP4; + if (is_nvfp4) { + // NVFP4 always uses the dequant fallback, so profile against the dense (A16) tactic. + wtype = dtype; + } else if (is_fp4) { + // fp4_sm80_prefill runs the e2m1 weights through the SM80 fused-dequant grouped GEMM, + // whose scratch + groupwise scale layout match INT4-groupwise (4-bit weight + activation-dtype + // scales). Profile against kINT4 so the workspace is sized for the groupwise path + // rather than the FP4 TMA block-scale layout; the launched kernel is still the e2m1 + // kernel (templated on the runner's WeightType), and the profiler's per-tactic + // try/catch discards tile shapes that fail can_implement. + wtype = fp4_sm80_prefill ? onnxruntime::llm::nvinfer::DataType::kINT4 + : (route_native_fp4 ? onnxruntime::llm::nvinfer::DataType::kFP4 : dtype); } else if (is_wfp4afp8) { // Native W4A8 path uses FP8 activation + FP4 weights through the block-scaled dispatch. // Profile against the FP4 weight tactic; fall back to dense dtype when the dequant path is selected. @@ -523,7 +797,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // (small M) and prefill (large M) each profile and select their own best tile shape. GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), fc1_out_size, static_cast(moe_params.hidden_size)); - mGemmProfiler.profileTactics(m_moe_runner.get(), dims, id1, compute_stream); + mGemmProfiler.profileTactics(active_runner, dims, id1, compute_stream); } config1 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id1); @@ -532,7 +806,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { if (!stream_is_capturing) { GemmDims dims(static_cast(moe_params.num_rows), static_cast(moe_params.num_rows), static_cast(moe_params.hidden_size), static_cast(moe_params.inter_size)); - mGemmProfiler.profileTactics(m_moe_runner.get(), dims, id2, compute_stream); + mGemmProfiler.profileTactics(active_runner, dims, id2, compute_stream); } config2 = mGemmProfiler.getBestConfig(static_cast(moe_params.num_rows), id2); @@ -540,7 +814,7 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // cached from a prior non-capturing run, use the runner's default tactic instead of leaving // the config unset. if (!config1 || !config2) { - auto tactics = m_moe_runner->getTactics(); + auto tactics = active_runner->getTactics(); if (!tactics.empty()) { if (!config1) { config1 = tactics[0]; @@ -551,10 +825,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } } - m_moe_runner->setTactic(config1, config2); + active_runner->setTactic(config1, config2); } - workspace_size = m_moe_runner->getWorkspaceSize( + workspace_size = active_runner->getWorkspaceSize( moe_params.num_rows, moe_params.hidden_size, moe_params.inter_size, moe_params.num_experts, k_, activation_type_, parallelism_config, use_awq); } @@ -843,7 +1117,20 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { transposed_fc2_scales_holder, transposed_fc2_zp_holder, transient_fc2_bias, p_fc2_scales, p_fc2_zp); onnxruntime::llm::kernels::cutlass_kernels::QuantParams quant_params; - if (is_fp4) { + if (is_fp4 && fp4_sm80_prefill) { + // SM80 FP4 grouped-GEMM prefill: the e2m1 weights are in the SM80 CUTLASS interleaved + // layout and the scales are the activation-dtype group scales (group=32, [E, K/32, N]) the GEMV path + // already built (gemv_fp4_fc*_scales_ = e8m0_block_scale * global_scale). Feed them through + // the generic fine-grained groupwise path — identical plumbing to INT4 block-wise QMoE. + quant_params = onnxruntime::llm::kernels::cutlass_kernels::QuantParams::GroupWise( + /*group_size=*/32, + gemv_fp4_fc1_scales_.get(), + gemv_fp4_fc2_scales_.get(), + nullptr, + nullptr, + nullptr, + nullptr); + } else if (is_fp4) { // FP4 quantization: use QuantParams::FP4 with block scales and global scales const void* p_fc1_block_scales = packed_fp4_fc1_block_scales_ ? packed_fp4_fc1_block_scales_.get() : (fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() : nullptr); @@ -855,14 +1142,14 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { : (fc2_global_scale ? fc2_global_scale->DataRaw() : nullptr); ORT_RETURN_IF_NOT(p_fc1_block_scales && p_fc1_global_scale && p_fc2_block_scales && p_fc2_global_scale, "QMoE quant_type='fp4' requires fc1_scales, fc2_scales, fc1_global_scale, and fc2_global_scale."); - if (!use_fp4_dequant_fallback_) { - using NVFP4ElementSF = onnxruntime::llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::NVFP4ElementSF; - quant_params = onnxruntime::llm::kernels::cutlass_kernels::QuantParams::FP4( - nullptr, // fc1_act_global_scale (no activation quantization for W4A16) - static_cast(p_fc1_block_scales), + // Only the native FP4 runner consumes block scales via quant_params; the dense fallback runner + // receives weights already dequantized to FP16/BF16 and leaves quant_params empty. + if (route_native_fp4) { + using MXFPXElementSF = onnxruntime::llm::kernels::cutlass_kernels::TmaWarpSpecializedGroupedGemmInput::MXFPXElementSF; + quant_params = onnxruntime::llm::kernels::cutlass_kernels::QuantParams::MXFP8MXFP4( + static_cast(p_fc1_block_scales), static_cast(p_fc1_global_scale), - nullptr, // fc2_act_global_scale - static_cast(p_fc2_block_scales), + static_cast(p_fc2_block_scales), static_cast(p_fc2_global_scale)); } } else if (is_wfp4afp8) { @@ -933,19 +1220,242 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { Tensor* output = context->Output(0, input->Shape()); + // --------------------------------------------------------------------------- + // Fused MXFP4 GEMV (W4A16) decode fast path. Default-on (opt-out via ORT_ENABLE_FP4_GEMV=0) + // on the SM<120 fallback regime. Instead of dequantizing every active expert's MXFP4 weights + // to dense BF16/FP16 HBM, route small-decode shapes through a standalone fused pipeline: + // build expert maps -> expand permuted activations -> fc1 SwiGLU GEMV -> fc2 GEMV -> + // finalize routing. The pre-packed [E,n,k/2] weights and [E,k/32,n] scales are produced + // by PrePack/TryBuildGemvFp4Scales. Unsupported shapes (prefill / large batch / missing + // pre-pack buffers) fall through to the dequant fallback below. + // MXFP4 uses a fixed group size of 32 (intrinsic to the e2m1 + e8m0 block format); the + // block_size_ attribute is unset (-1) for fp4, so it is not part of the gate. + // When native CUTLASS WFP4A16 is enabled, GEMV serves only the decode regime + // (num_rows < fp4_prefill_min_tokens_); prefill (M >= threshold) falls through to native. + const bool fp4_decode_regime = + use_fp4_dequant_fallback_ || + (enable_fp4_cutlass_gemm_ && fp4_prefill_min_tokens_ > 0 && + moe_params.num_rows < fp4_prefill_min_tokens_); + const bool fp4_gemv_buffers_ready = + (enable_fp4_sm80_gemm_ + ? (gemv_fp4_fc1_weights_decode_ != nullptr && gemv_fp4_fc2_weights_decode_ != nullptr) + : (gemv_fp4_fc1_weights_ != nullptr && gemv_fp4_fc2_weights_ != nullptr)) && + gemv_fp4_fc1_scales_ != nullptr && gemv_fp4_fc2_scales_ != nullptr; + if (is_fp4_family && fp4_decode_regime && enable_fp4_gemv_ && is_fused_swiglu && + fp4_gemv_buffers_ready) { + namespace gemv = onnxruntime::llm::kernels::moe_gemv; + namespace ck = onnxruntime::llm::kernels::cutlass_kernels; + const int num_experts = static_cast(moe_params.num_experts); + const int64_t num_rows = moe_params.num_rows; + const int64_t expanded = num_rows * static_cast(k_); + const int64_t hidden = moe_params.hidden_size; + const int64_t inter = moe_params.inter_size; + const int64_t fc1_n = inter * 2; + // MXFP4 uses block size 32 (e2m1 + e8m0); NVFP4 uses block size 16 (e2m1 + e4m3). + const int gemv_group_size = is_nvfp4 ? 16 : 32; + const bool fc1_gemv_supported = gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, hidden, gemv_group_size); + const bool fc2_gemv_supported = gemv::is_moe_gemv_fp4_supported(sm_, expanded, hidden, inter, gemv_group_size); + if (num_rows > 0 && num_rows <= 256 && expanded > 0 && fc1_gemv_supported && fc2_gemv_supported) { + // When the SM80 grouped-GEMM prefill repurposes gemv_fp4_fc*_weights_ for its pair- + // interleaved layout, the decode GEMV reads its own ColToRow copy in + // gemv_fp4_fc*_weights_decode_ instead (the SM80 layout is not GEMV-decodable). + const uint8_t* gemv_fc1_weight = static_cast( + (enable_fp4_sm80_gemm_ ? gemv_fp4_fc1_weights_decode_ : gemv_fp4_fc1_weights_).get()); + const uint8_t* gemv_fc2_weight = static_cast( + (enable_fp4_sm80_gemm_ ? gemv_fp4_fc2_weights_decode_ : gemv_fp4_fc2_weights_).get()); + auto p_r2u_buf = GetScratchBuffer(expanded, GetComputeStream(context)); + auto p_exp_buf = GetScratchBuffer(expanded, GetComputeStream(context)); + auto p_efto_buf = GetScratchBuffer(num_experts + 1, GetComputeStream(context)); + const size_t elt = is_fp16_ ? sizeof(half) : sizeof(__nv_bfloat16); + auto p_act_buf = GetScratchBuffer(SafeInt(expanded) * hidden * elt, GetComputeStream(context)); + auto p_fc1_buf = GetScratchBuffer(SafeInt(expanded) * inter * elt, GetComputeStream(context)); + auto p_fc2_buf = GetScratchBuffer(SafeInt(expanded) * hidden * elt, GetComputeStream(context)); + int* p_r2u = p_r2u_buf.get(); + int* p_exp = p_exp_buf.get(); + int64_t* p_efto = p_efto_buf.get(); + + ck::ActivationParams act_params(activation_type_); + act_params.alpha = activation_alpha_; + act_params.beta = activation_beta_; + act_params.swiglu_fusion = swiglu_fusion; + act_params.limit = swiglu_limit_; + + ck::fusedBuildExpertMapsSortFirstToken( + expert_indices, p_r2u, unpermuted_row_to_permuted_row, p_exp, p_efto, + num_rows, num_experts, static_cast(k_), 0, num_experts, stream); + + const void* fc1_bias = fc1_experts_bias_optional ? fc1_experts_bias_optional->DataRaw() : nullptr; + const void* fc2_bias = fc2_experts_bias_optional ? fc2_experts_bias_optional->DataRaw() : nullptr; + + using MoeGemvConfig = gemv::MoeGemvConfig; + + // Choose the fc1 (SwiGLU) and fc2 GEMV tiling configs. CtaN/Threads are pure tiling + // knobs (numerically bit-exact), so the only goal is picking the fastest. Reuse a + // cached per-shape result when available; otherwise profile on a non-captured (warmup) + // call and freeze the choice for CUDA-graph replay. During capture (or when autotune is + // off) fall back to the default tiling. + MoeGemvConfig fc1_config = MoeGemvConfig::kDefault; + MoeGemvConfig fc2_config = MoeGemvConfig::kDefault; + + const int64_t row_bucket = + onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler::bucketM(expanded); + const Fp4GemvTuneKey tune_key{is_fp16_, row_bucket, hidden, inter, sm_}; + bool have_tune = false; + { + std::lock_guard lock(fp4_gemv_tune_cache_mutex_); + const auto cached_tune = fp4_gemv_tune_cache_.find(tune_key); + have_tune = cached_tune != fp4_gemv_tune_cache_.end(); + if (have_tune) { + fc1_config = cached_tune->second.fc1_config; + fc2_config = cached_tune->second.fc2_config; + } + } + + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_CALL_THROW(cudaStreamIsCapturing(stream, &capture_status)); + const bool is_capturing = capture_status != cudaStreamCaptureStatusNone; + const bool do_tune = enable_fp4_gemv_autotune_ && !have_tune && !is_capturing; + + auto run_fused = [&](auto* t_ptr) { + using T = std::remove_pointer_t; + ck::expandInputRowsKernelLauncher( + static_cast(input->DataRaw()), static_cast(p_act_buf.get()), + nullptr, nullptr, p_r2u, num_rows, hidden, static_cast(k_), num_experts, + quant_params, false, p_efto, nullptr, nullptr, nullptr, stream); + + auto launch_fc1 = [&](MoeGemvConfig cfg) { + gemv::launch_moe_gemv_fp4_symmetric_interleaved_swiglu( + static_cast(p_act_buf.get()), + gemv_fc1_weight, + static_cast(gemv_fp4_fc1_scales_.get()), + static_cast(fc1_bias), static_cast(p_fc1_buf.get()), + p_efto, p_exp, num_experts, expanded, inter, hidden, gemv_group_size, sm_, act_params, cfg, stream); + }; + auto launch_fc2 = [&](MoeGemvConfig cfg) { + gemv::launch_moe_gemv_fp4_symmetric( + static_cast(p_fc1_buf.get()), + gemv_fc2_weight, + static_cast(gemv_fp4_fc2_scales_.get()), + static_cast(fc2_bias), static_cast(p_fc2_buf.get()), + p_efto, p_exp, num_experts, expanded, hidden, inter, gemv_group_size, sm_, cfg, stream); + }; + + if (do_tune) { + constexpr MoeGemvConfig kCandidates[] = { + MoeGemvConfig::kDefault, MoeGemvConfig::kCtaN16, MoeGemvConfig::kThreads64}; + constexpr int kWarmup = 3; + constexpr int kIters = 20; + + cudaEvent_t start_event = nullptr; + cudaEvent_t stop_event = nullptr; + CUDA_CALL_THROW(cudaEventCreate(&start_event)); + std::unique_ptr start_event_guard(start_event, cudaEventDestroy); + CUDA_CALL_THROW(cudaEventCreate(&stop_event)); + std::unique_ptr stop_event_guard(stop_event, cudaEventDestroy); + + auto time_launch = [&](auto&& launch_fn) -> float { + for (int i = 0; i < kWarmup; ++i) { + launch_fn(); + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + CUDA_CALL_THROW(cudaEventRecord(start_event, stream)); + for (int i = 0; i < kIters; ++i) { + launch_fn(); + } + CUDA_CALL_THROW(cudaEventRecord(stop_event, stream)); + CUDA_CALL_THROW(cudaEventSynchronize(stop_event)); + float elapsed_ms = 0.0f; + CUDA_CALL_THROW(cudaEventElapsedTime(&elapsed_ms, start_event, stop_event)); + return elapsed_ms; + }; + + // fc1 reads p_act_buf (populated by the expand above). + const bool log_tune = enable_fp4_gemv_autotune_log_; + float best_fc1 = std::numeric_limits::max(); + for (MoeGemvConfig cfg : kCandidates) { + if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, hidden, gemv_group_size, cfg)) { + continue; + } + const float ms = time_launch([&] { launch_fc1(cfg); }); + if (log_tune) { + LOGS_DEFAULT(WARNING) << "FP4 GEMV autotune candidate: fc1 expanded=" << expanded + << " n=" << fc1_n << " k=" << hidden << " cfg=" << QMoEGemvConfigName(cfg) + << " " << ms << "ms/" << kIters << "it"; + } + if (ms < best_fc1) { + best_fc1 = ms; + fc1_config = cfg; + } + } + // fc2 reads p_fc1_buf; populate it once with the chosen fc1 config before timing fc2. + launch_fc1(fc1_config); + float best_fc2 = std::numeric_limits::max(); + for (MoeGemvConfig cfg : kCandidates) { + if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, hidden, inter, gemv_group_size, cfg)) { + continue; + } + const float ms = time_launch([&] { launch_fc2(cfg); }); + if (log_tune) { + LOGS_DEFAULT(WARNING) << "FP4 GEMV autotune candidate: fc2 expanded=" << expanded + << " n=" << hidden << " k=" << inter << " cfg=" << QMoEGemvConfigName(cfg) + << " " << ms << "ms/" << kIters << "it"; + } + if (ms < best_fc2) { + best_fc2 = ms; + fc2_config = cfg; + } + } + + { + std::lock_guard lock(fp4_gemv_tune_cache_mutex_); + fp4_gemv_tune_cache_.try_emplace(tune_key, Fp4GemvTuneResult{fc1_config, fc2_config}); + } + if (log_tune) { + LOGS_DEFAULT(WARNING) << "FP4 GEMV autotune: is_fp16=" << is_fp16_ << " expanded=" << expanded + << " hidden=" << hidden << " inter=" << inter << " fc1=" + << QMoEGemvConfigName(fc1_config) << " (" << best_fc1 << "ms/" << kIters + << "it) fc2=" << QMoEGemvConfigName(fc2_config) << " (" << best_fc2 << "ms/" + << kIters << "it)"; + } + } + + launch_fc1(fc1_config); + launch_fc2(fc2_config); + ck::finalizeMoeRoutingKernelLauncher( + static_cast(p_fc2_buf.get()), static_cast(output->MutableDataRaw()), + nullptr, expert_scales, unpermuted_row_to_permuted_row, p_r2u, expert_indices, + p_efto, num_rows, hidden, static_cast(k_), num_experts, + parallelism_config, false, stream); + }; + + if (is_fp16_) { + run_fused(static_cast(nullptr)); + } else { + run_fused(static_cast<__nv_bfloat16*>(nullptr)); + } + return Status::OK(); + } + } + const void* fc1_weight_data = fc1_experts_weights ? fc1_experts_weights->DataRaw() : nullptr; const void* fc2_weight_data = fc2_experts_weights ? fc2_experts_weights->DataRaw() : nullptr; - if (is_wfp4afp8 && !use_wfp4afp8_dequant_fallback_) { - // The native CUTLASS WFP4AFP8 path consumes weights in the repacked FP4 + if (fp4_sm80_prefill) { + // SM80 FP4 grouped GEMM: consume the e2m1 weights in the SM80 CUTLASS interleaved layout + // that PrePack produced into the GEMV interleaved-layout buffers (same layout the INT4 SM80 + // grouped GEMM uses). The activation-dtype group scales are wired via quant_params above. + fc1_weight_data = gemv_fp4_fc1_weights_.get(); + fc2_weight_data = gemv_fp4_fc2_weights_.get(); + } else if ((is_fp4 && route_native_fp4) || (is_wfp4afp8 && !use_wfp4afp8_dequant_fallback_)) { + // The native CUTLASS FP4 paths consume weights in the repacked FP4 // layout produced by PrePack. If PrePack never ran (e.g. // ``session.disable_prepacking`` is set) the repacked buffers stay null and // falling through to the raw initializer bytes would feed a non-CUTLASS // layout to the runner, producing silently wrong output. Fail loudly. if (packed_fp4_fc1_weights_ == nullptr || packed_fp4_fc2_weights_ == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "QMoE wfp4afp8 requires PrePack to run, but the repacked FP4 weight " + "QMoE native FP4 requires PrePack to run, but the repacked FP4 weight " "buffers were not produced (is session.disable_prepacking set?). " - "Enable prepacking to use the native WFP4AFP8 path."); + "Enable prepacking to use the native FP4 path."); } fc1_weight_data = packed_fp4_fc1_weights_.get(); fc2_weight_data = packed_fp4_fc2_weights_.get(); @@ -963,15 +1473,25 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { } IAllocatorUniquePtr dequant_fc1_weights; IAllocatorUniquePtr dequant_fc2_weights; - // FP4 (W4A16) and WFP4AFP8 (W4A8) share the MXFP4 weight format. When the native CUTLASS path - // is unavailable on the current SM, dequantize MXFP4 weights to FP16/BF16 and run the dense A16 runner. - if ((is_fp4 && use_fp4_dequant_fallback_) || (is_wfp4afp8 && use_wfp4afp8_dequant_fallback_)) { - const void* p_fc1_block_scales = packed_fp4_fc1_block_scales_ ? packed_fp4_fc1_block_scales_.get() - : (fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() : nullptr); + // FP4 (W4A16) and WFP4AFP8 (W4A8) share the MXFP4 weight format (Float8E8M0 block scales, block 32). + // NVFP4 (W4A16) uses Float8E4M3FN block scales with block 16 and always runs the dequant fallback. + // When the native CUTLASS path is unavailable on the current SM (always for NVFP4), or when native + // FP4 routes this call to the dense fallback for the large per-expert-M regime, dequantize the E2M1 + // weights to FP16/BF16 and run the dense A16 runner. + if (is_nvfp4 || + (((is_fp4 && !route_native_fp4) || (is_wfp4afp8 && use_wfp4afp8_dequant_fallback_)) && !fp4_sm80_prefill)) { + // The dequant kernel expects raw [E, n, k_blocks] block scales. When native FP4 is enabled + // (this is the large per-expert-M fallback), packed_fp4_*_block_scales_ holds the TMA-swizzled + // layout, so use the raw copy kept in gemv_fp4_*_block_raw_ instead. NVFP4 and the SM<90 + // dequant-only build have no raw copy, so packed_fp4_*_block_scales_ already holds the raw scales. + const void* p_fc1_block_scales = gemv_fp4_fc1_block_raw_ ? gemv_fp4_fc1_block_raw_.get() + : (packed_fp4_fc1_block_scales_ ? packed_fp4_fc1_block_scales_.get() + : (fp4_fc1_block_scales ? fp4_fc1_block_scales->DataRaw() : nullptr)); const void* p_fc1_global_scale = packed_fc1_global_scale_ ? packed_fc1_global_scale_.get() : (fc1_global_scale ? fc1_global_scale->DataRaw() : nullptr); - const void* p_fc2_block_scales = packed_fp4_fc2_block_scales_ ? packed_fp4_fc2_block_scales_.get() - : (fp4_fc2_block_scales ? fp4_fc2_block_scales->DataRaw() : nullptr); + const void* p_fc2_block_scales = gemv_fp4_fc2_block_raw_ ? gemv_fp4_fc2_block_raw_.get() + : (packed_fp4_fc2_block_scales_ ? packed_fp4_fc2_block_scales_.get() + : (fp4_fc2_block_scales ? fp4_fc2_block_scales->DataRaw() : nullptr)); const void* p_fc2_global_scale = packed_fc2_global_scale_ ? packed_fc2_global_scale_.get() : (fc2_global_scale ? fc2_global_scale->DataRaw() : nullptr); ORT_RETURN_IF_NOT(p_fc1_block_scales && p_fc1_global_scale && p_fc2_block_scales && p_fc2_global_scale, @@ -988,25 +1508,41 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { dequant_fc1_weights = GetScratchBuffer(fc1_bytes, GetComputeStream(context)); dequant_fc2_weights = GetScratchBuffer(fc2_bytes, GetComputeStream(context)); - if (is_fp16_) { - LaunchQMoEDequantizeFp4Weights(static_cast(fc1_experts_weights->DataRaw()), - static_cast(p_fc1_block_scales), - static_cast(p_fc1_global_scale), - static_cast(dequant_fc1_weights.get()), num_experts, fc1_n, fc1_k, stream); - LaunchQMoEDequantizeFp4Weights(static_cast(fc2_experts_weights->DataRaw()), - static_cast(p_fc2_block_scales), - static_cast(p_fc2_global_scale), - static_cast(dequant_fc2_weights.get()), num_experts, fc2_n, fc2_k, stream); - } else { - LaunchQMoEDequantizeFp4Weights(static_cast(fc1_experts_weights->DataRaw()), - static_cast(p_fc1_block_scales), - static_cast(p_fc1_global_scale), - static_cast<__nv_bfloat16*>(dequant_fc1_weights.get()), num_experts, fc1_n, fc1_k, stream); - LaunchQMoEDequantizeFp4Weights(static_cast(fc2_experts_weights->DataRaw()), - static_cast(p_fc2_block_scales), - static_cast(p_fc2_global_scale), - static_cast<__nv_bfloat16*>(dequant_fc2_weights.get()), num_experts, fc2_n, fc2_k, stream); - } + // Choose the FP4 (MXFP4 / E8M0, block 32) or NVFP4 (E4M3, block 16) dequant launcher. + auto dequant = [&](const uint8_t* weights, const uint8_t* block_scales, const float* global_scale, + void* out, int n, int k) { + if (is_fp16_) { + half* out_h = static_cast(out); + if (is_nvfp4) { + LaunchQMoEDequantizeNvfp4Weights(weights, block_scales, global_scale, out_h, num_experts, n, k, stream); + } else { + LaunchQMoEDequantizeFp4Weights(weights, block_scales, global_scale, out_h, num_experts, n, k, stream); + } + } else { + __nv_bfloat16* out_b = static_cast<__nv_bfloat16*>(out); + if (is_nvfp4) { + LaunchQMoEDequantizeNvfp4Weights(weights, block_scales, global_scale, out_b, num_experts, n, k, stream); + } else { + LaunchQMoEDequantizeFp4Weights(weights, block_scales, global_scale, out_b, num_experts, n, k, stream); + } + } + }; + // This dequant fallback is reachable only for the FP4 family (fp4/nvfp4) and the + // WFP4AFP8 dequant path -- never for quant_type='int'. Only the int path nulls out + // fc*_experts_weights (int_weights_consumed_by_prepack), so the raw weight pointers are + // guaranteed live here. Enforce that invariant explicitly so a future mode-guard change + // that lets int fall through cannot silently dereference a null weight tensor. + ORT_ENFORCE(fc1_experts_weights != nullptr && fc2_experts_weights != nullptr, + "QMoE FP4/NVFP4 dequant fallback requires the raw expert-weight tensors; got null " + "(this path must not be reached in int-weight prepack mode)."); + dequant(static_cast(fc1_experts_weights->DataRaw()), + static_cast(p_fc1_block_scales), + static_cast(p_fc1_global_scale), + dequant_fc1_weights.get(), fc1_n, fc1_k); + dequant(static_cast(fc2_experts_weights->DataRaw()), + static_cast(p_fc2_block_scales), + static_cast(p_fc2_global_scale), + dequant_fc2_weights.get(), fc2_n, fc2_k); fc1_weight_data = dequant_fc1_weights.get(); fc2_weight_data = dequant_fc2_weights.get(); } else if (is_fp8 && use_fp8_dequant_fallback_) { @@ -1050,8 +1586,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // Set tactic and run MoE. Must hold the mutex since setTactic mutates runner state. { std::lock_guard profiler_lock(mGemmProfilerMutex); - m_moe_runner->setTactic(config1, config2); - m_moe_runner->runMoe( + active_runner->setTactic(config1, config2); + active_runner->runMoe( input->DataRaw(), nullptr, expert_indices, @@ -1113,12 +1649,62 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, #define DUMP_PACK_TENSOR(name, packed_scales, scales) #endif - if (input_idx == 2 && quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { + if (input_idx == 2 && ((quant_type_ == "fp4" && !use_fp4_dequant_fallback_) || + (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_))) { PrePackRepackFP4Weights(tensor, stream, alloc, packed_fp4_fc1_weights_, is_packed); + // Native CUTLASS + GEMV coexist: also pre-pack the GEMV layout for decode (interleaved + // layout when ORT_FP4_GEMV_INTERLEAVED=1, else the [E,n,k/2] row-major layout). + if (quant_type_ == "fp4" && enable_fp4_gemv_) { + bool local_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc1_weights_, local_packed, + onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved()); + } is_packed = false; - } else if (input_idx == 5 && quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { + } else if (input_idx == 5 && ((quant_type_ == "fp4" && !use_fp4_dequant_fallback_) || + (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_))) { PrePackRepackFP4Weights(tensor, stream, alloc, packed_fp4_fc2_weights_, is_packed); + if (quant_type_ == "fp4" && enable_fp4_gemv_) { + bool local_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc2_weights_, local_packed, + onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved()); + } is_packed = false; + } else if (input_idx == 2 && (quant_type_ == "fp4" || quant_type_ == "nvfp4") && enable_fp4_gemv_) { + // Fused FP4 GEMV: lay out fc1 weights as [E, 2*inter, hidden/2] row-major. Keep + // is_packed = false so the raw [E, hidden, n/2] initializer remains available for the + // dequant fallback used by shapes the GEMV does not support. When the SM80 grouped-GEMM + // port is enabled (MXFP4 only), force the SM80 CUTLASS ColumnMajorTileInterleave layout + // so this buffer feeds the prefill grouped GEMM; the decode GEMV then reads its + // own ColToRow copy in gemv_fp4_fc1_weights_decode_ (packed below) since it cannot consume + // that layout. NVFP4 (block 16) has no native/SM80 path and always uses the plain ColToRow + // layout the non-interleaved GEMV consumes. + const bool nvfp4 = (quant_type_ == "nvfp4"); + const bool use_interleave = + !nvfp4 && (onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved() || enable_fp4_sm80_gemm_); + const bool sm80_pair = !nvfp4 && enable_fp4_sm80_gemm_; + bool local_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc1_weights_, local_packed, + use_interleave, /*sm80_pair_interleave=*/sm80_pair); + if (enable_fp4_sm80_gemm_) { + bool decode_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc1_weights_decode_, decode_packed, + onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved(), + /*sm80_pair_interleave=*/false); + } + } else if (input_idx == 5 && (quant_type_ == "fp4" || quant_type_ == "nvfp4") && enable_fp4_gemv_) { + const bool nvfp4 = (quant_type_ == "nvfp4"); + const bool use_interleave = + !nvfp4 && (onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved() || enable_fp4_sm80_gemm_); + const bool sm80_pair = !nvfp4 && enable_fp4_sm80_gemm_; + bool local_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc2_weights_, local_packed, + use_interleave, /*sm80_pair_interleave=*/sm80_pair); + if (enable_fp4_sm80_gemm_) { + bool decode_packed = false; + PrePackRepackFP4Weights(tensor, stream, alloc, gemv_fp4_fc2_weights_decode_, decode_packed, + onnxruntime::llm::kernels::moe_gemv::Fp4MoeGemvUseInterleaved(), + /*sm80_pair_interleave=*/false); + } } else if (input_idx == 2 && quant_type_ == "int" && !weights_prepacked_) { // Caller opted in (``weights_prepacked=0`` attribute) to having ORT // do the CUTLASS fpA_intB layout transform internally, instead of @@ -1134,24 +1720,72 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, PrePackIntExpertWeights(tensor, stream, alloc, packed_fc2_weights_, is_packed); } else if (input_idx == 3) { // fc1_scales DUMP_TENSOR("fc1_scales", tensor); - if (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { + if (quant_type_ == "nvfp4") { + ORT_ENFORCE(tensor.IsDataType() && tensor.Shape().NumDimensions() == 3, + "QMoE NVFP4 fc1_scales must be a 3-D float8e4m3fn tensor."); + } + if (quant_type_ == "fp4" && !use_fp4_dequant_fallback_) { + PrePackFp4ScalesForTmaWs(tensor, stream, alloc, packed_fp4_fc1_block_scales_, is_packed); + // Native CUTLASS swizzles the block scales; keep a raw copy + dims so GEMV decode can + // build its combined scale layout. + if (enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + bool raw_packed = false; + PrePackCopyToGpu(tensor, stream, alloc, gemv_fp4_fc1_block_raw_, raw_packed); + gemv_fp4_fc1_scale_e_ = tensor.Shape()[0]; + gemv_fp4_fc1_scale_n_ = tensor.Shape()[1]; + gemv_fp4_fc1_scale_kb_ = tensor.Shape()[2]; + TryBuildGemvFp4Scales(1, stream, alloc); + } + } else if (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { PrePackSwizzleBlockScales(tensor, stream, alloc, packed_fp4_fc1_block_scales_, is_packed); - } else if (quant_type_ == "fp4" || quant_type_ == "wfp4afp8") { + } else if (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "wfp4afp8") { PrePackCopyToGpu(tensor, stream, alloc, packed_fp4_fc1_block_scales_, is_packed); + if ((quant_type_ == "fp4" || quant_type_ == "nvfp4") && enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + gemv_fp4_fc1_scale_e_ = tensor.Shape()[0]; + gemv_fp4_fc1_scale_n_ = tensor.Shape()[1]; + gemv_fp4_fc1_scale_kb_ = tensor.Shape()[2]; + TryBuildGemvFp4Scales(1, stream, alloc); + } } else if (quant_type_ == "int") { PrePackTransposeAndPack(tensor, stream, alloc, packed_fc1_scales_, is_packed); DUMP_PACK_TENSOR("packed_fc1_scales", packed_fc1_scales_, tensor); } + if (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "wfp4afp8") { + is_packed = false; + } } else if (input_idx == 6) { // fc2_scales DUMP_TENSOR("fc2_scales", tensor); - if (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { + if (quant_type_ == "nvfp4") { + ORT_ENFORCE(tensor.IsDataType() && tensor.Shape().NumDimensions() == 3, + "QMoE NVFP4 fc2_scales must be a 3-D float8e4m3fn tensor."); + } + if (quant_type_ == "fp4" && !use_fp4_dequant_fallback_) { + PrePackFp4ScalesForTmaWs(tensor, stream, alloc, packed_fp4_fc2_block_scales_, is_packed); + if (enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + bool raw_packed = false; + PrePackCopyToGpu(tensor, stream, alloc, gemv_fp4_fc2_block_raw_, raw_packed); + gemv_fp4_fc2_scale_e_ = tensor.Shape()[0]; + gemv_fp4_fc2_scale_n_ = tensor.Shape()[1]; + gemv_fp4_fc2_scale_kb_ = tensor.Shape()[2]; + TryBuildGemvFp4Scales(2, stream, alloc); + } + } else if (quant_type_ == "wfp4afp8" && !use_wfp4afp8_dequant_fallback_) { PrePackSwizzleBlockScales(tensor, stream, alloc, packed_fp4_fc2_block_scales_, is_packed); - } else if (quant_type_ == "fp4" || quant_type_ == "wfp4afp8") { + } else if (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "wfp4afp8") { PrePackCopyToGpu(tensor, stream, alloc, packed_fp4_fc2_block_scales_, is_packed); + if ((quant_type_ == "fp4" || quant_type_ == "nvfp4") && enable_fp4_gemv_ && tensor.Shape().NumDimensions() == 3) { + gemv_fp4_fc2_scale_e_ = tensor.Shape()[0]; + gemv_fp4_fc2_scale_n_ = tensor.Shape()[1]; + gemv_fp4_fc2_scale_kb_ = tensor.Shape()[2]; + TryBuildGemvFp4Scales(2, stream, alloc); + } } else if (quant_type_ == "int") { PrePackTransposeAndPack(tensor, stream, alloc, packed_fc2_scales_, is_packed); DUMP_PACK_TENSOR("packed_fc2_scales", packed_fc2_scales_, tensor); } + if (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "wfp4afp8") { + is_packed = false; + } } else if (input_idx == 11) { // fc1_zeros DUMP_TENSOR("fc1_zeros", tensor); PrePackComputeBias(tensor, stream, alloc, packed_fc1_scales_, packed_fc1_bias_, is_packed); @@ -1161,12 +1795,19 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, PrePackComputeBias(tensor, stream, alloc, packed_fc2_scales_, packed_fc2_bias_, is_packed); DUMP_PACK_TENSOR("packed_fc2_bias", packed_fc2_bias_, tensor); } else if ((input_idx == 15 || input_idx == 16) && - (quant_type_ == "fp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8")) { + (quant_type_ == "fp4" || quant_type_ == "nvfp4" || quant_type_ == "fp8" || quant_type_ == "wfp4afp8")) { + ORT_ENFORCE(tensor.IsDataType() && tensor.Shape().NumDimensions() == 1, + "QMoE ", quant_type_, " global scales must be 1-D float tensors."); if (input_idx == 15) { PrePackCopyToGpu(tensor, stream, alloc, packed_fc1_global_scale_, is_packed); + fp4_fc1_global_scale_e_ = tensor.Shape()[0]; + TryBuildGemvFp4Scales(1, stream, alloc); } else { PrePackCopyToGpu(tensor, stream, alloc, packed_fc2_global_scale_, is_packed); + fp4_fc2_global_scale_e_ = tensor.Shape()[0]; + TryBuildGemvFp4Scales(2, stream, alloc); } + is_packed = false; } else if ((input_idx == 17 || input_idx == 18) && quant_type_ == "wfp4afp8") { if (input_idx == 17) { PrePackCopyToGpu(tensor, stream, alloc, packed_fc1_act_scale_, is_packed); @@ -1404,11 +2045,58 @@ void QMoE::PrePackSwizzleBlockScales(const Tensor& tensor, cudaStream_t stream, is_packed = true; } +void QMoE::PrePackFp4ScalesForTmaWs(const Tensor& tensor, cudaStream_t stream, AllocatorPtr alloc, + IAllocatorUniquePtr& packed_buf, bool& is_packed) { + auto shape = tensor.Shape(); + ORT_ENFORCE(shape.NumDimensions() == 3, "Expected 3D FP4 block scales for WFP4A16 native prepack"); + ORT_ENFORCE(tensor.IsDataType(), "Expected Float8E8M0 FP4 block scales for WFP4A16 native prepack"); + + const int64_t experts = shape[0]; + const int64_t rows = shape[1]; + const int64_t k_blocks = shape[2]; + ORT_ENFORCE(experts > 0 && rows > 0 && k_blocks > 0, + "FP4 block scales must have positive dimensions, got ", shape.ToString()); + ORT_ENFORCE(experts <= std::numeric_limits::max() && rows <= std::numeric_limits::max() && + k_blocks <= std::numeric_limits::max(), + "FP4 block-scale dimensions exceed CUDA launch int range, got ", shape.ToString()); + + const size_t bytes = tensor.SizeInBytes(); + const void* p_src = tensor.DataRaw(); + IAllocatorUniquePtr temp_src_gpu; + if (tensor.Location().device.Type() == OrtDevice::CPU) { + temp_src_gpu = IAllocator::MakeUniquePtr(alloc, bytes, true); + CUDA_CALL_THROW(cudaMemcpyAsync(temp_src_gpu.get(), p_src, bytes, cudaMemcpyHostToDevice, stream)); + p_src = temp_src_gpu.get(); + } + + // LaunchQMoEPackFp4ScalesForTmaWs pads the k-block count up to a multiple of the packed + // K-tile (8). Allocate the output buffer for the padded count so the GEMM's last partial + // CTA-K-tile can read a full packed scale group. The kernel writes every padded element + // (real blocks copied, tail blocks zeroed), so the buffer is fully initialized. + constexpr int64_t kPackedScalesPerKTile = 8; + const int64_t k_blocks_padded = + ((k_blocks + kPackedScalesPerKTile - 1) / kPackedScalesPerKTile) * kPackedScalesPerKTile; + const size_t out_bytes = + SafeInt(experts) * SafeInt(rows) * SafeInt(k_blocks_padded); + packed_buf = IAllocator::MakeUniquePtr(alloc, out_bytes, true); + LaunchQMoEPackFp4ScalesForTmaWs( + static_cast(p_src), + static_cast(packed_buf.get()), + static_cast(experts), + static_cast(rows), + static_cast(k_blocks), + stream); + + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + is_packed = true; +} + // --------------------------------------------------------------------------- // PrePack helper: Repack column-major FP4 weights to row-major using GPU kernel. // --------------------------------------------------------------------------- void QMoE::PrePackRepackFP4Weights(const Tensor& tensor, cudaStream_t stream, AllocatorPtr alloc, - IAllocatorUniquePtr& packed_buf, bool& is_packed) { + IAllocatorUniquePtr& packed_buf, bool& is_packed, + bool gemv_interleaved, bool sm80_pair_interleave) { auto shape = tensor.Shape(); ORT_ENFORCE(shape.NumDimensions() == 3, "Expected 3D FP4 weights for WFP4AFP8 native prepack"); ORT_ENFORCE(tensor.IsDataType(), "Expected uint8 FP4 weights for WFP4AFP8 native prepack"); @@ -1434,6 +2122,58 @@ void QMoE::PrePackRepackFP4Weights(const Tensor& tensor, cudaStream_t stream, Al packed_buf = IAllocator::MakeUniquePtr(alloc, bytes, true); + if (gemv_interleaved) { + // Interleaved layout: produce the INT4-style ColumnMajorInterleaved FP4 weight layout instead of the + // [E, n, k/2] row-major ColToRow layout. The source per expert is [k, n/2] bytes == a + // [K, N] row-major W4 (e2m1) tensor, which is exactly what the CUTLASS fpA_intB SM80 W4_A16 + // preprocessor consumes (shape {k, n}). The preprocessor's layout-only steps 1-3 (row- + // permute + subbyte-transpose + column-interleave) apply to e2m1 unchanged; step 4 (the + // integer +bias + pair-interleave) MUST be skipped (apply_bias_interleave=false) because it + // would corrupt the floating-point nibbles. Block scales are unchanged (kStepK=32 == MXFP4 + // block size), so TryBuildGemvFp4Scales's [E,k/32,n] layout already maps one scale per + // column per K-block. + using onnxruntime::llm::kernels::weight_only::QuantType; + // The fused FP4 GEMV always runs on the SM80 fpA_intB layout (mixed W4 + fp16/bf16 act is + // not a Hopper TMA specialisation), so preprocess for SM80 regardless of the runtime device. + const int packing_sm = + onnxruntime::llm::kernels::weight_only::get_arch_for_mixed_gemm_weight_preprocess(80); + const size_t per_expert_bytes = static_cast(k) * static_cast(n) / 2; + int8_t* dst_all = reinterpret_cast(packed_buf.get()); + const int8_t* src_all = reinterpret_cast(p_src); + + // The preprocessor mutates its input buffers in place (it ping-pongs between the dst and a + // writable src across the layout steps), so feed it a writable per-expert scratch copy of + // the (const) source rather than the source initializer itself. Reused across experts. + IAllocatorUniquePtr src_scratch = this->GetTransientScratchBuffer(per_expert_bytes); + int8_t* src_scratch_ptr = reinterpret_cast(src_scratch.get()); + IAllocatorUniquePtr permutation_map = this->GetTransientScratchBuffer(32); + + for (int64_t e = 0; e < experts; ++e) { + CUDA_CALL_THROW(cudaMemcpyAsync(src_scratch_ptr, src_all + static_cast(e) * per_expert_bytes, + per_expert_bytes, cudaMemcpyDeviceToDevice, stream)); + // synchronize=false: one host-blocking sync after the loop (stream ordering guarantees + // expert e finishes before e+1 reuses src_scratch). apply_bias_interleave=false for e2m1. + // When this buffer feeds the SM80 MoE grouped GEMM (sm80_pair_interleave), additionally + // apply step 4's nibble pair-interleave WITHOUT the +8 bias, which is the layout that + // GEMM's e2m1 dequant converter inverts. The fused GEMV decode kernel consumes the plain + // steps-1-3 layout, so it packs its own copy with sm80_pair_interleave=false. + onnxruntime::llm::kernels::weight_only::preprocess_weights_for_mixed_gemm_cuda( + stream, + packing_sm, + dst_all + static_cast(e) * per_expert_bytes, + src_scratch_ptr, + permutation_map.get(), + {static_cast(k), static_cast(n)}, + QuantType::W4_A16, + /*synchronize=*/false, + /*apply_bias_interleave=*/false, + /*interleave_without_bias=*/sm80_pair_interleave); + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + is_packed = true; + return; + } + LaunchQMoERepackFP4ColToRow( static_cast(p_src), static_cast(packed_buf.get()), @@ -1444,6 +2184,64 @@ void QMoE::PrePackRepackFP4Weights(const Tensor& tensor, cudaStream_t stream, Al is_packed = true; } +void QMoE::TryBuildGemvFp4Scales(int fc, cudaStream_t stream, AllocatorPtr alloc) { + const bool is_nvfp4 = (quant_type_ == "nvfp4"); + if (!enable_fp4_gemv_ || (quant_type_ != "fp4" && !is_nvfp4)) { + return; + } + IAllocatorUniquePtr& block = (fc == 1) ? packed_fp4_fc1_block_scales_ : packed_fp4_fc2_block_scales_; + IAllocatorUniquePtr& raw_block = (fc == 1) ? gemv_fp4_fc1_block_raw_ : gemv_fp4_fc2_block_raw_; + // Native CUTLASS swizzles packed_fp4_*_block_scales_; in that regime read the raw e8m0 copy. + IAllocatorUniquePtr& src_block = (raw_block != nullptr) ? raw_block : block; + IAllocatorUniquePtr& global = (fc == 1) ? packed_fc1_global_scale_ : packed_fc2_global_scale_; + IAllocatorUniquePtr& out = (fc == 1) ? gemv_fp4_fc1_scales_ : gemv_fp4_fc2_scales_; + const int64_t experts = (fc == 1) ? gemv_fp4_fc1_scale_e_ : gemv_fp4_fc2_scale_e_; + const int64_t n = (fc == 1) ? gemv_fp4_fc1_scale_n_ : gemv_fp4_fc2_scale_n_; + const int64_t k_blocks = (fc == 1) ? gemv_fp4_fc1_scale_kb_ : gemv_fp4_fc2_scale_kb_; + const int64_t global_scale_e = (fc == 1) ? fp4_fc1_global_scale_e_ : fp4_fc2_global_scale_e_; + + // Build exactly once, and only when both inputs are present and dimensions are known. + if (out != nullptr || src_block == nullptr || global == nullptr || experts <= 0 || n <= 0 || k_blocks <= 0 || + global_scale_e != experts) { + return; + } + + const size_t element_size = is_fp16_ ? sizeof(half) : sizeof(__nv_bfloat16); + const size_t bytes = SafeInt(experts) * SafeInt(n) * SafeInt(k_blocks) * element_size; + out = IAllocator::MakeUniquePtr(alloc, bytes, true); + + if (is_fp16_) { + if (is_nvfp4) { + LaunchQMoECombineNvfp4ScalesForGemv( + static_cast(src_block.get()), + static_cast(global.get()), + static_cast(out.get()), + static_cast(experts), static_cast(n), static_cast(k_blocks), stream); + } else { + LaunchQMoECombineFp4ScalesForGemv( + static_cast(src_block.get()), + static_cast(global.get()), + static_cast(out.get()), + static_cast(experts), static_cast(n), static_cast(k_blocks), stream); + } + } else { + if (is_nvfp4) { + LaunchQMoECombineNvfp4ScalesForGemv( + static_cast(src_block.get()), + static_cast(global.get()), + static_cast<__nv_bfloat16*>(out.get()), + static_cast(experts), static_cast(n), static_cast(k_blocks), stream); + } else { + LaunchQMoECombineFp4ScalesForGemv( + static_cast(src_block.get()), + static_cast(global.get()), + static_cast<__nv_bfloat16*>(out.get()), + static_cast(experts), static_cast(n), static_cast(k_blocks), stream); + } + } + CUDA_CALL_THROW(cudaStreamSynchronize(stream)); +} + // --------------------------------------------------------------------------- // PrePack helper: Compute bias from zero-points and scales. // --------------------------------------------------------------------------- diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index 91c84f919bc82..03ea77fad1da8 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -8,8 +8,10 @@ #include "contrib_ops/cuda/moe/moe_base.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_kernels.h" #include "contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h" +#include "contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h" #include +#include namespace onnxruntime { namespace contrib { @@ -35,8 +37,22 @@ class QMoE final : public CudaKernel, public MoEBase { IAllocatorUniquePtr& packed_buf, bool& is_packed); void PrePackSwizzleBlockScales(const Tensor& tensor, cudaStream_t stream, AllocatorPtr alloc, IAllocatorUniquePtr& packed_buf, bool& is_packed); + void PrePackFp4ScalesForTmaWs(const Tensor& tensor, cudaStream_t stream, AllocatorPtr alloc, + IAllocatorUniquePtr& packed_buf, bool& is_packed); + // ``gemv_interleaved`` selects the CUTLASS fpA_intB interleaved layout (steps 1-3) instead of + // the plain [E, n, k/2] row-major ColToRow layout. ``sm80_pair_interleave`` additionally applies + // the SM80 grouped-GEMM nibble pair-interleave (step 4, no +8 bias) on top of the interleaved + // layout; it is meaningful only when ``gemv_interleaved`` is true. The decode GEMV consumes the + // ColToRow (or steps-1-3) layout, so it must pass ``sm80_pair_interleave = false``; only the SM80 + // grouped-GEMM prefill buffer sets it true. void PrePackRepackFP4Weights(const Tensor& tensor, cudaStream_t stream, AllocatorPtr alloc, - IAllocatorUniquePtr& packed_buf, bool& is_packed); + IAllocatorUniquePtr& packed_buf, bool& is_packed, + bool gemv_interleaved = false, bool sm80_pair_interleave = false); + // Builds the fused MXFP4 GEMV scale buffer for fc (1 or 2) once both the e8m0 block + // scales (inputs 3/6) and the per-expert global scale (inputs 15/16) have been staged + // to GPU. Order-independent: invoked from both PrePack handlers; the call that completes + // the pair performs the combine. No-op unless the fused FP4 GEMV path is enabled. + void TryBuildGemvFp4Scales(int fc, cudaStream_t stream, AllocatorPtr alloc); // Prepacks int4/int8 expert weights into the CUTLASS fpA_intB layout so the // QMoE runner can consume them directly. Mirrors what MatMulNBits.PrePack // does, looped over the E expert dimension. ``tensor`` is the 3-D @@ -79,10 +95,18 @@ class QMoE final : public CudaKernel, public MoEBase { // WFP4AFP8 (W4A8) requires SM100+ (Blackwell) block-scaled tensor ops. On older GPUs we // dequantize MXFP4 weights to FP16/BF16 and run the dense A16 MoE runner. bool use_wfp4afp8_dequant_fallback_ = false; - std::string quant_type_; // "int", "fp4", "fp8", or "wfp4afp8" + std::string quant_type_; // "int", "fp4", "nvfp4", "fp8", or "wfp4afp8" std::unique_ptr m_moe_runner; + // Dense A16 (FP16/BF16) fallback runner, constructed only when the native FP4 CUTLASS path is + // enabled. The native FP4 grouped GEMM wins when each expert processes few tokens, but the dense + // grouped GEMM (FP4 weights dequantized to FP16/BF16) is faster once the per-expert token count + // is large (compute-bound regime). ComputeInternal routes per call based on average tokens per + // expert vs fp4_native_max_tokens_per_expert_. Null when native FP4 is not enabled. + std::unique_ptr + m_fp4_dense_fallback_runner_; + // Pre-packed buffers // Note: For QMoE, we need both Scales (for dequant) and Bias (derived from ZP/Scale) during inference. // PrePack logic: @@ -107,6 +131,71 @@ class QMoE final : public CudaKernel, public MoEBase { IAllocatorUniquePtr packed_fp4_fc1_block_scales_; IAllocatorUniquePtr packed_fp4_fc2_block_scales_; + // Fused MXFP4 GEMV (W4A16) decode path. Default-on (opt-out via ORT_ENABLE_FP4_GEMV=0) on + // the SM<120 dequant-fallback regime. When enabled, PrePack additionally lays out the MXFP4 + // weights in the GEMV-consumed [E, n, k/2] row-major layout and combines the e8m0 block + // scales with the per-expert global scale into the + // [E, k/32, n] activation-dtype scale layout. ComputeInternal routes small-decode shapes + // through a standalone fused GEMV pipeline (prologue -> expand -> fc1 SwiGLU GEMV -> + // fc2 GEMV -> finalize) instead of dequantizing to dense weights. Falls back to the + // dequant path for unsupported shapes (prefill / large batch). + bool enable_fp4_gemv_ = false; + bool enable_fp4_cutlass_gemm_ = false; + // Default-on (set ORT_FP4_SM80_GEMM=0 to disable) port of the INT4 SM80 fused-dequant grouped + // GEMM to MXFP4 (wfp4a16: e2m1 weight + FP16/BF16 activation). On H200 the native SM90 + // TMA FP4 path is ~50x slower than vLLM at prefill; routing prefill through the Ampere/SM80 + // DqMma grouped GEMM (the same kernel INT4 uses) closes most of the gap. When set, PrePack lays the e2m1 + // weights out in the SM80 CUTLASS ColumnMajorTileInterleave layout (reusing the GEMV + // interleaved-layout buffers) and ComputeInternal routes prefill through the FP4 runner with + // QuantParams::GroupWise(32, ...) activation-dtype group scales. Because that layout is + // incompatible with the decode GEMV kernel, PrePack also packs a separate ColToRow copy of + // the e2m1 weights (gemv_fp4_fc*_weights_decode_) that the fused GEMV decode path consumes. + bool enable_fp4_sm80_gemm_ = false; + // When native CUTLASS WFP4A16 is enabled, GEMV is also pre-packed and used for decode shapes + // (M < this threshold); prefill (M >= threshold) runs the native grouped GEMM. 0 disables the + // split (pure GEMV regime). Overridable via ORT_FP4_PREFILL_MIN_TOKENS. + int64_t fp4_prefill_min_tokens_ = 0; + // Upper bound on average tokens per expert (num_rows * top_k / num_experts) for the native FP4 + // grouped GEMM. Above it, the prefill routes to the dense A16 fallback runner, which is faster in + // the compute-bound regime (measured crossover ~128 tokens/expert on H200). 0 disables the upper + // bound (always native for prefill). Overridable via ORT_FP4_NATIVE_MAX_TOKENS_PER_EXPERT. + int64_t fp4_native_max_tokens_per_expert_ = 0; + // Per-shape autotune of the fused FP4 GEMV tiling. Both knobs are read once in the constructor + // (default-off autotune and logging) rather than fresh on every inference call, so the + // decision cannot change underneath a session whose runner/layout was already built — matching + // the constructor-plumbed pattern used by ORT_FP4_SM80_GEMM. Overridable via + // ORT_FP4_GEMV_AUTOTUNE / ORT_FP4_GEMV_AUTOTUNE_LOG. + bool enable_fp4_gemv_autotune_ = false; + bool enable_fp4_gemv_autotune_log_ = false; + IAllocatorUniquePtr gemv_fp4_fc1_weights_; // [E, 2*inter, hidden/2] row-major e2m1 + IAllocatorUniquePtr gemv_fp4_fc2_weights_; // [E, hidden, inter/2] row-major e2m1 + // When enable_fp4_sm80_gemm_ repurposes gemv_fp4_fc*_weights_ for the SM80 grouped-GEMM + // prefill (SM80 pair-interleaved layout, which the decode GEMV kernel cannot read), these + // hold the decode GEMV's own copy of the e2m1 weights in the GEMV-consumed layout + // (ColToRow, or the interleaved layout's preprocessor steps 1-3). Null when SM80 GEMM is disabled -- then the decode GEMV + // reads gemv_fp4_fc*_weights_ directly. + IAllocatorUniquePtr gemv_fp4_fc1_weights_decode_; + IAllocatorUniquePtr gemv_fp4_fc2_weights_decode_; + IAllocatorUniquePtr gemv_fp4_fc1_scales_; // [E, hidden/32, 2*inter] activation dtype + IAllocatorUniquePtr gemv_fp4_fc2_scales_; // [E, inter/32, hidden] activation dtype + // Raw [E, n, k_blocks] e8m0 block scales kept for GEMV when the native CUTLASS path has + // already swizzled packed_fp4_*_block_scales_ into the TMA layout. Empty in the pure-GEMV + // regime, where TryBuildGemvFp4Scales reads packed_fp4_*_block_scales_ directly. + IAllocatorUniquePtr gemv_fp4_fc1_block_raw_; + IAllocatorUniquePtr gemv_fp4_fc2_block_raw_; + // Block-scale dimensions captured at PrePack time so TryBuildGemvFp4Scales can size and + // launch the combine kernel once the global scale also arrives. [E, n, k_blocks]. + int64_t gemv_fp4_fc1_scale_e_ = 0; + int64_t gemv_fp4_fc1_scale_n_ = 0; + int64_t gemv_fp4_fc1_scale_kb_ = 0; + int64_t gemv_fp4_fc2_scale_e_ = 0; + int64_t gemv_fp4_fc2_scale_n_ = 0; + int64_t gemv_fp4_fc2_scale_kb_ = 0; + // Global-scale lengths captured during PrePack. TryBuildGemvFp4Scales must not + // launch until they match the staged block-scale expert dimension. + int64_t fp4_fc1_global_scale_e_ = 0; + int64_t fp4_fc2_global_scale_e_ = 0; + // Per-expert global weight scales used by FP4 and FP8 modes. IAllocatorUniquePtr packed_fc1_global_scale_; IAllocatorUniquePtr packed_fc2_global_scale_; @@ -118,6 +207,45 @@ class QMoE final : public CudaKernel, public MoEBase { mutable onnxruntime::llm::kernels::cutlass_kernels::MoeGemmProfiler mGemmProfiler; mutable std::mutex mGemmProfilerMutex; + + struct Fp4GemvTuneKey { + bool is_fp16 = false; + int64_t row_bucket = 0; + int64_t hidden = 0; + int64_t inter = 0; + int sm = 0; + + bool operator==(const Fp4GemvTuneKey& other) const { + return is_fp16 == other.is_fp16 && row_bucket == other.row_bucket && hidden == other.hidden && + inter == other.inter && sm == other.sm; + } + }; + + struct Fp4GemvTuneKeyHash { + size_t operator()(const Fp4GemvTuneKey& key) const { + size_t hash = 1469598103934665603ULL; + auto combine = [&hash](auto value) { + hash ^= static_cast(value); + hash *= 1099511628211ULL; + }; + combine(key.is_fp16); + combine(key.row_bucket); + combine(key.hidden); + combine(key.inter); + combine(key.sm); + return hash; + } + }; + + struct Fp4GemvTuneResult { + onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig fc1_config = + onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig::kDefault; + onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig fc2_config = + onnxruntime::llm::kernels::moe_gemv::MoeGemvConfig::kDefault; + }; + + mutable std::mutex fp4_gemv_tune_cache_mutex_; + mutable std::unordered_map fp4_gemv_tune_cache_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 43420c04b1a8e..44e4285519d16 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu @@ -1020,6 +1020,123 @@ void LaunchQMoEDequantizeFp4Weights( LaunchQMoEDequantizeFp4WeightsImpl(packed_weights, block_scales, global_scales, output, num_experts, n, k, stream); } +template +__global__ void QMoECombineFp4ScalesForGemvKernel( + const uint8_t* block_scales, + const float* global_scales, + T* output, + int experts, + int n, + int k_blocks) { + int64_t total = static_cast(experts) * n * k_blocks; + int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) { + return; + } + + int64_t expert_stride = static_cast(n) * k_blocks; + int expert = static_cast(index / expert_stride); + int64_t offset = index - static_cast(expert) * expert_stride; + int row = static_cast(offset / k_blocks); + int k_block = static_cast(offset - static_cast(row) * k_blocks); + + int64_t output_index = (static_cast(expert) * k_blocks + k_block) * n + row; + float scale = DecodeUE8M0(block_scales[index]) * global_scales[expert]; + output[output_index] = static_cast(scale); +} + +template +void LaunchQMoECombineFp4ScalesForGemvImpl( + const uint8_t* block_scales, + const float* global_scales, + T* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + int64_t total = static_cast(experts) * n * k_blocks; + constexpr int block = 256; + int grid = onnxruntime::narrow((total + block - 1) / block); + QMoECombineFp4ScalesForGemvKernel<<>>( + block_scales, global_scales, output, experts, n, k_blocks); + CUDA_CALL_THROW(cudaGetLastError()); +} + +void LaunchQMoECombineFp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + half* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + LaunchQMoECombineFp4ScalesForGemvImpl(block_scales, global_scales, output, experts, n, k_blocks, stream); +} + +void LaunchQMoECombineFp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + LaunchQMoECombineFp4ScalesForGemvImpl(block_scales, global_scales, output, experts, n, k_blocks, stream); +} + +__global__ void QMoEPackFp4ScalesForTmaWsKernel( + const uint8_t* input, + uint8_t* output, + int n, + int k_blocks, + int k_blocks_padded, + int packed_scales_per_k_tile, + int64_t total) { + int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) { + return; + } + + int inner = static_cast(index % packed_scales_per_k_tile); + int row = static_cast((index / packed_scales_per_k_tile) % n); + int64_t packed_k_tile = (index / packed_scales_per_k_tile) / n; + int packed_k_tiles = k_blocks_padded / packed_scales_per_k_tile; + int expert = static_cast(packed_k_tile / packed_k_tiles); + int tile = static_cast(packed_k_tile - static_cast(expert) * packed_k_tiles); + int k_block = tile * packed_scales_per_k_tile + inner; + + // Tail k-blocks added by padding k_blocks up to a multiple of packed_scales_per_k_tile are + // zero-filled. They are only read by the GEMM's last partial CTA-K-tile, where the matching + // A/B K elements are TMA-zeroed, so their scale value does not affect the result. + output[index] = (k_block < k_blocks) + ? input[(static_cast(expert) * n + row) * k_blocks + k_block] + : static_cast(0); +} + +void LaunchQMoEPackFp4ScalesForTmaWs( + const uint8_t* input, + uint8_t* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + constexpr int kPackedScalesPerKTile = 8; + // Pad k_blocks up to a multiple of the packed K-tile so the GEMM's last (possibly partial) + // CTA-K-tile can read a full packed scale group. The caller must allocate the output buffer + // for the padded k-block count (experts * n * k_blocks_padded bytes). + const int k_blocks_padded = + ((k_blocks + kPackedScalesPerKTile - 1) / kPackedScalesPerKTile) * kPackedScalesPerKTile; + int64_t total = static_cast(experts) * n * k_blocks_padded; + if (total <= 0) { + return; + } + constexpr int block = 256; + int grid = onnxruntime::narrow((total + block - 1) / block); + QMoEPackFp4ScalesForTmaWsKernel<<>>( + input, output, n, k_blocks, k_blocks_padded, kPackedScalesPerKTile, total); + CUDA_CALL_THROW(cudaGetLastError()); +} + __device__ __forceinline__ float DecodeFloat8E4M3FN(uint8_t code) { // ONNX float8e4m3fn has no infinities. The only NaN payloads are 0x7F/0xFF; // finite values, including the max finite code 0x7E, use the normal E4M3 formula. @@ -1101,6 +1218,155 @@ void LaunchQMoEDequantizeFp8Weights( LaunchQMoEDequantizeFp8WeightsImpl(weights, global_scales, output, num_experts, n, k, stream); } +// NVFP4 dequantization. Identical structure to QMoEDequantizeFp4WeightsKernel (MXFP4) except: +// (a) block size is 16 (scale_k = k / 16) instead of 32, +// (b) the per-block scale is a Float8E4M3FN byte decoded via DecodeFloat8E4M3FN +// instead of a Float8E8M0 byte decoded via DecodeUE8M0. +// Weight layout [E, K, N/2] (N-packed) and nibble selection are identical to the MXFP4 kernel. +// value = DecodeFp4E2M1(fp4_code) * DecodeE4M3(block_scale[n, k/16]) * global_scales[expert] +template +__global__ void QMoEDequantizeNvfp4WeightsKernel( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + T* output, + int num_experts, + int n, + int k) { + int64_t total = static_cast(num_experts) * n * k; + int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) { + return; + } + + int64_t expert_stride = static_cast(n) * k; + int expert = static_cast(index / expert_stride); + int64_t offset = index - static_cast(expert) * expert_stride; + int row = static_cast(offset / k); + int col = static_cast(offset - static_cast(row) * k); + + int packed_n = n / 2; + uint8_t packed = packed_weights[(static_cast(expert) * k + col) * packed_n + row / 2]; + uint8_t fp4_code = (row & 1) == 0 ? (packed & 0x0F) : (packed >> 4); + + constexpr int kNvfp4BlockSize = 16; + int scale_k = k / kNvfp4BlockSize; + uint8_t scale_code = block_scales[(static_cast(expert) * n + row) * scale_k + col / kNvfp4BlockSize]; + float value = DecodeFp4E2M1(fp4_code) * DecodeFloat8E4M3FN(scale_code) * global_scales[expert]; + output[index] = static_cast(value); +} + +template +void LaunchQMoEDequantizeNvfp4WeightsImpl( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + T* output, + int num_experts, + int n, + int k, + cudaStream_t stream) { + int64_t total = static_cast(num_experts) * n * k; + constexpr int block = 256; + int grid = onnxruntime::narrow((total + block - 1) / block); + QMoEDequantizeNvfp4WeightsKernel<<>>( + packed_weights, block_scales, global_scales, output, num_experts, n, k); + CUDA_CALL_THROW(cudaGetLastError()); +} + +void LaunchQMoEDequantizeNvfp4Weights( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + half* output, + int num_experts, + int n, + int k, + cudaStream_t stream) { + LaunchQMoEDequantizeNvfp4WeightsImpl(packed_weights, block_scales, global_scales, output, num_experts, n, k, stream); +} + +void LaunchQMoEDequantizeNvfp4Weights( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int num_experts, + int n, + int k, + cudaStream_t stream) { + LaunchQMoEDequantizeNvfp4WeightsImpl(packed_weights, block_scales, global_scales, output, num_experts, n, k, stream); +} + +// NVFP4 counterpart of QMoECombineFp4ScalesForGemvKernel. Identical [E, n, k_blocks] -> +// [E, k_blocks, n] transpose and per-expert global-scale fold, but the per-block scale byte is a +// Float8E4M3FN code decoded via DecodeFloat8E4M3FN (block size 16) instead of a Float8E8M0 code +// decoded via DecodeUE8M0 (block size 32). Produces the TypeA folded scales the fused MoE GEMV +// consumes for NVFP4 (quant_type="nvfp4"). +template +__global__ void QMoECombineNvfp4ScalesForGemvKernel( + const uint8_t* block_scales, + const float* global_scales, + T* output, + int experts, + int n, + int k_blocks) { + int64_t total = static_cast(experts) * n * k_blocks; + int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= total) { + return; + } + + int64_t expert_stride = static_cast(n) * k_blocks; + int expert = static_cast(index / expert_stride); + int64_t offset = index - static_cast(expert) * expert_stride; + int row = static_cast(offset / k_blocks); + int k_block = static_cast(offset - static_cast(row) * k_blocks); + + int64_t output_index = (static_cast(expert) * k_blocks + k_block) * n + row; + float scale = DecodeFloat8E4M3FN(block_scales[index]) * global_scales[expert]; + output[output_index] = static_cast(scale); +} + +template +void LaunchQMoECombineNvfp4ScalesForGemvImpl( + const uint8_t* block_scales, + const float* global_scales, + T* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + int64_t total = static_cast(experts) * n * k_blocks; + constexpr int block = 256; + int grid = onnxruntime::narrow((total + block - 1) / block); + QMoECombineNvfp4ScalesForGemvKernel<<>>( + block_scales, global_scales, output, experts, n, k_blocks); + CUDA_CALL_THROW(cudaGetLastError()); +} + +void LaunchQMoECombineNvfp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + half* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + LaunchQMoECombineNvfp4ScalesForGemvImpl(block_scales, global_scales, output, experts, n, k_blocks, stream); +} + +void LaunchQMoECombineNvfp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream) { + LaunchQMoECombineNvfp4ScalesForGemvImpl(block_scales, global_scales, output, experts, n, k_blocks, stream); +} + // Repack column-major FP4 packed weights to row-major layout. // Input: [experts, k, n/2] packed col-major (each byte holds 2 values along n). // Output: [experts, n, k/2] packed row-major (each byte holds 2 values along k). diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h index dfa4b4364f42a..6aba099ffae89 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h +++ b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h @@ -227,6 +227,17 @@ void LaunchQMoEDequantizeFp4Weights( int k, cudaStream_t stream); +// Packs MXFP4 e8m0 block scales from [experts, n, k_blocks] into the SM90 TMA WS +// WFP4A16 layout. The currently dispatched native WFP4A16 K tile is 256, so one +// TMA scale element contains 8 adjacent k_blocks for one output row. +void LaunchQMoEPackFp4ScalesForTmaWs( + const uint8_t* input, + uint8_t* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream); + void LaunchQMoEDequantizeFp8Weights( const uint8_t* weights, const float* global_scales, @@ -245,6 +256,29 @@ void LaunchQMoEDequantizeFp8Weights( int k, cudaStream_t stream); +// NVFP4 weight dequantization: E2M1 4-bit weights with Float8E4M3FN block scales +// (block size 16) and per-expert float32 global scales. Weight layout [E, K, N/2], +// block-scale layout [E, N, K/16]. Mirrors LaunchQMoEDequantizeFp4Weights (MXFP4). +void LaunchQMoEDequantizeNvfp4Weights( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + half* output, + int num_experts, + int n, + int k, + cudaStream_t stream); + +void LaunchQMoEDequantizeNvfp4Weights( + const uint8_t* packed_weights, + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int num_experts, + int n, + int k, + cudaStream_t stream); + // Repack column-major FP4 packed weights to row-major layout on GPU. // Input shape interpretation: [experts, k, n/2] (col-major packed), // output: [experts, n, k/2] (row-major packed). @@ -257,6 +291,42 @@ void LaunchQMoERepackFP4ColToRow( int64_t n, cudaStream_t stream); +void LaunchQMoECombineFp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + half* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream); + +void LaunchQMoECombineFp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream); + +void LaunchQMoECombineNvfp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + half* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream); + +void LaunchQMoECombineNvfp4ScalesForGemv( + const uint8_t* block_scales, + const float* global_scales, + __nv_bfloat16* output, + int experts, + int n, + int k_blocks, + cudaStream_t stream); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 065fc41b4c173..2db00f4933f10 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -1507,16 +1507,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("block_size", "Size of each quantization block along the K (input feature) dimension. " "Must be power of two and ≥ 16 (e.g., 16, 32, 64, 128). " - "If provided, both hidden_size and inter_size must be divisible by the block size. " - "Otherwise, there is no blocking and a whole column shares one scaling factor. ", + "Both hidden_size and inter_size must be divisible by the block size. " + "The FP4 modes always use blocking: MXFP4 ('fp4'/'wfp4afp8') is normalized to block_size 32 " + "and NVFP4 ('nvfp4') to block_size 16, even when block_size is omitted. " + "For integer quantization ('int'), omitting block_size means there is no blocking " + "and a whole column shares one scaling factor. ", AttributeProto::INT, OPTIONAL_VALUE) .Attr("quant_type", "Quantization type: 'int' for integer quantization (default), 'fp4' for MXFP4 quantization, " - "'fp8' for FP8 e4m3 weight-only quantization, " + "'nvfp4' for NVFP4 quantization, 'fp8' for FP8 e4m3 weight-only quantization, " "or 'wfp4afp8' for MXFP4 weight with FP8 activation. " - "When quant_type is 'fp4', weights are stored in MXFP4 format (2 values per byte), " - "fc*_scales inputs contain MXFP4 block scales, and fc*_global_scale inputs must be provided.", + "When quant_type is 'fp4' or 'nvfp4', weights are stored in E2M1 FP4 format (2 values per byte), " + "fc*_scales inputs contain the FP4 block scales, and fc*_global_scale inputs must be provided. " + "'fp4' uses Float8E8M0 block scales with block_size 32; 'nvfp4' uses Float8E4M3FN block scales " + "with block_size 16.", AttributeProto::STRING, std::string("int")) .Attr("weights_prepacked", @@ -1546,7 +1551,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(num_experts, fusion_size * inter_size), or a 3D tensor with shape " "(num_experts, fusion_size * inter_size, hidden_size / block_size) when block_size is provided. " "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " - "(num_experts, fusion_size * inter_size, hidden_size / 32). Not used for quant_type='fp8'.", + "(num_experts, fusion_size * inter_size, hidden_size / 32). " + "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " + "(num_experts, fusion_size * inter_size, hidden_size / 16). Not used for quant_type='fp8'.", "T2", OpSchema::Optional) .Input(4, @@ -1562,7 +1569,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(num_experts, hidden_size), or a 3D tensor with shape " "(num_experts, hidden_size, inter_size / block_size) when block_size is provided. " "For quant_type='fp4' or 'wfp4afp8', this is a float8e8m0 MXFP block-scale tensor with shape " - "(num_experts, hidden_size, inter_size / 32). Not used for quant_type='fp8'.", + "(num_experts, hidden_size, inter_size / 32). " + "For quant_type='nvfp4', this is a float8e4m3fn NVFP4 block-scale tensor with shape " + "(num_experts, hidden_size, inter_size / 16). Not used for quant_type='fp8'.", "T2", OpSchema::Optional) .Input(7, @@ -1620,13 +1629,13 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Input(15, "fc1_global_scale", "1D optional tensor with shape (num_experts,). " - "Per-expert global weight scale for FC1. Required when quant_type is 'fp4', 'fp8', or 'wfp4afp8'.", + "Per-expert global weight scale for FC1. Required when quant_type is 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8'.", "T4", OpSchema::Optional) .Input(16, "fc2_global_scale", "1D optional tensor with shape (num_experts,). " - "Per-expert global weight scale for FC2. Required when quant_type is 'fp4', 'fp8', or 'wfp4afp8'.", + "Per-expert global weight scale for FC2. Required when quant_type is 'fp4', 'nvfp4', 'fp8', or 'wfp4afp8'.", "T4", OpSchema::Optional) .Input(17, @@ -1656,9 +1665,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") .TypeConstraint("T1", {"tensor(uint8)", "tensor(float8e4m3fn)"}, "Constrain quantized weight types. Integer and FP4 weights use uint8. FP8 weights use float8e4m3fn.") - .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(float8e8m0)"}, + .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)", "tensor(float8e8m0)", "tensor(float8e4m3fn)"}, "Constrain scale types. Float tensors are used for integer quantization scales. " - "Float8e8m0 tensors are used for MXFP block scales.") + "Float8e8m0 tensors are used for MXFP4 block scales; float8e4m3fn tensors are used for NVFP4 block scales.") .TypeConstraint("T4", {"tensor(float)"}, "Constrain FP4 global scale type to float32 tensors.") .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput)); diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index 8da74afd7b6c1..a563df7db0cf5 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -1861,9 +1861,10 @@ def test_swiglu_qmoe_int_partial_ktile_rejected(self): # NaN-hardening regression: the INT4/INT8 weight-only path stores B in the column-interleaved # layout, whose CUTLASS K iterator requires each GEMM reduction dim to be a whole multiple of the # 64-element interleave tile (fc1.K == hidden_size, fc2.K == inter_size). A partial final K tile - # is read past the valid range and silently produces garbage/NaN. QMoE now rejects such shapes up - # front with a clear error instead of computing wrong results. Here inter_size 544 (== 17*32) is - # block-quant valid (block_size=32) but 544 % 64 == 32, so the op must raise. + # is read past the valid range and silently produces garbage/NaN. The CUTLASS mixed-GEMM weight + # prepacking (shared by the offline CudaQuantizer and ORT's PrePack) therefore rejects such shapes + # up front instead of computing wrong results. Here inter_size 544 (== 17*32) is block-quant valid + # (block_size=32) but 544 % 64 == 32, so building the quantized model must raise. torch.manual_seed(4321) numpy.random.seed(4321) @@ -1879,12 +1880,11 @@ def test_swiglu_qmoe_int_partial_ktile_rejected(self): use_asymmetric_quant=False, ) - # Build the ONNX model + session (the interleaved-layout guard fires at run time in - # ComputeInternal, not during session creation), then assert the run is rejected. - self.assertTrue(swiglu_moe.recreate_onnx_model()) - hidden_states = torch.randn(1, 1, config.hidden_size).to(device).to(torch.float16) - with self.assertRaisesRegex(Exception, "inter_size to be a multiple of 64"): - swiglu_moe.ort_forward(hidden_states) + # The partial-K-tile guard fires while prepacking the expert weights into the CUTLASS + # column-interleaved layout, so recreating the ONNX model is rejected before a session + # is ever created. + with self.assertRaisesRegex(Exception, "incompatible with column-interleave tiling"): + swiglu_moe.recreate_onnx_model() @parameterized.expand(swiglu_test_cases) def test_swiglu_qmoe_parity_bf16(self, batch_size, sequence_length, quant_bits): diff --git a/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py index 1814885543149..9e3db03c76d68 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py @@ -14,6 +14,7 @@ # -------------------------------------------------------------------------- import math +import os import unittest import numpy @@ -352,6 +353,7 @@ def _run_fp4_moe_test( onnx_dtype, use_swiglu=False, block_size=32, + gemv_mode=None, ): self._skip_if_no_fp4() @@ -414,6 +416,12 @@ def _run_fp4_moe_test( # ── create ORT session ──────────────────────────────────── opts = onnxruntime.SessionOptions() opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + # gemv_mode toggles the fused FP4 GEMV decode path (read once in the QMoE op ctor during + # session creation): "1" forces it on, "0" forces the dequant fallback, None leaves the + # default (on). Restore the previous value right after the session is built. + prev_gemv_env = os.environ.get("ORT_ENABLE_FP4_GEMV") + if gemv_mode is not None: + os.environ["ORT_ENABLE_FP4_GEMV"] = gemv_mode try: session = onnxruntime.InferenceSession( onnx_model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] @@ -422,6 +430,12 @@ def _run_fp4_moe_test( if "FP4" in str(e) or "ENABLE_FP4" in str(e) or "SM" in str(e): self.skipTest(f"FP4 not supported in this build: {e}") raise + finally: + if gemv_mode is not None: + if prev_gemv_env is None: + os.environ.pop("ORT_ENABLE_FP4_GEMV", None) + else: + os.environ["ORT_ENABLE_FP4_GEMV"] = prev_gemv_env # ── run inference ────────────────────────────────────────── input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch_dtype) @@ -526,6 +540,57 @@ def _compute_reference(input_tensor, router_logits, fc1_deq, fc2_deq, num_expert # Dimensions must be multiples of 128 for MXFP4 alignment # (MinKDimAlignmentMXFPX = 128, MinNDimAlignmentMXFPX = 128) + def test_fp4_rejects_non_32_multiple_hidden_size(self): + """Reject truncated MXFP4 block-scale shapes before launching kernels.""" + self._skip_if_no_fp4() + + hidden_size = 258 + inter_size = 256 + num_experts = 2 + top_k = 1 + num_tokens = 1 + onnx_dtype = TensorProto.FLOAT16 + + fc1_weights = torch.zeros((num_experts, hidden_size, inter_size // 2), dtype=torch.uint8, device=device) + fc2_weights = torch.zeros((num_experts, inter_size, hidden_size // 2), dtype=torch.uint8, device=device) + fc1_block_scales = torch.ones((num_experts, inter_size, hidden_size // 32), dtype=torch.uint8, device=device) + fc2_block_scales = torch.ones((num_experts, hidden_size, inter_size // 32), dtype=torch.uint8, device=device) + global_scale = torch.ones(num_experts, dtype=torch.float32, device=device) + + onnx_model = create_fp4_moe_onnx_graph( + num_tokens=num_tokens, + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + onnx_dtype=onnx_dtype, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + fc1_block_scales=fc1_block_scales, + fc1_global_scale=global_scale, + fc2_block_scales=fc2_block_scales, + fc2_global_scale=global_scale, + ) + + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + opts.add_session_config_entry("session.disable_prepacking", "1") + session = onnxruntime.InferenceSession( + onnx_model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + + input_tensor = torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.float16) + router_logits = torch.ones((num_tokens, num_experts), device=device, dtype=torch.float16) + output_tensor = torch.empty((num_tokens, hidden_size), device=device, dtype=torch.float16) + + iobinding = session.io_binding() + iobinding.bind_input("input", "cuda", 0, onnx_dtype, input_tensor.shape, input_tensor.data_ptr()) + iobinding.bind_input("router_probs", "cuda", 0, onnx_dtype, router_logits.shape, router_logits.data_ptr()) + iobinding.bind_output("output", "cuda", 0, onnx_dtype, output_tensor.shape, output_tensor.data_ptr()) + + with self.assertRaisesRegex(Exception, "hidden_size to be a multiple of 32"): + session.run_with_iobinding(iobinding) + def test_fp4_fp16_silu_basic(self): """Basic FP16 + SiLU activation.""" self._run_fp4_moe_test( @@ -623,6 +688,142 @@ def test_fp4_fp16_larger_dims(self): onnx_dtype=TensorProto.FLOAT16, ) + @parameterized.expand( + [ + (TensorProto.FLOAT16, 1, 4), + (TensorProto.FLOAT16, 2, 4), + (TensorProto.BFLOAT16, 1, 4), + (TensorProto.BFLOAT16, 2, 4), + ] + ) + def test_fp4_decode_swiglu_gemv(self, onnx_dtype, num_tokens, top_k): + """Decode-shaped SwiGLU (hidden=inter=512, expanded_rows = num_tokens*top_k <= 8). + + This shape satisfies is_moe_gemv_fp4_supported (n,k >= 512, expanded_rows in (0, 8]), + so forcing ORT_ENABLE_FP4_GEMV=1 exercises the fused MXFP4 W4A16 GEMV decode path. + test_fp4_decode_swiglu_fallback runs the identical shape with the dequant fallback + (ORT_ENABLE_FP4_GEMV=0); together they give an on/off parity check for the fused GEMV. + """ + self._run_fp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=top_k, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="1", + ) + + @parameterized.expand( + [ + (TensorProto.FLOAT16, 1, 4), + (TensorProto.BFLOAT16, 1, 4), + ] + ) + def test_fp4_decode_swiglu_fallback(self, onnx_dtype, num_tokens, top_k): + """Same decode-shaped SwiGLU as test_fp4_decode_swiglu_gemv but with the fused GEMV + disabled (ORT_ENABLE_FP4_GEMV=0), so it exercises the dequant fallback on a shape the + GEMV path also supports.""" + self._run_fp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=8, + top_k=top_k, + num_tokens=num_tokens, + onnx_dtype=onnx_dtype, + use_swiglu=True, + gemv_mode="0", + ) + + def test_fp4_native_cutlass_row_varying_scales(self): + """Native SM90 WFP4A16 scale prepack preserves per-output-row MXFP4 scales.""" + self._skip_if_no_fp4() + if _cuda_sm() != 90: + self.skipTest(f"Native FP4 CUTLASS GEMM is currently enabled only on SM90, got SM{_cuda_sm()}") + + hidden_size = 512 + inter_size = 512 + num_experts = 1 + top_k = 1 + num_tokens = 1 + onnx_dtype = TensorProto.FLOAT16 + + fc1_weights = torch.full((num_experts, hidden_size, inter_size // 2), 0x22, dtype=torch.uint8, device=device) + fc2_codes = torch.zeros((hidden_size, inter_size), dtype=torch.uint8, device=device) + diagonal = torch.arange(hidden_size, device=device) + fc2_codes[diagonal, diagonal] = 2 + fc2_codes_kn = fc2_codes.T.contiguous() + fc2_weights = ((fc2_codes_kn[:, 1::2] << 4) | fc2_codes_kn[:, 0::2])[None] + + row_codes = torch.where( + (torch.arange(inter_size, device=device) % 2) == 0, + torch.tensor(126, dtype=torch.uint8, device=device), + torch.tensor(128, dtype=torch.uint8, device=device), + ) + fc1_block_scales = row_codes[None, :, None].expand(num_experts, inter_size, hidden_size // 32).contiguous() + fc2_block_scales = torch.full( + (num_experts, hidden_size, inter_size // 32), 127, dtype=torch.uint8, device=device + ) + global_scale = torch.ones(num_experts, dtype=torch.float32, device=device) + + onnx_model = create_fp4_moe_onnx_graph( + num_tokens=num_tokens, + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + onnx_dtype=onnx_dtype, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + fc1_block_scales=fc1_block_scales, + fc1_global_scale=global_scale, + fc2_block_scales=fc2_block_scales, + fc2_global_scale=global_scale, + use_swiglu=False, + ) + + input_tensor = torch.full((num_tokens, hidden_size), 1.0 / hidden_size, device=device, dtype=torch.float16) + router_logits = torch.ones((num_tokens, num_experts), device=device, dtype=torch.float16) + + def run(enable_native): + old_cutlass = os.environ.get("ORT_ENABLE_FP4_CUTLASS_GEMM") + old_unsafe = os.environ.get("ORT_ENABLE_FP4_CUTLASS_UNSAFE") + os.environ["ORT_ENABLE_FP4_CUTLASS_GEMM"] = "1" if enable_native else "0" + if enable_native: + os.environ["ORT_ENABLE_FP4_CUTLASS_UNSAFE"] = "1" + else: + os.environ.pop("ORT_ENABLE_FP4_CUTLASS_UNSAFE", None) + try: + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + session = onnxruntime.InferenceSession( + onnx_model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + finally: + if old_cutlass is None: + os.environ.pop("ORT_ENABLE_FP4_CUTLASS_GEMM", None) + else: + os.environ["ORT_ENABLE_FP4_CUTLASS_GEMM"] = old_cutlass + if old_unsafe is None: + os.environ.pop("ORT_ENABLE_FP4_CUTLASS_UNSAFE", None) + else: + os.environ["ORT_ENABLE_FP4_CUTLASS_UNSAFE"] = old_unsafe + + output_tensor = torch.empty((num_tokens, hidden_size), device=device, dtype=torch.float16) + iobinding = session.io_binding() + iobinding.bind_input("input", "cuda", 0, onnx_dtype, input_tensor.shape, input_tensor.data_ptr()) + iobinding.bind_input("router_probs", "cuda", 0, onnx_dtype, router_logits.shape, router_logits.data_ptr()) + iobinding.bind_output("output", "cuda", 0, onnx_dtype, output_tensor.shape, output_tensor.data_ptr()) + session.run_with_iobinding(iobinding) + iobinding.synchronize_outputs() + return output_tensor.float().cpu() + + fallback_output = run(enable_native=False) + native_output = run(enable_native=True) + max_diff = (native_output - fallback_output).abs().max().item() + self.assertEqual(max_diff, 0.0) + # ============================================================================ # Standalone packing utility tests diff --git a/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py new file mode 100644 index 0000000000000..b5a5edb14d900 --- /dev/null +++ b/onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py @@ -0,0 +1,638 @@ +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +# +# Tests for QMoE NVFP4 quantization on CUDA — W4A16 mode. +# +# NVFP4 format: E2M1 4-bit weights (2 values per byte, same as MXFP4), block +# size 16, per-block scales stored as Float8E4M3FN, per-expert float32 global +# scale (weight_scale_2). Dequant: +# w = DecodeFp4E2M1(code) * DecodeE4M3(block_scale) * global_scale[expert] +# +# Two decode paths are exercised: the dequant-to-A16 fallback (native block-scaled +# CUTLASS GEMM is Blackwell-only, so this is the general path) and, for small-decode +# SwiGLU shapes, the fused FP4 GEMV kernel (forced on via ORT_ENABLE_FP4_GEMV=1; a +# gemv_mode="0" companion checks the fallback on the same shape). Both paths are +# SM-agnostic; the tests require SM80+ / CUDA / an ENABLE_FP4 + USE_FP4_QMOE build. +# -------------------------------------------------------------------------- + +import os +import unittest + +import numpy +import torch +import torch.nn.functional as F +from cuda_plugin_ep_helper import resolve_cuda_plugin_ep +from onnx import helper +from parameterized import parameterized + +import onnxruntime + +try: + from onnx import TensorProto + + has_onnx = True +except ImportError: + has_onnx = False + +onnxruntime.preload_dlls() + +build_info = onnxruntime.get_build_info() +has_fp4_qmoe = ", fp4-qmoe=" in build_info + +device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + +torch.manual_seed(42) +numpy.random.seed(42) + +# ============================================================================ +# NVFP4 (E2M1 + E4M3 block scale) quantization utilities +# ============================================================================ + +# Positive FP4 e2m1 representable values (codes 0-7). Negative uses codes 8-15. +FP4_POS_VALUES = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) +FP4_MAX = 6.0 +E4M3_MAX = 448.0 +NVFP4_BLOCK_SIZE = 16 + + +def fp4_e2m1_quantize(values): + """Quantize float values to nearest FP4 e2m1 value; return (quantized, 4-bit codes).""" + dev = values.device + pos_vals = FP4_POS_VALUES.to(device=dev, dtype=torch.float32) + + flat = values.float().reshape(-1) + sign = flat.sign() + abs_val = flat.abs().clamp(max=FP4_MAX) + + diffs = (abs_val.unsqueeze(-1) - pos_vals.unsqueeze(0)).abs() + nearest_idx = diffs.argmin(dim=-1) # code 0-7 + + quantized = sign * pos_vals[nearest_idx] + + codes = nearest_idx.to(torch.uint8) + codes[sign < 0] += 8 + codes[flat == 0] = 0 + + return quantized.reshape(values.shape), codes.reshape(values.shape) + + +def quantize_weight_to_nvfp4(weight, block_size=NVFP4_BLOCK_SIZE): + """ + Quantize a per-expert weight matrix [N, K] to NVFP4. + + Two-level scaling: a per-expert float32 global scale (weight_scale_2) plus a + per-block Float8E4M3FN scale. The kernel reconstructs each element as + fp4_value * decode_e4m3(block_scale) * global_scale. + + Returns: + packed_col_major: [K, N//2] uint8 — column-major packed FP4 codes + block_scale_bytes: [N, K//block_size] uint8 — raw Float8E4M3FN bytes + global_scale: float scalar (per expert) + dequantized: [N, K] float — reference dequantized weights + """ + n, k = weight.shape + assert k % block_size == 0, f"K={k} must be divisible by block_size={block_size}" + assert n % 2 == 0, f"N={n} must be even for FP4 packing" + + w = weight.float() + num_blocks = k // block_size + blocks = w.reshape(n, num_blocks, block_size) + + block_amax = blocks.abs().amax(dim=-1) # [N, num_blocks] + global_amax = w.abs().amax().item() + + # Per-expert global scale (weight_scale_2). Chosen so that per-block E4M3 + # scales land within [0, E4M3_MAX]. + global_scale = global_amax / (FP4_MAX * E4M3_MAX) + if global_scale <= 0: + global_scale = 1.0 + + # Ideal per-block scale before E4M3 quantization. + block_scale_f = block_amax / FP4_MAX / global_scale # [N, num_blocks], all <= E4M3_MAX + block_scale_e4m3 = block_scale_f.to(torch.float8_e4m3fn) + block_scale_deq = block_scale_e4m3.float() # DecodeE4M3(byte) + + # Effective per-element scale (matches the kernel): block_scale_deq * global_scale. + eff_scale = (block_scale_deq * global_scale).unsqueeze(-1) # [N, num_blocks, 1] + eff_scale_safe = torch.where(eff_scale > 0, eff_scale, torch.ones_like(eff_scale)) + + scaled = blocks / eff_scale_safe + quantized_vals, fp4_codes = fp4_e2m1_quantize(scaled) + quantized_vals = quantized_vals.reshape(n, num_blocks, block_size) + fp4_codes = fp4_codes.reshape(n, num_blocks, block_size) + + dequantized = (quantized_vals * eff_scale).reshape(n, k) + + # Pack codes [N, K] -> transpose -> [K, N] -> pack pairs along N -> [K, N//2]. + codes_nk = fp4_codes.reshape(n, k) + codes_kn = codes_nk.T.contiguous() # [K, N] + low = codes_kn[:, 0::2].to(torch.uint8) # even N-index -> low nibble + high = codes_kn[:, 1::2].to(torch.uint8) # odd N-index -> high nibble + packed = (high << 4) | low # [K, N//2] + + block_scale_bytes = block_scale_e4m3.view(torch.uint8).contiguous() # [N, K//block_size] + + return packed, block_scale_bytes, float(global_scale), dequantized + + +# ============================================================================ +# SwiGLU activation reference +# ============================================================================ + + +def swiglu_ref(x, alpha=1.702, limit=7.0): + """SwiGLU activation matching the QMoE kernel implementation.""" + dim = x.shape[-1] + x = x.view(-1, dim // 2, 2) + g, l_val = x[..., 0], x[..., 1] + if limit is not None: + g = g.clamp(max=limit) + l_val = l_val.clamp(min=-limit, max=limit) + return g * torch.sigmoid(alpha * g) * (l_val + 1) + + +# ============================================================================ +# ONNX graph builder for NVFP4 QMoE +# ============================================================================ + + +def create_nvfp4_moe_onnx_graph( + num_tokens, + hidden_size, + inter_size, + num_experts, + top_k, + onnx_dtype, + fc1_weights, # [E, K1, N1/2] uint8 packed FP4 (column-major) + fc2_weights, # [E, K2, N2/2] uint8 packed FP4 (column-major) + fc1_block_scales, # [E, N1, K1//16] uint8 (Float8E4M3FN bytes) + fc1_global_scale, # [E] float32 + fc2_block_scales, # [E, N2, K2//16] uint8 (Float8E4M3FN bytes) + fc2_global_scale, # [E] float32 + block_size=NVFP4_BLOCK_SIZE, + use_swiglu=False, +): + """Build ONNX model with QMoE operator in NVFP4 mode.""" + inputs = [ + "input", # 0 + "router_probs", # 1 + "fc1_weights", # 2: uint8 packed FP4 + "fc1_scales", # 3: Float8E4M3FN NVFP4 block scales + "", # 4: fc1_bias + "fc2_weights", # 5: uint8 packed FP4 + "fc2_scales", # 6: Float8E4M3FN NVFP4 block scales + "", # 7: fc2_bias + "", # 8: fc3_weights + "", # 9: fc3_scales + "", # 10: fc3_bias + "", # 11: fc1_zero_points + "", # 12: fc2_zero_points + "", # 13: fc3_zero_points + "", # 14: router_weights + "fc1_global_scale", # 15 + "fc2_global_scale", # 16 + ] + + activation = "swiglu" if use_swiglu else "silu" + + nodes = [ + helper.make_node( + "QMoE", + inputs, + ["output"], + "QMoE_NVFP4", + k=top_k, + normalize_routing_weights=1, + activation_type=activation, + expert_weight_bits=4, + quant_type="nvfp4", + block_size=block_size, + swiglu_fusion=1 if use_swiglu else 0, + swiglu_limit=7.0, + activation_alpha=1.702, + activation_beta=1.0, + domain="com.microsoft", + ), + ] + + initializers = [] + + # FC1 / FC2 packed weights [E, K, N/2] uint8 + for name, tensor in [("fc1_weights", fc1_weights), ("fc2_weights", fc2_weights)]: + arr = numpy.ascontiguousarray(tensor.cpu().numpy().astype(numpy.uint8)) + initializers.append(helper.make_tensor(name, TensorProto.UINT8, list(tensor.shape), arr.tobytes(), raw=True)) + + # NVFP4 block scales [E, N, K//16] float8e4m3fn (stored as raw bytes) + for name, tensor in [("fc1_scales", fc1_block_scales), ("fc2_scales", fc2_block_scales)]: + arr = numpy.ascontiguousarray(tensor.cpu().numpy().astype(numpy.uint8)) + initializers.append( + helper.make_tensor(name, TensorProto.FLOAT8E4M3FN, list(tensor.shape), arr.tobytes(), raw=True) + ) + + # Per-expert global scales [E] float32 (T4) + for name, tensor in [("fc1_global_scale", fc1_global_scale), ("fc2_global_scale", fc2_global_scale)]: + vals = tensor.cpu().float().flatten().tolist() + initializers.append(helper.make_tensor(name, TensorProto.FLOAT, list(tensor.shape), vals, raw=False)) + + graph_inputs = [ + helper.make_tensor_value_info("input", onnx_dtype, [num_tokens, hidden_size]), + helper.make_tensor_value_info("router_probs", onnx_dtype, [num_tokens, num_experts]), + ] + graph_outputs = [ + helper.make_tensor_value_info("output", onnx_dtype, [num_tokens, hidden_size]), + ] + + graph = helper.make_graph(nodes, "QMoE_NVFP4_Test", graph_inputs, graph_outputs, initializers) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 22), helper.make_opsetid("com.microsoft", 1)], + ) + return model.SerializeToString() + + +# ============================================================================ +# Test class +# ============================================================================ + + +def _cuda_sm(): + """Return SM version (e.g. 90 for Hopper).""" + if not torch.cuda.is_available(): + return 0 + cc = torch.cuda.get_device_capability() + return cc[0] * 10 + cc[1] + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available") +@unittest.skipIf(not has_onnx, "ONNX not available") +@unittest.skipIf(not has_fp4_qmoe, "CUDA QMoE FP4 kernels not enabled in this build") +@unittest.skipIf(not hasattr(torch, "float8_e4m3fn"), "PyTorch build does not expose torch.float8_e4m3fn") +class TestQMoENVFP4(unittest.TestCase): + """Tests for W4A16 NVFP4 MoE quantization (dequant fallback).""" + + def _skip_if_no_fp4(self): + # NVFP4 always uses the dequant-to-A16 fallback (native block-scaled CUTLASS GEMM is + # Blackwell-only). That path is SM-agnostic and runs on any CUDA GEMM-capable GPU, so + # SM80 (A100) is sufficient here even though the production target is H200/SM90. + sm = _cuda_sm() + if sm < 80: + self.skipTest(f"NVFP4 QMoE requires SM80+, got SM{sm}") + + def _run_nvfp4_moe_test( + self, + hidden_size, + inter_size, + num_experts, + top_k, + num_tokens, + onnx_dtype, + use_swiglu=False, + block_size=NVFP4_BLOCK_SIZE, + gemv_mode=None, + ): + self._skip_if_no_fp4() + + torch.manual_seed(42) + numpy.random.seed(42) + + torch_dtype = torch.float16 if onnx_dtype == TensorProto.FLOAT16 else torch.bfloat16 + onnx_elem = TensorProto.FLOAT16 if torch_dtype == torch.float16 else TensorProto.BFLOAT16 + + fc1_n = 2 * inter_size if use_swiglu else inter_size + fc1_k = hidden_size + fc2_n = hidden_size + fc2_k = inter_size + + fc1_packed, fc1_bs, fc1_gs, fc1_deq = [], [], [], [] + fc2_packed, fc2_bs, fc2_gs, fc2_deq = [], [], [], [] + + for _ in range(num_experts): + w1 = torch.randn(fc1_n, fc1_k, device=device) * 0.1 + p1, b1, g1, d1 = quantize_weight_to_nvfp4(w1, block_size) + fc1_packed.append(p1) + fc1_bs.append(b1) + fc1_gs.append(torch.tensor(g1, dtype=torch.float32)) + fc1_deq.append(d1) + + w2 = torch.randn(fc2_n, fc2_k, device=device) * 0.1 + p2, b2, g2, d2 = quantize_weight_to_nvfp4(w2, block_size) + fc2_packed.append(p2) + fc2_bs.append(b2) + fc2_gs.append(torch.tensor(g2, dtype=torch.float32)) + fc2_deq.append(d2) + + fc1_weights = torch.stack(fc1_packed, dim=0) # [E, K, N/2] + fc2_weights = torch.stack(fc2_packed, dim=0) # [E, K, N/2] + fc1_block_scales = torch.stack(fc1_bs, dim=0) # [E, N, K//16] + fc2_block_scales = torch.stack(fc2_bs, dim=0) # [E, N, K//16] + fc1_global_scale = torch.stack(fc1_gs) # [E] + fc2_global_scale = torch.stack(fc2_gs) # [E] + fc1_deq_all = torch.stack(fc1_deq, dim=0) # [E, N, K] + fc2_deq_all = torch.stack(fc2_deq, dim=0) # [E, N, K] + + onnx_model = create_nvfp4_moe_onnx_graph( + num_tokens=num_tokens, + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + onnx_dtype=onnx_elem, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + fc1_block_scales=fc1_block_scales, + fc1_global_scale=fc1_global_scale, + fc2_block_scales=fc2_block_scales, + fc2_global_scale=fc2_global_scale, + block_size=block_size, + use_swiglu=use_swiglu, + ) + + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + # gemv_mode toggles the fused FP4 GEMV decode path (read once in the QMoE op ctor during + # session creation): "1" forces it on, "0" forces the dequant fallback, None leaves the + # default. Restore the previous value right after the session is built. + prev_gemv_env = os.environ.get("ORT_ENABLE_FP4_GEMV") + if gemv_mode is not None: + os.environ["ORT_ENABLE_FP4_GEMV"] = gemv_mode + try: + session = onnxruntime.InferenceSession( + onnx_model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + except Exception as e: + if "ENABLE_FP4" in str(e) or "requires USE_FP4_QMOE" in str(e): + self.skipTest(f"NVFP4 not supported in this build: {e}") + raise + finally: + if gemv_mode is not None: + if prev_gemv_env is None: + os.environ.pop("ORT_ENABLE_FP4_GEMV", None) + else: + os.environ["ORT_ENABLE_FP4_GEMV"] = prev_gemv_env + + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch_dtype) + router_logits = torch.randn(num_tokens, num_experts, device=device, dtype=torch_dtype) + output_tensor = torch.zeros(num_tokens, hidden_size, device=device, dtype=torch_dtype) + + iobinding = session.io_binding() + iobinding.bind_input("input", "cuda", 0, onnx_elem, input_tensor.shape, input_tensor.data_ptr()) + iobinding.bind_input("router_probs", "cuda", 0, onnx_elem, router_logits.shape, router_logits.data_ptr()) + iobinding.bind_output("output", "cuda", 0, onnx_elem, output_tensor.shape, output_tensor.data_ptr()) + + iobinding.synchronize_inputs() + try: + session.run_with_iobinding(iobinding) + except Exception as e: + msg = str(e) + if "ENABLE_FP4" in msg or "requires USE_FP4_QMOE" in msg or "stubbed out" in msg: + self.skipTest(f"NVFP4 kernel not available in this build: {e}") + raise + iobinding.synchronize_outputs() + + ort_output = output_tensor.clone() + + ref_output = self._compute_reference( + input_tensor, + router_logits, + fc1_deq_all, + fc2_deq_all, + num_experts, + top_k, + use_swiglu, + torch_dtype, + ) + + max_diff = (ort_output.float() - ref_output.float()).abs().max().item() + dtype_tag = "FP16" if torch_dtype == torch.float16 else "BF16" + act_tag = "SwiGLU" if use_swiglu else "SiLU" + print( + f"NVFP4 MoE test: {dtype_tag} {act_tag} " + f"tokens={num_tokens} experts={num_experts} " + f"hidden={hidden_size} inter={inter_size} " + f"max_diff={max_diff:.6f}" + ) + + atol = 0.15 if torch_dtype == torch.bfloat16 else 0.12 + self.assertLess( + max_diff, + atol, + f"NVFP4 MoE parity check failed: max_diff={max_diff:.6f} > atol={atol}", + ) + + def _assert_invalid_nvfp4_model( + self, block_size=NVFP4_BLOCK_SIZE, truncate_fc1_scales=False, truncate_fc1_global_scale=False + ): + self._skip_if_no_fp4() + num_experts = 2 + hidden_size = 64 + inter_size = 64 + fc1_weights = torch.zeros(num_experts, hidden_size, inter_size // 2, dtype=torch.uint8) + fc2_weights = torch.zeros(num_experts, inter_size, hidden_size // 2, dtype=torch.uint8) + fc1_scales = torch.zeros(num_experts, inter_size, hidden_size // NVFP4_BLOCK_SIZE, dtype=torch.uint8) + fc2_scales = torch.zeros(num_experts, hidden_size, inter_size // NVFP4_BLOCK_SIZE, dtype=torch.uint8) + if truncate_fc1_scales: + fc1_scales = fc1_scales[:, :, :-1] + model = create_nvfp4_moe_onnx_graph( + num_tokens=1, + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=1, + onnx_dtype=TensorProto.FLOAT16, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + fc1_block_scales=fc1_scales, + fc1_global_scale=torch.ones(num_experts - int(truncate_fc1_global_scale), dtype=torch.float32), + fc2_block_scales=fc2_scales, + fc2_global_scale=torch.ones(num_experts, dtype=torch.float32), + block_size=block_size, + ) + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + with self.assertRaisesRegex(Exception, "block_size|fc1_scales|fc1_global_scale"): + session = onnxruntime.InferenceSession( + model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + session.run( + None, + { + "input": numpy.zeros((1, hidden_size), dtype=numpy.float16), + "router_probs": numpy.zeros((1, num_experts), dtype=numpy.float16), + }, + ) + + def test_nvfp4_rejects_wrong_block_size(self): + self._assert_invalid_nvfp4_model(block_size=32) + + def test_nvfp4_rejects_malformed_prepacked_block_scales(self): + self._assert_invalid_nvfp4_model(truncate_fc1_scales=True) + + def test_nvfp4_rejects_malformed_prepacked_global_scale(self): + self._assert_invalid_nvfp4_model(truncate_fc1_global_scale=True) + + @staticmethod + def _compute_reference(input_tensor, router_logits, fc1_deq, fc2_deq, num_experts, top_k, use_swiglu, torch_dtype): + """Reference MoE forward pass using dequantized weights.""" + num_tokens = input_tensor.shape[0] + hidden_size = input_tensor.shape[1] + + x = input_tensor.float() + logits = router_logits.float() + + topk_vals, topk_idx = torch.topk(logits, top_k, dim=-1) + routing_weights = F.softmax(topk_vals, dim=1) + + output = torch.zeros(num_tokens, hidden_size, device=x.device, dtype=torch.float32) + expert_mask = F.one_hot(topk_idx, num_classes=num_experts).permute(2, 1, 0) + + for e in range(num_experts): + idx, top_x = torch.where(expert_mask[e]) + if top_x.shape[0] == 0: + continue + + tokens = x[top_x] + w1 = fc1_deq[e].float() + w2 = fc2_deq[e].float() + + h = tokens @ w1.T + h = swiglu_ref(h) if use_swiglu else F.silu(h) + h = h @ w2.T + h = h * routing_weights[top_x, idx, None] + + output.index_add_(0, top_x, h) + + return output.to(torch_dtype) + + # ================================================================ + # Test cases + # ================================================================ + + def test_nvfp4_fp16_silu_basic(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.FLOAT16, + ) + + def test_nvfp4_bf16_silu_basic(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.BFLOAT16, + ) + + def test_nvfp4_fp16_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + ) + + def test_nvfp4_bf16_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.BFLOAT16, + use_swiglu=True, + ) + + @parameterized.expand( + [ + (8,), + (64,), + (128,), + ] + ) + def test_nvfp4_fp16_token_counts(self, num_tokens): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=4, + top_k=2, + num_tokens=num_tokens, + onnx_dtype=TensorProto.FLOAT16, + ) + + def test_nvfp4_fp16_more_experts(self): + self._run_nvfp4_moe_test( + hidden_size=64, + inter_size=64, + num_experts=8, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.FLOAT16, + ) + + def test_nvfp4_fp16_larger_dims(self): + self._run_nvfp4_moe_test( + hidden_size=128, + inter_size=256, + num_experts=4, + top_k=2, + num_tokens=32, + onnx_dtype=TensorProto.FLOAT16, + ) + + # ================================================================ + # Fused FP4 GEMV decode fast path (block size 16). The GEMV support window requires + # n, k >= 512 and expanded rows (num_tokens * top_k) <= 8, plus SwiGLU fusion, so these + # decode-shaped SwiGLU cases route through the NVFP4 GEMV kernel (gemv_mode="1"). The + # gemv_mode="0" companion forces the dequant fallback on the identical shape; both must + # match the exact dequantized reference. + # ================================================================ + + def test_nvfp4_fp16_gemv_decode_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=4, + top_k=2, + num_tokens=2, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="1", + ) + + def test_nvfp4_bf16_gemv_decode_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=4, + top_k=2, + num_tokens=2, + onnx_dtype=TensorProto.BFLOAT16, + use_swiglu=True, + gemv_mode="1", + ) + + def test_nvfp4_fp16_gemv_disabled_swiglu(self): + self._run_nvfp4_moe_test( + hidden_size=512, + inter_size=512, + num_experts=4, + top_k=2, + num_tokens=2, + onnx_dtype=TensorProto.FLOAT16, + use_swiglu=True, + gemv_mode="0", + ) + + +if __name__ == "__main__": + unittest.main()