Back to Onnxruntime

MoE and QMoE — CUDA Operator Documentation

docs/contrib_ops/cuda/moe_qmoe.md

1.29.086.3 KB
Original Source

MoE and QMoE — CUDA Operator Documentation

This document describes the design, schema, kernel dispatch, weight formats, and current implementation status of the MoE (com.microsoft::MoE) and QMoE (com.microsoft::QMoE) operators on the CUDA execution provider.

The CUTLASS kernels are derived from TensorRT-LLM (CUTLASS 4.4.2, commit 346018db87) and have been significantly modified for ONNX Runtime — see §16 Differences vs. TensorRT-LLM.


Table of Contents

  1. Overview & Operator Set
  2. Operator Schema
  3. Quantization Modes
  4. Architecture Dispatch & Kernel Paths
  5. PrePack Transformations
  6. Weight Formats
  7. Cross-Architecture Packing Compatibility
  8. SwiGLU Fusion
  9. FP4 (MXFP4) Details
  1. FP8 (W8A16) Details
  2. WFP4AFP8 Details
  3. Future / Deferred Modes
  1. Testing
  2. Build Configuration
  1. Limitations & Known Issues
  2. Differences vs. TensorRT-LLM

1. Overview & Operator Set

Two contrib ops are registered in the com.microsoft domain:

OperatorPurposeSource
MoEStandard (non-quantized) Mixture-of-Experts. FP16/BF16/FP32 weights.onnxruntime/contrib_ops/cuda/moe/moe.cc
QMoEQuantized Mixture-of-Experts. INT4/INT8/FP8/MXFP4 weights, FP16/BF16/FP8 activations.onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc

Both ops share the same CUTLASS-based runner (CutlassMoeFCRunner), routing engine, and sort/permute infrastructure. They differ only in how the weight tensors are interpreted.

The execution pipeline is:

input tokens → router (top-k softmax) → permute by expert
  → GEMM1 (per-expert) → activation (SiLU/GeLU/ReLU/SwiGLU)
  → GEMM2 (per-expert) → un-permute → weighted sum

2. QMoE Operator Schema

2.1 Attributes

AttributeTypeDefaultDescription
kint1Top-K experts selected per token.
activation_typestring"relu"One of "relu", "gelu", "silu", "swiglu", "identity". These are the only values accepted end-to-end (attribute parsing); other kernel-internal types are not reachable from ONNX.
normalize_routing_weightsint0Re-normalize the top-k weights to sum to 1.
use_sparse_mixerint0Enable sparse-mixer routing variant.
swiglu_fusionint00=no fusion, 1=interleaved (Gate/Value), 2=block (Gate;Value). See §8.
activation_alphafloat1.0SwiGLU alpha. Default 1.0 (Standard SwiGLU); GPT-OSS uses 1.702.
activation_betafloat0.0SwiGLU beta. Default 0.0 (Standard SwiGLU); GPT-OSS uses 1.0.
swiglu_limitfloatunset (+inf)SwiGLU clamp limit. Unset means no clamp (Standard SwiGLU); GPT-OSS uses 7.0.
expert_weight_bits (QMoE only)int44 (INT4/MXFP4) or 8 (INT8/FP8).
block_size (QMoE only)int-1Group 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.
weights_prepacked (QMoE only)int-1Tri-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.

2.2 Type Constraints

ConstraintAllowed TypesUsed By
Tfloat, float16, bfloat16input, output, biases, router
T1uint8, float8e4m3fnquantized weights and zero points: INT4/INT8/FP4/NVFP4 weights use uint8; FP8 weights use float8e4m3fn
T2float, float16, bfloat16, uint8, float8e4m3fnINT4/INT8 weight scales use floating-point tensors; MXFP block scales use uint8 (float8e8m0) storage; NVFP4 block scales use float8e4m3fn
T4floatper-expert global scales, FP8 activation scales

2.3 Inputs

The schema is unified across all quant_type values. Inputs that are not relevant to the selected quant_type are simply omitted (most are Optional).

IdxNameTypeShapeUsed by quant_type
0inputT(num_tokens, hidden_size)all
1router_probsT(num_tokens, num_experts)all
2fc1_experts_weightsT1(E, fusion×inter, hidden/pack)all
3fc1_scalesT2 (Opt)varies — see §2.4int, fp4, nvfp4, wfp4afp8
4fc1_experts_biasT (Opt)(E, fusion×inter)optional
5fc2_experts_weightsT1(E, hidden, inter/pack)all
6fc2_scalesT2 (Opt)variesint, fp4, nvfp4, wfp4afp8
7fc2_experts_biasT (Opt)(E, hidden)optional
8fc3_experts_weightsT1 (Opt)(E, inter, hidden/pack)optional (SwiGLU split-weight)
9fc3_scalesT2 (Opt)variesoptional
10fc3_experts_biasT (Opt)(E, inter)optional
11fc1_zero_pointsT1 (Opt)matches fc1_scalesint only
12fc2_zero_pointsT1 (Opt)matches fc2_scalesint only
13fc3_zero_pointsT1 (Opt)matches fc3_scalesoptional, int only
14router_weightsT (Opt)(num_tokens, num_experts)optional (DeepSeek noaux_tc)
15fc1_global_scaleT4 (Opt)(num_experts,)fp4, nvfp4, fp8, wfp4afp8
16fc2_global_scaleT4 (Opt)(num_experts,)fp4, nvfp4, fp8, wfp4afp8
17fc1_act_scaleT4 (Opt)(1,) or (num_experts,)wfp4afp8 (Variant A)
18fc2_act_scaleT4 (Opt)(1,) or (num_experts,)wfp4afp8 (Variant A)
19fc1_act_block_scaleT2 (Opt, float8e8m0)(E, M_pad, K/32)wfp4afp8 (Variant B)
20fc2_act_block_scaleT2 (Opt, float8e8m0)(E, M_pad, inter/32)wfp4afp8 (Variant B)

E = num_experts. pack = 8 / expert_weight_bits for INT/MXFP4 weights; pack = 1 for FP8 weights. fusion = 2 for swiglu_fusion=1, otherwise 1.

router_weights (input 14) enables DeepSeek-style routing where router_probs is used only for top-K selection and router_weights provides the mixing weights gathered at the selected expert indices. When omitted, router_probs is used for both (backward compatible).

2.4 Input 3/6/9 interpretation by quant_type

quant_typedtypeShapeSemantics
"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 (both MXFP4 and NVFP4) are symmetric formats with no zero-point.


3. Quantization Modes

quant_typeNotationActivationWeightNative SMFallbackBuild gate
"int" (4-bit)W4A16FP16/BF16INT4 group-wiseSM75+ (Ampere GemmGrouped)always
"int" (8-bit)W8A16FP16/BF16INT8 group-wiseSM75+always
"fp8"W8A16-fp8BF16/FP16FP8 e4m3 (no packing)SM90+ nativedequant→A16 on SM<90ENABLE_FP8 (CUDA ≥ 11.8)
"fp4"W4A16-MXFP4BF16/FP16MXFP4 e2m1, group=32SM120+ nativedequant→A16 on SM<120ENABLE_FP4 + USE_FP4_QMOE (CUDA ≥ 12.8)
"nvfp4"W4A16-NVFP4BF16/FP16NVFP4 e2m1, group=16, float8e4m3fn block scale + per-expert FP32 global scaleSM120/SM121 native (block-scaled FP4×FP4 prefill) + fused GEMV decodedequant→A16 on other SMsENABLE_FP4 + USE_FP4_QMOE (CUDA ≥ 12.8)
"wfp4afp8"W4A8-MXFP4×FP8FP8 e4m3 (quantized in-runner)MXFP4 e2m1, group=32SM100+ nativedequant→A16 on SM<100ENABLE_FP4 + USE_FP4_QMOE + ENABLE_FP8

Selection logic (see moe_quantization.cc):

cpp
if (quant_type_ == "fp4")      use_fp4_dequant_fallback_      = (sm_ < 120);
if (quant_type_ == "nvfp4")    use_fp4_dequant_fallback_      = !enable_nvfp4_cutlass_gemm_;  // native SM120+ block-scaled FP4xFP4 prefill (ORT_ENABLE_NVFP4_CUTLASS_GEMM, shape-gated); fused GEMV decode still 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, nvfp4, wfp4afp8 → must be 4
  • fp8 → must be 8

When the build is configured without the corresponding flags, quant_type values that require them are rejected at construction time:

cpp
#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_ != "wfp4afp8", ...);
#endif
#if !defined(ENABLE_FP8)
  ORT_ENFORCE(quant_type_ != "wfp4afp8", "...");
#endif

4. Architecture Dispatch & Kernel Paths

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

PathCUTLASS classUsed forSM range
MoE GEMV fast pathfpA_intB_gemv-based custom kernelINT4/INT8 per-column WA16 and symmetric INT4/INT8 block-wise WA16, and MXFP4 (group-32) / NVFP4 (group-16) W4A16 decode, with FP16 or BF16 activations and true decode row countsSM80+
Ampere GemmGroupedcutlass::gemm::kernel::GemmGroupedINT4/INT8 W*A16, FP8 W8A16 dequant fallback, FP32SM75–SM89, plus all mixed-input on SM90/SM120
TMA Warp-Specialized (mixed-input)CollectiveBuilderMixedInputSame-type FP16×FP16 / BF16×BF16, native MXFP4 W4A16SM90 (same-type), SM120 (FP4 W4A16)
Block-Scaled Tensor OpOpClassBlockScaledTensorOpNative FP8×MXFP4 (wfp4afp8)SM100+ (Blackwell)

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

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

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

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

4.1 Per-mode dispatch matrix

ModeSM75-89 (Ampere/Ada)SM90 (Hopper)SM100 (Blackwell)SM120 (RTX 5090)
INT4/INT8 W*A16Ampere GemmGroupedAmpere GemmGrouped (TMA WS rejects mixed-type INT)Ampere GemmGroupedAmpere GemmGrouped
FP16/BF16 (no quant, MoE op)Ampere GemmGroupedTMA WS (same-type)TMA WS / valid Blackwell specTMA WS / Ampere fallback
FP8 W8A16 nativedequant fallbackTMA WSTMA WSSM89 FP8 kernel redirect
FP4 W4A16 nativedequant fallbackdequant fallbackdequant fallbackTMA WS mixed-input FP4
NVFP4 W4A16 (group-16)dequant fallback + fused GEMV decodedequant fallback + fused GEMV decodedequant fallback + fused GEMV decodeTMA WS block-scaled FP4×FP4 prefill + fused GEMV decode
WFP4AFP8 nativedequant fallbackdequant fallbackBlock-scaled tensor opBlock-scaled tensor op
FP32Ampere GemmGrouped (forced)samesamesame

4.2 Minimum dimension constraint (min_dim)

  • Both hidden_size and inter_size must be ≥ 16.
  • TMA WS path: smallest tile is 128×16×128B (N=16 for FP16). K residues handled by TMA.
  • Ampere GemmGrouped path: smallest instantiated tile N=128, but CUTLASS predicates N < tile_N.
  • Alignment to 128 bits is enforced separately (e.g., dimensions must be multiples of 8 for FP16).

4.3 Dequant-to-A16 fallback

When the requested native path is not available on the running GPU, the QMoE op decodes the quantized weights into FP16/BF16 once and feeds them to the dense A16 runner. Helper kernels:

  • LaunchQMoEDequantizeFp4Weights — MXFP4 → FP16/BF16
  • LaunchQMoEDequantizeFp8Weights — FP8 e4m3 → FP16/BF16

The decoded buffers are owned by the QMoE op for the lifetime of the session.

MXFP4 exception. With ORT_FP4_SM80_GEMM=1 (default on SM80–SM119) both prefill and decode are served by pre-packed e2m1 buffers, so LaunchQMoEDequantizeFp4Weights is never reached and the raw initializers are released in PrePack (§9.11). Set ORT_FP4_SM80_GEMM=0 to restore the dequant fallback.

4.4 Target hardware (developer matrix)

RTX 3090 (SM86), RTX 4090 (SM89), H200 (SM90), GB200/B200 (SM100), RTX 5090 (SM120).


5. PrePack Transformations

QMoE::PrePack (moe_quantization.cc) copies constant inputs to GPU once and, for INT4/INT8, derives a pre-scaled bias from the zero points so the kernel can apply asymmetric quantization with no extra subtraction.

5.1 Weights (input 2 / 5 / 8)

INT4/INT8 weight layout is controlled by the weights_prepacked attribute (§2.1). The prepacked layouts selected by -1 and 1 are determined by the execution provider:

  • weights_prepacked=-1 (default) — the fc1/fc2 weights are already in the EP's default prepacked layout (e.g. packed offline by pack_weights_for_cuda_mixed_gemm for the CUDA EP). They are copied to GPU and consumed as-is.
  • weights_prepacked=1 — the fc1/fc2 weights are already in the EP's SM90 (Hopper) prepacked layout (reserved; see the note below).
  • weights_prepacked=0 — the fc1/fc2 weights are raw, schema-conformant [E, N, K/pack] tensors as produced by quantize_matmul_{4,8}bits. PrePack runs the CUTLASS layout transform itself via PrePackIntExpertWeights, removing the offline pre-pack dependency. This makes integer QMoE symmetric with MatMulNBits::PrePack_B.

Single layout on the CUDA EP. The CUDA EP INT4/INT8 MoE GEMM always dispatches to the Ampere (SM80) grouped-GEMM kernel — even on SM90 — because mixed int-weight + fp16/bf16 activation is not a valid Hopper TMA warp-specialized specialisation (isValidHopperMOESpecialisation is false). This matches TensorRT-LLM, which likewise routes W4A16/W8A16 MoE to the SM80 kernel on Hopper; its Hopper TMA-WS mixed-dtype MoE kernel is reserved for W4A8 (FP8 activation) and WFP4A16 (FP4 weight). Consequently the CUDA EP consumes the SM80 fpA_intB layout on every GPU, PrePack always packs for SM80, and weights_prepacked=-1 and =1 are equivalent today. 1 is accepted and reserved for a possible future Hopper-specific layout (e.g. W4A8). There is therefore no architecture-match constraint: SM80-format weights run correctly on SM90 via the SM80 kernel.

PrePackIntExpertWeights loops over the E experts and, per expert, applies the same transpose + row-permutation / column-interleave / bias / pair-interleave transform as pack_weights_for_cuda_mixed_gemm (see §6.1), always targeting the SM80 layout. SM75+ is required. The source [E, N, K/pack] initializers are released after their shapes are cached (fc1_weights_shape_ / fc2_weights_shape_), so peak weight memory stays ~1×. The prepacked GPU buffers (packed_fc1_weights_ / packed_fc2_weights_) are then preferred by ComputeInternal. If prepacking is disabled at the session level (session.disable_prepacking), the buffers stay null and the raw initializer pointers are read at compute time instead.

Note: weights_prepacked=0 is the only path that triggers an in-PrePack layout transform for INT weights. FP4 / FP8 / WFP4AFP8 weight handling is unaffected.

MXFP4 weights must be packed by pack_fp4_weights_for_cuda_moe_gemm. FP8 weights are stored as raw e4m3 bytes (no packing).

MXFP4 on SM80–SM119. When ORT_FP4_SM80_GEMM is on (the default), PrePack produces a single pre-packed e2m1 buffer per FC that serves both the grouped-GEMM prefill and the fused GEMV decode, and reports is_packed = true so ORT releases the raw [E, K, N/2] initializers. Peak MXFP4 weight memory is ~1×; without this it was ~3× (raw initializer + prefill layout + decode layout). See §9.11.

5.2 INT4/INT8 scales + zero-point → bias

The kernels use a pre-calculated additive bias to avoid the per-element zero-point subtraction.

  • 8-bit: weights are shifted uint8 → int8 (− 128). Bias compensates:
    bias = (128 − ZP) × scale
    
    Effectively computes (W_stored + 128 − ZP) × scale = (W_orig − ZP) × scale.
  • 4-bit: zero points are unpacked from nibbles (2 per byte) and converted to scaled biases:
    bias = (8 − ZP) × scale
    
    Equivalent to (W − (ZP − 8)) × scale.
  • Symmetric (no fc*_zero_points): bias = 0.

Kernels: LaunchQMoEPrePackOffsetBias, LaunchQMoEPrePackPacked4BitZPKernel. Output buffer (packed_bias) has the scale dtype (float16 / bfloat16 / float).

5.3 FP8 / FP4 / WFP4AFP8 PrePack

For floating-point quantization modes, PrePack simply copies constant tensors to GPU memory:

Input idxMemberUsed by
15/16/17packed_fc{1,2,3}_global_scale_fp4, fp8, wfp4afp8
18/19packed_fc{1,2}_act_scale_wfp4afp8 (Variant A)

Block scales (inputs 3/6/9 for fp4/wfp4afp8 and 20/21 for wfp4afp8 Variant B) that are constant initializers are also copied to GPU; otherwise they are read directly from context->Input at runtime.


6. Weight Formats

This section covers the five distinct weight encodings supported by QMoE.

6.1 INT4 group-wise (quant_type="int", expert_weight_bits=4)

Logical and packed shapes

TensorLogicalPacked storage
FC1 weight[E, N, K] (N = fusion × inter)[E, N, K/2] bytes
FC1 scales[E, N, K/group_size] (T2)
FC1 zero-points (asymmetric)[E, N, K/group_size/2] packed (T1)

INT4 packing layout within a byte: [high_nibble | low_nibble] = [elt_1 | elt_0]. Each INT4 element is in [-8, 7] (signed) before bias, [0, 15] after the +8 bias.

Preprocessing pipeline (offline pack_weights_for_cuda_mixed_gemm, or in-PrePack via PrePackIntExpertWeights)

This is the layout transform applied either offline by pack_weights_for_cuda_mixed_gemm, or per-expert inside PrePack when weights_prepacked=0 (see §5.1).

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

  1. Input layout: [N, K] per expert (Out × In), 2 elements per byte for INT4.
  2. Transpose & signed conversion:
    • Unpack uint4 [0, 15] → subtract 8 → int8 [-8, 7].
    • Transpose row-major [K, N] → column-major [N, K] with nibble-level swaps.
  3. Row permutation (LDSM) for SM75+ tensor cores. INT4 uses a 32-row pattern:
    {0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27,
     4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31}
    
    (kPerm_W4_A16 in fpA_intB_gemm_preprocessors_impl.h).
  4. Column interleaving: ColumnMajorTileInterleave<64, 4> on Ampere/Ada/Blackwell. RowsPerTile=64 (K direction), ColumnsInterleaved=4 (N direction).
  5. Bias addition + register interleaving: add 128 (so storage is uint8), reorder elements within a 32-bit word from [7,6,5,4,3,2,1,0] to [7,5,3,1,6,4,2,0] to minimize shift/mask cost in the kernel (add_bias_and_interleave_int8s_inplace_kernel).

Dequantization

cpp
// Symmetric (no zero-point):  W_stored is uint8 in [0, 15]
float W = (float)(W_stored - 8) * scale;

// Asymmetric (with zero tensor):
float W = (float)W_stored * scale + zero;     // zero is the scaled bias from PrePack

6.2 INT8 group-wise (quant_type="int", expert_weight_bits=8)

TensorLogicalPacked storage
FC1 weight[E, N, K][E, N, K] (no packing)
FC1 scales[E, N, K/group_size] (T2)

Preprocessing pipeline differences from INT4:

  • Row permutation: 16-row pattern {0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15} (kPerm_W8_A16).
  • Column interleaving: ColumnMajorTileInterleave<64, 2> (RowsPerTile=64, ColumnsInterleaved=2).
  • Bias: +128 shift (signed [-128,127] → unsigned [0,255]).
  • Register interleaving: [3, 1, 2, 0] per 32-bit word.

Dequantization (symmetric): W = (W_stored - 128) * scale.

6.3 INT4 vs INT8 comparison

AspectINT4INT8
Elements per byte21
Elements per int32 word84
Value range (signed)[-8, 7][-128, 127]
Bias offset+8+128
Row permutation size32 rows16 rows
Packed shape[E, N, K/2][E, N, K]
Column interleave<64, 4><64, 2>

6.4 FP8 e4m3 (quant_type="fp8")

  • Storage: [E, N, K] float8e4m3fn (Float8E4M3FN in ORT; __nv_fp8_e4m3 in CUDA), 1 byte per value.
  • Packing: pack_size = 1 — no offline packing required.
  • Scales: per-expert global scale only — fc1_global_scale (input 15) of shape (E,), T4 float32. No block scales (inputs 3/6/9 omitted).
  • Zero-points: not applicable (FP8 is symmetric); inputs 11/12/13 must be absent.
  • Dequantization (applied in the GEMM epilogue): W_bf16 = fp8_to_bf16(W_fp8) × global_scale.

6.5 MXFP4 e2m1 (quant_type="fp4" and "wfp4afp8")

  • Storage: [E, N, K/2] uint8, reinterpreted as __nv_fp4_e2m1 (2 values per byte).
  • Packer: pack_fp4_weights_for_cuda_moe_gemm (Python binding in onnxruntime_pybind_quant.cc). No Ampere-style row permutation or column interleaving — SM90+ TMA-based FP4 kernels expect a simpler column-major packed layout:
    1. Input: [N, K/2] FP4 (2 per byte along K, row-major per expert)
    2. Nibble-level transpose [N, K][K, N]
    3. Output: [K, N/2] bytes (2 per byte along N, column-major packed)
  • Block scales: fc1_scales (input 3) — (E, N, K/32) uint8 storage, semantically float_ue8m0_t (8-bit power-of-2 exponent).
  • Global scale: fc1_global_scale (input 15) — (E,) float32.
  • Dequantization (applied during the GEMM in registers): W_float ≈ fp4_to_float(W_fp4) × ue8m0_to_float(block_scale) × global_scale.

6.6 Supported INT group sizes

ArchitectureActivationSupported block_size
SM75–89 (Turing/Ampere/Ada)FP16/BF1632, 64, 128
SM90 (Hopper)FP16/BF16falls back to Ampere — 32, 64, 128
SM100/120 (Blackwell)FP16/BF16falls back to Ampere — 32, 64, 128

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


7. Cross-Architecture Packing Compatibility

Weight packing is architecture-aware. The following table summarizes which packed weights are interchangeable across SMs:

Target SMCompatible packed weights from…Notes
SM70 (Volta)Not supported (no INT8 LDSM).
SM75 (Turing)SM75/80/86/89/100/120LDSM permutation + column interleaving.
SM80 (Ampere)SM75/80/86/89/100/120Same.
SM86/89 (Ada/Lovelace)SM75/80/86/89/100/120Same.
SM90 (Hopper)SM90 onlyHopper skips column interleaving (uses Permuted-Linear).
SM100/120 (Blackwell)SM75/80/86/89/100/120Falls back to SM80 packing for INT4/INT8.

Summary groups

  • Group A (universal INT4/INT8): SM75, SM80, SM86, SM89, SM100, SM120.
  • Group B (Hopper INT4/INT8): SM90 only.
  • MXFP4: separate format (§6.5) — does not use pack_weights_for_cuda_mixed_gemm.
  • FP8: no packing.

QMoE uses Group A on every GPU. The table above describes the layouts the pack_weights_for_cuda_mixed_gemm preprocessor can emit. The QMoE INT4/INT8 MoE GEMM, however, always dispatches to the Ampere (SM80) grouped-GEMM kernel — even on SM90 — because mixed int-weight + fp16/bf16 activation is not a valid Hopper TMA warp-specialized specialisation (the same is true in TensorRT-LLM). It therefore consumes the Group A (SM80) layout on all architectures, including Hopper. For QMoE, always pack INT4/INT8 weights for SM80 (arch=80), and PrePackIntExpertWeights (weights_prepacked=0) does exactly that regardless of the runtime device SM. Group B (SM90) layout is currently unused by QMoE.


8. SwiGLU Fusion

SwiGLU formula:

SwiGLU(x) = G × Sigmoid(alpha × G) × (L + beta)
    G = clamp(Gate,  max=limit)
    L = clamp(Value, min=-limit, max=limit)

Gate and Value are the two halves of the FC1 output. The behavior is controlled by activation_alpha (alpha), activation_beta (beta) and swiglu_limit (limit).

ParameterAttributeDefaultStandard SwiGLUGPT-OSS SwiGLU
alphaactivation_alpha1.01.01.702
betaactivation_beta0.00.01.0
limitswiglu_limitunset → +inf (no clamp)unset / +inf7.0

The attribute defaults are exactly Standard SwiGLU, which reduces to:

SwiGLU(x) = Gate × Sigmoid(Gate) × Value = SiLU(Gate) × Value

This is the activation used by Llama- and Gemma-style MoE. GPT-OSS SwiGLU instead uses alpha=1.702, beta=1.0, limit=7.0. Both variants run on the same CUDA kernel; when alpha=1.0, beta=0.0 and limit=+inf, the kernel takes the plain SiLU(Gate) × Value path (no clamping).

The operator supports three fusion modes via the swiglu_fusion attribute:

swiglu_fusionInputsFC1 layoutNotes
0fc1, fc2, fc3separate Gate / Value / UpConceptually three GEMMs.
1 (interleaved)fc1, fc2[Gate_0, Value_0, Gate_1, Value_1, …][E, 2×inter, hidden]GPT-OSS layout.
2 (block)fc1, fc2`[Gate_0…Gate_NValue_0…Value_N][E, 2×inter, hidden]`

Backward-compatibility remap (swiglu_fusion=0 + SwiGLU): The published gpt-oss-20b model before June 2025 uses interleaved SwiGLU layout but swiglu_fusion attribute to 0. To keep those models working, when activation_type="swiglu" and swiglu_fusion=0, the CUDA op treats the FC1 weights as interleaved (i.e. as if swiglu_fusion=1): unconditionally for QMoE (which never has a separate fc3). A one-time warning is logged. Consequently a SwiGLU model that genuinely intended the non-interleaved split must provide a separate fc3 (standard MoE) rather than rely on swiglu_fusion=0. New exporters should set swiglu_fusion explicitly.

CPU note: The CPU MoE/QMoE implementation only supports the interleaved SwiGLU layout (swiglu_fusion=1). The concatenated layout (swiglu_fusion=2) throws ORT_NOT_IMPLEMENTED on CPU; use the CUDA EP for concatenated SwiGLU.

Standard MoE runtime fc3 fusion

The non-quantized MoE operator (not QMoE) accepts an optional fc3_experts_weights input. When present, the op allocates a temporary buffer and concatenates fc1 (Gate) with fc3 (Value) per expert at runtime, simulating swiglu_fusion=2. This makes it easy to feed Mixtral-style models without offline fusion.

Note: This runtime fusion is only in standard MoE. For QMoE, weights must be fused offline before quantization+packing.


9. FP4 (MXFP4) Details

The QMoE operator supports MXFP4 quantized weights with FP16/BF16 activations (W4A16) via quant_type="fp4". The kernel path is the mixed-input TMA warp-specialized CUTLASS kernel.

9.1 Why "mixed input"?

Block-scaled tensor ops (OpClassBlockScaledTensorOp) require both operands to use block scaling (FP4×FP4 or FP8×FP4). W4A16 has full-precision activations paired with narrow FP4 weights — that is the mixed-input configuration. The dispatch flips on the use_wfp4a16 flag in moe_gemm_kernels.h:

cpp
static constexpr bool use_wfp4a16 = weight_fp4 &&
    (std::is_same_v<T, half> || std::is_same_v<T, __nv_bfloat16>);

9.2 Dispatch flow

CutlassMoeFCRunner<half, __nv_fp4_e2m1, half>::dispatchToArch()
  └─ use_wfp4a16 == true
     └─ select fusion from hopper_inputs.fusion (NONE or FINALIZE)
        └─ select K tile: inputs.k % 256 == 0 → PackedScalesNum=1 (K=256)
                           else              → PackedScalesNum=2 (K=128)
           └─ sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass<..., FUSION, PackedScalesNum>()
              └─ Ktile = PackedScalesNum==2 ? 128 : 256
                 └─ dispatch on tile_config_sm90 enum (M×N from heuristic)
                    └─ sm90_dispatch_moe_mixed_dtype_gemm_config<..., FUSION, Shape<M, N, Ktile>>()
                       └─ dispatch on cluster_shape
                          └─ sm90_dispatch_mainloop_schedules<..., FUSION>()
                             └─ sm90_generic_mixed_moe_gemm_kernelLauncher()
                                ├─ ElementA = cutlass::half_t  (activation)
                                ├─ ElementB = cutlass::float_e2m1_t  (weight, stored as FP4)
                                ├─ group_size = 32 (mxfp4_group_size)
                                ├─ ElementScale = cutlass::float_ue8m0_t
                                ├─ CollectiveBuilderMixedInput (FP4→FP16 upconvert in registers)
                                └─ Epilogue: NONE (per-expert output) or FINALIZE (fused scatter+scale)

Note: H100/H200 (SM90) does not have native FP4 tensor core instructions. The kernel uses FP4 purely as a compressed storage format — weights are loaded via TMA and upconverted to FP16/BF16 in shared memory/registers by CollectiveBuilderMixedInput before the actual MMA runs on FP16 tensor cores. This is a memory bandwidth optimization (4x compression), not a compute throughput feature. Native FP4 MMA is available on Blackwell (SM100+) via the separate block-scaled tensor op path (see §11).

Native FP4 path triggers when sm_ >= 120 (use_fp4_dequant_fallback_ = sm_ < 120). On older SMs, MXFP4 weights are decoded via LaunchQMoEDequantizeFp4Weights and fed to the dense A16 runner.

9.3 W4A16 vs W4A8-INT4 differences

PropertyW4A16 (MXFP4)W4A8-INT4
ElementAhalf_t / bfloat16_tfloat_e4m3_t
ElementBfloat_e2m1_tint4b_t
Group size32 (MXFP4)128 (INT4)
ElementScalefloat_ue8m0_t__nv_bfloat16 (SFA)
Epilogue α1 (no per-group α)0 (uses alpha_ptr_array)
Epilogue fusionNONE or FINALIZENONE or FINALIZE
M tiles64, 12864, 128
N tiles16, 32, 64, 12816, 32, 64, 128
K tiles128, 256128 × PackedScalesNum / sizeof(T)
Cluster shapes(1,1), (2,1), (1,2), (2,2)(1,1), (2,1), (1,2), (2,2)
Mainloop schedulesPingpong, CooperativePingpong, Cooperative

9.4 Mainloop modification

The CUTLASS collective mainloop uses a type-dependent group size:

cpp
static constexpr bool IsMXFP4 = cute::is_same_v<ElementA, cutlass::float_e2m1_t>;
static constexpr int  ScalingGroupSize = IsMXFP4 ? detail::mxfp4_group_size
                                                 : detail::int4_group_size;

This affects scale_k = K / ScalingGroupSize, NumMMAsPerChunk, and NumChunksPerTileK calculations.

9.5 Key data structures

cpp
// QuantParams::FP4Inputs (moe_kernels.h)
struct FP4Inputs {
  struct GemmInputs {
    bool use_per_expert_act_scale = false;
    float const* act_global_scale = nullptr;        // nullptr for W4A16
    NVFP4ElementSF const* weight_block_scale;       // (E, N, K/32) ue8m0 bytes
    float const* global_scale;                      // (E,) float
  };
  GemmInputs fc1, fc2;
};

// Block scaling type (moe_gemm_kernels.h)
enum class FpXBlockScalingType { MXFPX /*32*/, NVFP4 /*16*/, NONE };

9.6 Constructor and ComputeInternal

cpp
// Constructor (sm_ >= 120, ENABLE_FP4 + USE_FP4_QMOE)
m_moe_runner = std::make_unique<CutlassMoeFCRunner<half, __nv_fp4_e2m1, half>>(
    sm_, activation_type_, has_fc3_, normalize_routing_weights_, use_sparse_mixer_);

// ComputeInternal
quant_params = QuantParams::FP4(
    /*fc1_act_global_scale*/ nullptr,
    fc1_block_scales, fc1_global_scale,
    /*fc2_act_global_scale*/ nullptr,
    fc2_block_scales, fc2_global_scale);

9.7 Kernel instantiation files

FileTemplate
moe_gemm/moe_gemm_kernels_fp16_fp4.cuMoeGemmRunner<half, __nv_fp4_e2m1, half>
moe_gemm/moe_gemm_kernels_bf16_fp4.cuMoeGemmRunner<__nv_bfloat16, __nv_fp4_e2m1, __nv_bfloat16>
moe_gemm/launchers/moe_gemm_tma_ws_sm90_fp4_instantiation.cuhInstantiation macros: ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_{PP,CO} (NONE fusion), ORT_MOE_GEMM_TMA_WS_SM90_FP4_INST_{PP,CO}_FINALIZE
moe_gemm/launchers/generate_moe_gemm_tma_ws_sm90_fp4.pyPython generator: produces 320 .generated.cu files across FP16/BF16, M={64,128}, N={16,32,64,128}, K={128,256}, 4 cluster shapes, PP/CO schedules, NONE/FINALIZE fusion
moe_gemm/launchers/moe_gemm_tma_ws_sm90_fp4_*.generated.cu320 generated SM90 mixed-input FP4 launcher instantiations (built when onnxruntime_USE_FP4_QMOE=ON)
moe_gemm/launchers/moe_gemm_tma_ws_sm120_fp4_*.generated.cuSM120 mixed-input FP4 launcher
moe_gemm/launchers/moe_gemm_tma_ws_sm120_fp8_fp4.generated.cuSM120 block-scaled FP8×FP4 launcher (WFP4AFP8)

Build note: When onnxruntime_USE_FP4_QMOE is OFF, the stub is also excluded and all moe_gemm_kernels_*_fp4.cu / moe_gemm_tma_ws_sm{90,120}_fp4_*.generated.cu files are filtered out. CUDA 13 PTXAS does not complete the FP4 M=128/N=64 pingpong specializations, so those specific generated units are also excluded (the dispatcher routes that tile through cooperative variants instead). See §14.

9.8 K=128, Epilogue Fusion & Expanded Tile Configs

This subsection documents the expanded SM90 W4A16 mixed-input FP4 MoE GEMM configuration that closes the gap with TRT-LLM.

Changes Summary

GapBeforeAfter
K tiles256 only{128, 256} — selected at runtime based on inputs.k % 256
Epilogue fusionNONE onlyNONE + FINALIZE — routed from hopper_inputs.fusion
N tiles accessibleOnly CtaShape128x32x128B + ClusterShape_1x1x1All instantiated tiles (N={16,32,64,128}, clusters=(1,1),(2,1),(1,2),(2,2))
Generated .cu files~80320
Mainloop schedulesPingpong only (for most tiles)Pingpong + Cooperative (for M=128 tiles)

K Tile Dispatch Mechanism

The CutlassTileConfigSM90 enum encodes K as "128B" (128 bytes), but for FP4 mixed-input the actual K tile in elements differs. The dispatch uses a PackedScalesNum encoding trick:

  • PackedScalesNum = 1 → K = 256 elements (selected when inputs.k % 256 == 0)
  • PackedScalesNum = 2 → K = 128 elements (selected otherwise)

Inside sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass:

cpp
constexpr int Ktile = is_wfp4a16 ? (PackedScalesNum == 2 ? 128 : 256) : 128 * PackedScalesNum / sizeof(T);

Epilogue Fusion

The mixed-input launcher now supports two epilogue modes, matching the same-type launcher pattern:

  • NONE: Per-expert intermediate output (standard grouped GEMM epilogue)
  • FINALIZE: Fused scatter + router-scale + bias epilogue using EpilogueMoeFusedFinalizeBuilder

The fusion is routed at runtime in dispatchToArch:

cpp
switch (hopper_inputs.fusion) {
  case EpilogueFusion::FINALIZE:
    sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass<..., FINALIZE, PackedScalesNum>(...);
    break;
  case EpilogueFusion::NONE:
  default:
    sm90_dispatch_moe_mixed_dtype_gemm_to_cutlass<..., NONE, PackedScalesNum>(...);
    break;
}

Files Modified

FileChanges
launchers/moe_gemm_tma_ws_mixed_input_launcher.hAdded EpilogueFusion FUSION template parameter
launchers/moe_gemm_tma_ws_mixed_input_launcher.inlAdded FINALIZE epilogue support (CollectiveEpilogueFinalize, make_epilogue_scalars/args lambdas)
launchers/moe_gemm_tma_ws_sm90_fp4_instantiation.cuhAdded _PP_FINALIZE and _CO_FINALIZE macros
launchers/generate_moe_gemm_tma_ws_sm90_fp4.pyAdded k and fusion fields; generates K={128,256} × NONE/FINALIZE
moe_gemm_template_dispatch_tma_ws_mixed_dtype.hFUSION param throughout; PackedScalesNum-based K tile; direct N tile mapping; workspace calc with Ntile=128
moe_gemm_template_dispatch.hFUSION 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.

VariableDefaultEffect
ORT_ENABLE_FP4_GEMVonFused 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_AUTOTUNE0Opt-in per-shape autotune of the GEMV CtaN/Threads tiling. Enabling it synchronizes the first uncached inference for each shape.
ORT_FP4_GEMV_AUTOTUNE_LOG0Set to 1 to log the chosen GEMV configs per shape.
ORT_FP4_GEMV_INTERLEAVED0Experimental, 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.
ORT_FP4_GEMV_INTERLEAVED_HALFACC0Override. 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_GEMM1Routes SM80–SM119 FP4 prefill through the fused-dequant grouped GEMM. Set to 0 to force dense fallback for debugging or comparison. Decode routes through the fused MXFP4 GEMV, reading the same pre-packed buffer as prefill. In this regime the raw e2m1 initializers are released after PrePack. See §9.11.
ORT_ENABLE_FP4_CUTLASS_GEMM0Opt-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_UNSAFE0Confirms 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_TOKENS64When 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 <typename T>
using Fp4GemvAccT = std::conditional_t<std::is_same<T, half>::value, half, float>;
ActivationAccumulatorRationale
fp16fp16 (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.
bf16fp32 (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, 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.

9.11 Single-copy SM80 weights (prefill + decode share one buffer)

Default on SM80–SM119 (ORT_FP4_SM80_GEMM=1). The SM80 grouped GEMM and the fused decode GEMV historically needed different e2m1 weight layouts, so a QMoE node kept up to three persistent copies of the expert weights:

copylayoutconsumer
raw initializer[E, K, N/2] schema layoutdequant fallback — unreachable in this regime
gemv_fp4_fc*_weights_SM80 ColumnMajorTileInterleave + nibble pair-interleavegrouped-GEMM prefill
gemv_fp4_fc*_weights_decode_ColToRow (or interleave steps 1–3)fused GEMV decode

For a 20B-class MXFP4 MoE each copy is ~9 GiB, which put post-load VRAM out of reach of 24 GB consumer cards. Two changes collapse this to a single copy:

  1. Release the raw initializers. In this regime every dispatch is served by a pre-packed buffer, so the dequant fallback — the only consumer of the raw layout — is unreachable. PrePack now caches fc*_weights_shape_ for CheckInputs and reports is_packed = true.
  2. Teach the GEMV to read the SM80 layout. The two layouts differ by exactly one preprocessor step: the [e0,e2,e4,e6,e1,e3,e5,e7] nibble pair-interleave applied by interleave_int4s_inplace_kernel. Inverting it in the decoder is a compile-time index remap of the same eight decode calls — no branches, no extra registers:
cpp
// onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h
template <typename AType, bool PairInterleaved = false>
struct Fp4I2FConverter { ... };            // PairInterleaved un-permutes the SM80 nibble order

// onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu
template <typename T>
using Fp4KernelDetailsSm80Pair =
    fiv::KernelDetails<typename Fp4ADetails<T>::Type, fiv::Fp4DetailsW,
                       fiv::ColumnMajorInterleaved, /*UseInterleavedConverter=*/true,
                       kTileSizeKFp4>;

ConverterWrapper now forwards kUseInterleavedConverter to the FP4 branch, so selecting Fp4KernelDetailsSm80Pair is all that is needed to decode the prefill buffer directly.

Shape gate. Because this reuses the interleaved GEMV (kInterleave=4, kStepK=32), it inherits the interleaved rules — MXFP4 group_size == 32, n % 16 == 0, k % 64 == 0 — checked at pack time by is_moe_gemv_fp4_sm80_layout_supported. Shapes that miss the gate transparently fall back to a dedicated gemv_fp4_fc*_weights_decode_ copy (logged at INFO); NVFP4 (block 16) always uses the plain ColToRow layout. gpt-oss-20b (k=2880, n=5760/2880) clears the gate.

Measuredgpt-oss-20b MXFP4, post-load device memory:

configurationpost-load
before (3 copies)32908 MiB
release initializers only (2 copies)23164 MiB
single copy (default)13996 MiB

Returning the memory to the device. Released initializers only shrink the process's device footprint when initializers bypass the BFC arena. Set the session option session.use_device_allocator_for_initializers = 1; otherwise the freed bytes are merely recycled inside the arena for later activation/KV allocations.


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 codesE2M1, 2 codes/byte (uint8)E2M1, 2 codes/byte (uint8) — same
Block size3216
Block scalefloat8e8m0 (E8M0, pow-of-two), stored uint8float8e4m3fn (E4M3)
Global scale1.0 (implicit)per-expert FP32 fc*_global_scale (weight_scale_2)
Reconstructw = 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).

Execution behavior: On SM120/SM121 (Blackwell) NVFP4 runs a native block-scaled CUTLASS FP4×FP4 grouped-GEMM prefill path (activations dynamically quantized to NVFP4 in-runner), gated by ORT_ENABLE_NVFP4_CUTLASS_GEMM (default on) and the CUTLASS shape requirements. On SM<120, or when the native path is disabled / shape-unsupported, NVFP4 falls back to the dequant-to-A16 path. The fused GEMV decode fast path additionally handles supported small-decode shapes on all SMs. Details in §9b.1.

9b.1 Dispatch

On SM120/SM121 NVFP4 selects a native block-scaled CUTLASS FP4×FP4 grouped-GEMM prefill path (CutlassMoeFCRunner<__nv_fp4_e2m1, __nv_fp4_e2m1, ...>, runner instantiated in moe_gemm_kernels_fp4_fp4.cu). Activations are dynamically quantized to NVFP4 (block-16 E4M3 scale, global scale 1.0) inside the runner (expandInputRowsKernel), and the prepacked weight block scales are consumed as the Blackwell 128×4 swizzled SF atom. The path is enabled when sm_ is 120 or 121, ORT_ENABLE_NVFP4_CUTLASS_GEMM is on (default), and the CUTLASS shape requirements hold (hidden/inter multiple of 64); it sets use_fp4_dequant_fallback_ = false. ORT_FP4_PREFILL_MIN_TOKENS (default 64) splits the workload: prefill-shaped batches take the native grouped GEMM while small-decode shapes still take the fused GEMV path; an A16 dense fallback runner covers batches routed away from native.

When the native path is unavailable (SM<120, disabled, or shape-unsupported), NVFP4 uses the dequant-to-A16 fallback (§4.3) for prefill / GEMV-unsupported shapes, and the fused GEMV decode fast path (§4) 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) 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).
  • NVFP4 uses only the non-interleaved ColumnMajor layout; the opt-in interleaved path (§9.10) is MXFP4-only because its kStepK=32 tile is tied to the block-32 scale layout.
  • QMoECombineNvfp4ScalesForGemv (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) 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):

kernelbeforeafterΔ
FC1 GEMV (SwiGLU-fused, n=1024 k=2048)20.6 µs13.25 µs−36%
FC2 GEMV (n=2048 k=512)10.7 µs7.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 (-k gemv compares ORT_ENABLE_FP4_GEMV=1 vs =0 vs a torch reference).
  • Profiling: profile_qmoe_gemv.py builds the QMoE node at configurable decode shapes and compares GEMV with the grouped-GEMM path (with ORT_FP4_GEMV_AUTOTUNE_LOG=1 for the chosen tiling).

9b.5 Native prefill accuracy & routing

The native FP4×FP4 prefill path quantizes both operands to 4-bit NVFP4: weights are pre-quantized offline, and activations are dynamically quantized to NVFP4 (block-16 E4M3 scale, global scale 1.0) inside expandInputRowsKernel. This is a true W4A4 GEMM, so it carries more quantization error than the weight-only (W4A16) dequant-fallback / GEMV paths, which keep activations in FP16/BF16. Against a full-precision-activation reference the native prefill output therefore differs by a small but non-negligible amount (order ~0.2 absolute on unit-scale outputs), coming entirely from the activation quantization — not a kernel bug. Forcing the same shape through the dequant fallback (ORT_FP4_PREFILL_MIN_TOKENS very large) reduces the difference to ~1e-3, which confirms the native GEMM arithmetic is correct.

Routing between the native and weight-only paths is governed by ORT_FP4_PREFILL_MIN_TOKENS (default 64):

  • num_rows >= ORT_FP4_PREFILL_MIN_TOKENS → native block-scaled FP4×FP4 prefill (when enabled and shape-supported on SM120/SM121).
  • num_rows < ORT_FP4_PREFILL_MIN_TOKENS → small-decode shapes stay on the fused GEMV decode / dequant fallback (W4A16), keeping the latency-sensitive decode path off the extra activation-quant error.

ORT_ENABLE_NVFP4_CUTLASS_GEMM=0 disables the native path entirely (all shapes route to GEMV/fallback). The parity test test_qmoe_nvfp4_cuda.py mirrors this gate: it applies a looser tolerance only to shapes that route native (SM120/SM121, native enabled, num_rows >= ORT_FP4_PREFILL_MIN_TOKENS) and keeps the strict bound for all decode/fallback shapes, so a genuinely broken native kernel (error order ~1.0+) is still caught.


10. FP8 (W8A16) Details

quant_type="fp8" supplies FP8 e4m3 weights with BF16/FP16 activations. This was added so H200 (SM90) has a working narrow-weight QMoE path that does not require the FP4 launcher.

10.1 Native dispatch (SM90+)

cpp
// Constructor — sm_ >= 90 with ENABLE_FP8
m_moe_runner = std::make_unique<CutlassMoeFCRunner<half, __nv_fp8_e4m3, half>>(...);
// or BF16 variant
m_moe_runner = std::make_unique<CutlassMoeFCRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>>(...);

The SM80 specialization in moe_gemm_template_dispatch.h is intentionally left uninstantiated for W8A16-FP8; the implementation comment states that the SM80 path is not supported and the native path is SM90 TMA WS. On Hopper, use_wfp8a16 is routed through the TMA warp-specialized dispatcher, which enforces inputs.gemm_config.is_tma_warp_specialized and applies the per-expert global scale via alpha_scale_ptr_array in the epilogue. On SM120, the code redirects W8A16-FP8 to the SM89 FP8 kernel implementations.

10.2 Scale wiring

QuantParams::FP8::dequant_fc1  (float*, num_experts)
       │
       ▼
computeFP8DequantScale()  →  alpha_scale_ptr_array[e] = &dequant_fc1[e]
       │
       ▼
GroupedGemm with EpilogueOpDefault:
   output[i] = fp8_to_bf16(gemm_accum[i]) * (*alpha_scale_ptr_array[expert])

computeFP8DequantScale and the Ampere FP8 epilogue already exist (moe_kernels.cu) — the QMoE op only needs to construct QuantParams::FP8(dequant_fc1, nullptr, dequant_fc2) from the per-expert global scales (inputs 15/16).

10.3 Dequant fallback (SM<90)

LaunchQMoEDequantizeFp8Weights decodes weights into BF16/FP16 and the dense A16 runner is used.

10.4 Kernel instantiation files

FileTemplate
moe_gemm/moe_gemm_kernels_fp16_fp8.cuMoeGemmRunner<half, __nv_fp8_e4m3, half>
moe_gemm/moe_gemm_kernels_bf16_fp8.cuMoeGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>

10.5 End-to-end data flow

Model input (BF16/FP16)
    │
    ▼
Router → top-k → permute
    │
    ▼  (still BF16/FP16 — no activation quantization)
GEMM1: bf16_act × fp8_weight → bf16_out  (×dequant_fc1 in epilogue)
    │
    ▼
Activation (SwiGLU / SiLU / ReLU / …)
    │
    ▼
GEMM2: bf16_act × fp8_weight → bf16_out  (×dequant_fc2 in epilogue)
    │
    ▼
Un-permute → weighted sum → output

11. WFP4AFP8 Details

quant_type="wfp4afp8" pairs MXFP4 weights with FP8 e4m3 activations. Unlike W4A16, both operands use block scaling, so this path uses CUTLASS's block-scaled tensor op primitive (OpClassBlockScaledTensorOp) — natively supported only on SM100+ (Blackwell).

11.1 Native dispatch (SM100+)

cpp
// Constructor — sm_ >= 100 with ENABLE_FP4 + USE_FP4_QMOE + ENABLE_FP8
m_moe_runner = std::make_unique<
    CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half, half>>(...);
// or BF16 output:
m_moe_runner = std::make_unique<
    CutlassMoeFCRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, __nv_bfloat16, __nv_bfloat16>>(...);

The runner is constructed with T = __nv_fp8_e4m3, WeightType = __nv_fp4_e2m1, OutputType = half/bf16, InputType = half/bf16. The user-facing input is BF16/FP16; the runner quantizes it to MXFP8 (FP8 + per-block ue8m0 scales) inside expandInputRowsKernel (MXFP8 branch). The MXFP8 branch is triggered when quant_params.mxfp8_mxfp4.fc{1,2}.weight_block_scale is non-null.

11.2 Two variants

VariantQuantParams factoryActivation scalingUsed inputs
A — global-scaled FP8 actQuantParams::FP8MXFP4per-tensor or per-expert float scaleweight side (3,15) + (6,16); act 18, 19
B — MXFP8 block-scaled actQuantParams::MXFP8MXFP4per-block ue8m0 scalesweight side (3,15) + (6,16); act 20, 21

The current build uses Variant B (QuantParams::MXFP8MXFP4) for the native SM100+ path; activation block scales are produced inside the runner by expandInputRowsKernel. The act_scale inputs (18/19) are validated and pre-packed for forward compatibility with Variant A but are not consumed by the current native plumbing.

11.3 Dequant fallback (SM<100)

cpp
use_wfp4afp8_dequant_fallback_ = (sm_ < 100);

When the fallback is selected, MXFP4 weights are decoded with LaunchQMoEDequantizeFp4Weights and fed into the dense BF16/FP16 MoE runner — exactly the same path used by quant_type="fp4" on SM<120. Verified working on SM90 (H200) using the bundled Python parity test.

11.4 Kernel instantiation files

FileTemplate
moe_gemm/moe_gemm_kernels_fp8_fp4.cuMoeGemmRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, half> and BF16 variant
moe_gemm/launchers/moe_gemm_tma_ws_sm120_fp8_fp4.generated.cuSM120 block-scaled tensor op launcher (FP8×FP4, 128×128×128 tile)

Built only when onnxruntime_USE_FP4_QMOE=ON (which implies ENABLE_FP4+ENABLE_FP8). The SM120 launcher additionally requires COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS (set by cmake when SM120 is in CMAKE_CUDA_ARCHITECTURES).

11.5 Why two CUTLASS paths

W4A16 :  FP16 act (full precision) × FP4 weight (block scaled)
        → mixed-input  (CollectiveBuilderMixedInput, group=32, ue8m0 scales)

W4A8  :  FP8 act (block scaled)    × FP4 weight (block scaled)
        → block-scaled tensor op  (OpClassBlockScaledTensorOp — native FP4×FP8 in tensor cores)

The block-scaled tensor op path is fundamentally more efficient because the hardware fuses dequantization with the matrix multiply, vs. the in-register software dequant of the mixed-input path.

MSVC note: Native SM90/SM120 TMA grouped MoE kernels are disabled in Windows/MSVC builds because CUDA 13 generates host stubs that MSVC rejects for over-aligned TMA parameters. See §14.1.


12. Future / Deferred Modes

ModeNotationActivationWeightStatus
W4AFP8INT4 weight + FP8 activationFP8 e4m3INT4 (uint4b_t)Deferred — fast path is gated to SM89 only in TRT-LLM (moe_gemm_template_dispatch.h); falls back to Ampere dequant on SM90+ which offers no advantage over W4A16-int. A proper SM90 W4AFP8 TMA WS kernel would be needed.
W8A8-fp8FP8 act + FP8 weightFP8 e4m3FP8 e4m3Not implemented. Targeted for SM89 (RTX 4090).
WFP4AFP8 native validation(above)Native path implemented; end-to-end validation requires SM100+ hardware.
WFP4AFP8 Variant Aglobal-scaled FP8 activationRequires QMoE op to accept pre-quantized FP8 input or wire a separate global-scaled BF16→FP8 prologue.

The schema reserves the necessary input slots (18–21) so adding these modes will not change the operator interface.

12.1 MoE GEMV Optimization Summary

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

Completed per-column INT4 work

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

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

Completed per-column INT8 work

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

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

Completed block-wise INT4/INT8 work

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

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

Completed BF16 enablement

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

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

Split-K2 SwiGLU GEMV default path

The fp16 INT4 interleaved-SwiGLU GEMV path uses a two-pass Split-K2 FC1 kernel by default for supported decode shapes. The first pass computes two K-split FP32 partials into QMoE workspace, and the second pass reduces those partials, adds optional bias, and applies the interleaved SwiGLU epilogue. FC2 stays on the regular moe_gemv_kernel path.

Set ORT_DISABLE_MOE_GEMV_SPLITK2_SWIGLU=1 before process start to force the previous single-kernel FC1 SwiGLU GEMV path for debugging, A/B benchmarking, or bisecting numerical differences. On GPT-OSS-20B, Split-K2 reduced FC1 kernel work from about 21.42 us to 19.98 us and improved repeated CUDA-graph decode throughput by about 0.9% to 1.6% with valid focused-helper output. A 1000-sample MMLU smoke matched the opt-out fallback within noise. A future autotuner can replace this hand-selected default with per-shape route selection.

bash
onnxruntime/test/python/transformers/profile_qmoe_gemv.py \
  --case gpt_oss_20b_m1_top4_fp16_2880x2880_e32 \
  --disable-splitk2-swiglu --warmup 5 --repeat 100 --nvtx

Accumulation policy

The QMoE GEMV fast path accumulates fp16 activations in fp16 by default. Set ORT_MOE_GEMV_FP32_ACCUM=1 before process start to restore the previous fp32 accumulation path for fp16 activations. BF16 activations always use fp32 accumulation because bf16 accumulation is too lossy.

On the GPT-OSS-20B decode-shaped helper case gpt_oss_20b_m1_top4_fp16_2880x2880_e32, the default fp16-accumulation path was 0.0708 ms versus 0.0812 ms with ORT_MOE_GEMV_FP32_ACCUM=1. In a full GPT-OSS CUDA-graph decode run, default fp16 accumulation reached 386.26 tok/s versus 353.70 tok/s with the fp32 fallback. A 1000-sample MMLU smoke test matched pooled accuracy at 0.8260 for both modes.

Experiments rejected after profiling

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

Remaining ideas not tried

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

Keep ORT_DISABLE_MOE_GEMV=1 available. It is useful for A/B testing, fallback validation, and bisecting numerical or performance regressions. For quick profiling, use profile_qmoe_gemv.sh:

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

13. Testing

Test fileCoverage
test_moe_cuda.pyStandard MoE on CUDA: FP16/BF16, SiLU/GeLU/SwiGLU, routing, GEMM parity. SwiGLU coverage includes both GPT-OSS (TestSwigluMoE: interleaved, alpha=1.702/beta=1.0/limit=7.0) and Standard/Llama-Gemma (TestStandardSwigluMoE: concatenated swiglu_fusion=2, alpha=1.0/beta=0.0/no limit → SiLU(Gate)×Value).
test_moe_cpu.pyStandard MoE on CPU (smoke).
test_qmoe_cuda.pyINT4/INT8 QMoE — primary regression signal for the production QMoE path. Exercises pack_weights_for_cuda_mixed_gemm and dequant-then-matmul reference. TestQMoEIntPrePackSmoke covers the raw-weight weights_prepacked=0 in-PrePack layout transform (smoke test: asserts finite output, not bit-parity).
test_qmoe_cpu.pyINT4/INT8 QMoE on CPU (smoke).
test_qmoe_fp4_cuda.pyMXFP4 QMoE: quantization utilities, packing, FP16/BF16, SiLU/SwiGLU, top-k and expert-count variants. End-to-end runs on SM120; on SM<120 the dequant fallback is exercised.
test_qmoe_fp8_cuda.pyFP8 W8A16 QMoE on SM90+ native path and SM<90 dequant fallback.
test_qmoe_wfp4afp8_cuda.pyWFP4AFP8 — native Blackwell path requires SM100+; SM<100 exercises the dequant fallback.

Reference computation

The "ground truth" is computed by dequantizing weights to FP16 in Python:

python
dequantized = (q_weight - zero_point) * scale     # INT
# or
dequantized = fp4_to_float(W) * ue8m0_to_float(block_scale) * global_scale  # FP4
reference = input @ dequantized.T

This validates the numerical correctness of the dequantization fusion.


14. Build Configuration

CMake gates relevant to MoE/QMoE (see cmake/CMakeLists.txt and cmake/onnxruntime_providers_cpu.cmake):

DefineSet whenEffect
ENABLE_BF16CUDA ≥ 11.0BF16 weight/activation paths.
ENABLE_FP8CUDA ≥ 11.8FP8 e4m3 instantiations and QuantParams::FP8.
ENABLE_FP4CUDA ≥ 12.8FP4 e2m1 type (__nv_fp4_e2m1) and FP4 traits.
onnxruntime_USE_FP4_QMOEuser opt-in (requires ENABLE_FP4)Enables FP4 / WFP4AFP8 kernel instantiations and CUTLASS launchers.
EXCLUDE_SM_100, EXCLUDE_SM_120architecture exclusionDrops the corresponding generated kernels.

CUDA architecture defaults:

  • CUDA 12.8+ : 60;70;75;80;86;89;90;100;120
  • CUDA 13.x : 75;80;86;89;90;100;120
  • SM90+ gets -a suffix (enables WGMMA, TMA, setmaxnreg).

CMake exclusion filters (current state)

cmake/onnxruntime_cuda_source_filters.cmake:

cmake
if(NOT onnxruntime_USE_FP4_QMOE)
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm90_fp4_.*\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm120_fp4_.*\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm120_fp8_fp4\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_kernels_(fp16|bf16)_fp4\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_kernels_fp4_fp4\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_kernels_fp8_fp4\\.cu")
else()
  # CUDA 13 PTXAS does not complete the FP4 M=128/N=64 pingpong specializations
  # in this build configuration. The dispatcher routes that tile through
  # cooperative mainloop variants instead, so exclude only those unused units.
  list(FILTER … EXCLUDE REGEX
       "moe_gemm_tma_ws_sm90_fp4_(fp16|bf16)_m128_n64_k[0-9]+_cm[12]_cn[12]_pp(_finalize)?\\.generated\\.cu")
endif()

if(NOT onnxruntime_USE_FP8_QMOE)
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm90_wfp8_.*\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm120_fp4_fp8_.*\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_tma_ws_sm120_fp8_fp4\\.generated\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_kernels_(fp16|bf16)_fp8\\.cu")
  list(FILTER … EXCLUDE REGEX "moe_gemm_kernels_fp8_fp4\\.cu")
endif()

14.1 MSVC and TMA grouped MoE GEMM

Windows/MSVC builds intentionally do not define the grouped TMA MoE compile switches:

  • COMPILE_HOPPER_TMA_GROUPED_GEMMS
  • COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS

The generated grouped TMA launchers pass CUTLASS TMA descriptor types through NVCC-generated host stubs. With CUDA 13 and MSVC, those stubs contain formal parameters with 128-byte alignment requirements, which triggers MSVC C2719: the requested alignment for a by-value formal parameter cannot be guaranteed. This affects the generated SM90/SM120 grouped MoE TMA launcher translation units, including the native SM120 QMoE FP4 / FP8×FP4 launchers.

The source files are still present in the build graph, but the generated launcher bodies are guarded by the compile switches above, so they become empty units on MSVC. Runtime dispatch mirrors this build-time choice:

  • Standard FP16/BF16 MoE may skip TMA configs and use the existing SM80/Ampere grouped GEMM fallback.
  • QMoE modes that require grouped TMA kernels do not silently fall back to SM80. This includes FP4/block-scaled modes such as native SM120 fp4, wfp4afp8, and other TMA-only mixed quantized paths. They fail with a clear error saying the required TMA grouped MoE GEMM was not compiled.
  • wfp4a16 on SM120 normally routes through the SM90 mixed-input TMA kernel set for forward compatibility, but it is also unavailable when the Hopper grouped TMA switch is disabled by MSVC.

The intent is to keep Windows CUDA packaging builds working while avoiding a misleading or invalid fallback for QMoE configurations whose data layout requires TMA/block-scaled kernels. Re-enable these switches for MSVC only after the CUDA host-stub alignment issue is fixed or the launcher ABI is changed to avoid over-aligned by-value parameters.

14.2 LLM object library and SM120 native SASS

The CUDA provider is split into per-architecture OBJECT libraries (cmake/onnxruntime_providers_cuda.cmake). The general LLM object library (onnxruntime_providers_cuda_llm, SM75+) holds the ordinary LLM/MoE kernels (MoE expand/finalize, fpA_intB GEMV/GEMM, NVFP4 dequant, ...). The SM120-specific TMA warp-specialized grouped-GEMM kernels are compiled separately in onnxruntime_providers_cuda_sm120_tma and their bodies are gated behind COMPILE_BLACKWELL_SM120_TMA_GROUPED_GEMMS, so the LLM library itself contains no SM120-only device code — building it at 120-real simply produces native SASS for the ordinary kernels instead of leaving them as compute_120 PTX that must be JIT-compiled on first use.

Architecture filtering of the LLM library (onnxruntime_filter_cuda_archs):

Platform / configLLM library archsRationale
Linux (any)includes 120-real (native SASS)LLM kernels run natively on SM120; no JIT warm-up; robust to real-only arch lists.
Windows/MSVC, USE_FP4_QMOE=OFFexcludes 120-real, keeps virtual compute_120 PTX (EXCLUDE_SM120_REAL)native sm_120a pulls CCCL tcgen05 PTX headers that fail to compile with the MSVC host toolchain; SM120 devices JIT from the kept PTX.
any platform, USE_FP4_QMOE=ONincludes 120-realthe NVFP4 native path emits cvt.e2m1x2 in expandInputRowsKernel, valid only for real sm_120a and not expressible in virtual compute_120 PTX.

CUDA 209 gotcha (real-only arch list). If the LLM library excludes 120-real and the configured arch list carries no virtual compute_120 PTX entry (e.g. CMAKE_CUDA_ARCHITECTURES="86-real;120-real"), the library ends up with no loadable SM120 image at all, and an SM120 GPU fails at EP load / first kernel launch with no kernel image is available for execution on the device (CUDA error 209). On Linux this cannot happen because the LLM library keeps 120-real. On Windows/MSVC, ensure a virtual compute_120 PTX arch is present (e.g. 86-real;120-real;120-virtual) so the JIT fallback works. Verify native SASS is present with cuobjdump libonnxruntime_providers_cuda.so | grep 'arch =' (use the toolkit-matched cuobjdump; an older one in PATH may misparse the multi-arch fatbin and abort).


15. Limitations & Known Issues

  • Row-wise INT quantization (block_size <= 0): does not currently support zero points in the QMoE operator.
  • Asymmetric INT zero points are supported only when block_size >= 64.
  • Minimum dimension: hidden_size and inter_size must be ≥ 16 (and aligned to 128 bits — multiples of 8 for FP16). See §4.2.
  • Float32 input: always uses the SM80 (Ampere) kernel path regardless of the actual device SM.
  • FP4 native path (SM90/SM100): although CUTLASS supports SM90 mixed-input FP4, the QMoE op currently routes only sm_ >= 120 through the native FP4 runner. SM90/SM100 fall back to dequantization. (Remove sm_ < 120 and rebuild to enable native FP4 on those SMs once validated.)
  • Windows/MSVC native TMA QMoE: grouped TMA MoE kernels are disabled on MSVC because CUDA 13 host stubs hit MSVC C2719 with over-aligned TMA parameters. Standard MoE can fall back to SM80 kernels; native QMoE FP4/block-scaled modes cannot. See §14.1.
  • Windows/MSVC SM120 LLM library: the general LLM object library excludes native sm_120a on MSVC (CCCL tcgen05 PTX headers do not compile with the MSVC host toolchain) and relies on virtual compute_120 PTX + JIT. As a result a real-only SM120 arch list can leave the library with no loadable image (CUDA error 209) — keep a compute_120 (virtual) arch on MSVC. NVFP4 QMoE (USE_FP4_QMOE) requires native sm_120a for cvt.e2m1x2 and is therefore a non-MSVC configuration in practice. See §14.2.
  • WFP4AFP8 native requires SM100+ hardware; only the dequant fallback path is validated end-to-end so far.
  • In-PrePack INT weight layout transform (weights_prepacked=0) is currently covered only by a smoke test (TestQMoEIntPrePackSmoke), not a bit-parity check: the existing offline pre-pack harness hardcodes force_arch=80 (the same SM80 layout consumed by the CUDA EP on all GPUs), so a separate parity harness for this path is still pending.
  • Hopper W4A8 (INT4 weight + FP8 activation) is not supported — TRT-LLM gates its fast path to SM89 only.
  • MXFP4 single-copy weights require group_size == 32, n % 16 == 0 and k % 64 == 0. Shapes outside that gate keep a second decode-layout copy (~2× weight memory); NVFP4 (block 16) always does. See §9.11.

16. Differences vs. TensorRT-LLM

The CUTLASS kernels are derived from TensorRT-LLM (CUTLASS 4.4.2, commit 346018db87) but have been significantly modified.

Modifications

  1. Pre-packed ZP/Bias optimizationPrePack derives (K − ZP) × scale biases offline so the kernel handles asymmetric quantization with no extra subtraction. (See §5.2.)
  2. SwiGLU interleaving — activation kernels support interleaved Gate/Value weights (§8).
  3. Sparse Mixer support via the use_sparse_mixer attribute.
  4. supportsTmaWarpSpecialized() exposed on CutlassMoeFCRunnerInterface to allow dynamic min_dim selection without knowing the concrete template type at the call site.
  5. MXFP4 in QMoE schema — extended schema and runner to accept MXFP4 weights plus per-expert global scales and ue8m0 block scales (§9).
  6. WFP4AFP8 (SM100+) — added MoeGemmRunner<__nv_fp8_e4m3, __nv_fp4_e2m1, …> with in-runner BF16/FP16→MXFP8 quantization in expandInputRowsKernel (§11).
  7. Backported bug fix (TRT-LLM 603ec03f) — moved griddepcontrol.launch_dependents to after computeTmaWarpSpecializedInputPointers in computeStridesTmaWarpSpecializedKernel to fix a pre-exit race.

Removed (not needed for ORT MoE/QMoE)

  • LoRA parameters (use_lora, LoraParams)
  • Min-latency mode (MoeMinLatencyParams)
  • AllToAll MoE paths (enable_alltoall)
  • DeepSeek FP8 block-scale GEMM mode (use_deepseek_fp8_block_scale, BlockScaleParams)
  • Deep Gemm, standalone FP4 GEMM, FP8 block-scale GEMM, fused gated GEMM directories (the relevant pieces are inlined into the MoE runner).