diff --git a/docs/contrib_ops/cuda/moe_qmoe.md b/docs/contrib_ops/cuda/moe_qmoe.md index 55d654c0f6376..1496e395f1a03 100644 --- a/docs/contrib_ops/cuda/moe_qmoe.md +++ b/docs/contrib_ops/cuda/moe_qmoe.md @@ -23,6 +23,7 @@ and have been significantly modified for ONNX Runtime — see 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) + - [9.11 Single-copy SM80 weights (prefill + decode share one buffer)](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer) 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) @@ -253,6 +254,12 @@ A16 runner. Helper kernels: 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](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer)). 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). @@ -315,6 +322,13 @@ pointers are read at compute time instead. 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](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer). + ### 5.2 INT4/INT8 scales + zero-point → bias @@ -777,7 +791,7 @@ debug switches. | `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_FP4_SM80_GEMM` | `1` | Routes 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](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer). | | `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. | @@ -818,6 +832,65 @@ The interleaved layout speeds up fp16 decode; bf16 gets no speedup (the fp32 reg 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: + +| copy | layout | consumer | +|------|--------|----------| +| raw initializer | `[E, K, N/2]` schema layout | dequant fallback — **unreachable** in this regime | +| `gemv_fp4_fc*_weights_` | SM80 `ColumnMajorTileInterleave` + nibble pair-interleave | grouped-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 +struct Fp4I2FConverter { ... }; // PairInterleaved un-permutes the SM80 nibble order + +// onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +template +using Fp4KernelDetailsSm80Pair = + fiv::KernelDetails::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. + +**Measured** — `gpt-oss-20b` MXFP4, post-load device memory: + +| configuration | post-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 @@ -1431,6 +1504,10 @@ Architecture filtering of the LLM library (`onnxruntime_filter_cuda_archs`): 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](#911-single-copy-sm80-weights-prefill--decode-share-one-buffer). --- 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 2504f12b57776..45c0089218144 100644 --- a/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h +++ b/onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h @@ -232,7 +232,20 @@ struct I2FConverter { } }; -template +// E2M1 (FP4) -> half/bf16 converter. +// +// ``PairInterleaved`` selects the nibble order of the packed source: +// false -- linear order, i.e. element ``i`` occupies nibble ``i`` of the 32-bit word. This is +// what ``LaunchQMoERepackFP4ColToRow`` emits and what the CUTLASS fpA_intB W4_A16 +// preprocessor's layout-only steps 1-3 preserve. +// true -- the ``[e0,e2,e4,e6,e1,e3,e5,e7]`` pair-interleave applied by +// ``interleave_int4s_inplace_kernel`` (preprocessor step 4 without the integer +8 +// bias). This is the layout the SM80 grouped GEMM consumes, so reading it here lets a +// single pre-packed weight buffer serve both the grouped-GEMM prefill and the fused +// GEMV decode instead of keeping two full copies of the expert weights. +// The un-permutation is a compile-time index remap of the same eight ``decode`` calls, so it +// costs no extra registers, branches or ALU work. +template struct Fp4I2FConverter { static_assert(std::is_same_v || std::is_same_v); @@ -271,14 +284,33 @@ struct Fp4I2FConverter { 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); + if constexpr (!PairInterleaved) { + static_assert(N % 2 == 0); #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)); + 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)); + } + } else { + // The pair-interleave permutes whole 32-bit words, so N must cover complete words. + static_assert(N % 8 == 0, "Pair-interleaved FP4 decode needs a multiple of 8 elements"); + // Packing writes element i to nibble slot (i even ? i/2 : (i - 1)/2 + 4), so logical + // element i is read back from slot kSlot[i]. Nibble slot j lives in byte j/2, low nibble + // for even j. kSlot is constexpr and the loops are unrolled, so every index below folds + // to a compile-time constant. + constexpr int kSlot[8] = {0, 4, 1, 5, 2, 6, 3, 7}; +#pragma unroll + for (int w = 0; w < N / 8; ++w) { + uint8_t const* word = s + w * 4; +#pragma unroll + for (int i = 0; i < 8; ++i) { + uint8_t const byte = word[kSlot[i] >> 1]; + d[w * 8 + i] = decode(static_cast((byte >> ((kSlot[i] & 1) * 4)) & 0x0F)); + } + } } } }; @@ -290,7 +322,7 @@ struct ConverterWrapper { static constexpr bool kUseInterleavedConverter = Details::kUseInterleavedConverter; using Converter = std::conditional_t< IsFp4Weight::value, - Fp4I2FConverter, + Fp4I2FConverter, I2FConverter>; }; 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 index ca5336b9fe0a6..1bbe9a8bbb461 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu @@ -90,6 +90,24 @@ using Fp4KernelDetailsInterleaved = fiv::KernelDetails::Type, fiv::Fp4DetailsW, fiv::ColumnMajorInterleaved, false, kTileSizeKFp4>; +// Same interleaved Details, but with UseInterleavedConverter = true so the e2m1 decoder inverts +// the `[e0,e2,e4,e6,e1,e3,e5,e7]` nibble pair-interleave on the fly. That is preprocessor step 4 +// (interleave_without_bias), i.e. exactly the layout the SM80 grouped GEMM consumes. Selecting +// this variant lets the decode GEMV read the *same* pre-packed buffer the SM80 prefill uses, so +// the MXFP4 expert weights are stored once instead of twice (~9 GiB saved for gpt-oss-20b). +// The un-permutation is a compile-time index remap, so it is free at runtime. +template +using Fp4KernelDetailsSm80Pair = + fiv::KernelDetails::Type, fiv::Fp4DetailsW, fiv::ColumnMajorInterleaved, true, + kTileSizeKFp4>; + +// Carries a Details type into a generic lambda so the interleaved and pair-interleaved launch +// bodies can be shared instead of duplicated. +template +struct TypeTag { + using type = U; +}; + // 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 @@ -103,7 +121,7 @@ using Fp4GemvAccT = std::conditional_t::value, half, float // 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) { + MoeGemvConfig config, bool sm80_pair_interleaved) { if (sm < 80) { return false; } @@ -123,7 +141,7 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int (n < kMinProfiledProblemDimForExpandedRowsAbove4 || k < kMinProfiledProblemDimForExpandedRowsAbove4)) { return false; } - if (Fp4MoeGemvUseInterleaved()) { + if (sm80_pair_interleaved || 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 @@ -154,6 +172,20 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int return true; } +bool is_moe_gemv_fp4_sm80_layout_supported(int64_t n, int64_t k, int group_size) { + // Same constraints the interleaved branch of is_moe_gemv_fp4_supported enforces, minus the + // row-count bounds (which PrePack cannot know): the interleaved kStepK=32 tile is tied to the + // MXFP4 block-32 scale layout, each block covers kInterleavedCtaN*kInterleave columns, and a + // complete interleaved K-tile is kStepK*kThreadsPerInterleavedTile = 64 wide. + return group_size == 32 && n % (kInterleavedCtaN * 4) == 0 && k % 64 == 0; +} + +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, + MoeGemvConfig config) { + return is_moe_gemv_fp4_supported(sm, expanded_num_rows, n, k, group_size, config, + /*sm80_pair_interleaved=*/false); +} + 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); } @@ -162,22 +194,33 @@ template void launch_moe_gemv_fp4_symmetric(const T* act, const uint8_t* weight, const T* scales, const T* bias, T* 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, - int sm, MoeGemvConfig config, cudaStream_t stream) { + int sm, MoeGemvConfig config, bool sm80_pair_interleaved, 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. + // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller + // CtaN. Taken either via the opt-in env knob or because the caller pre-packed a single + // SM80-grouped-GEMM buffer (sm80_pair_interleaved) that this kernel un-permutes while decoding. + // 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); + if (sm80_pair_interleaved || Fp4MoeGemvUseInterleaved()) { + auto launch_interleaved = [&](auto details_tag) { + using DetailsI = typename decltype(details_tag)::type; + 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); + } + }; + if (sm80_pair_interleaved) { + launch_interleaved(TypeTag>{}); } 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); + launch_interleaved(TypeTag>{}); } return; } @@ -205,26 +248,37 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const T* act, const uint8_t* weight, const T* scales, const T* bias, T* 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, int sm, - cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, cudaStream_t stream) { + cutlass_kernels::ActivationParams activation_params, MoeGemvConfig config, bool sm80_pair_interleaved, + 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); + // Interleaved path: ColumnMajorInterleaved layout + dtype-conditional accumulation + smaller + // CtaN, fusing SwiGLU. Taken either via the opt-in env knob or because the caller pre-packed a + // single SM80-grouped-GEMM buffer (sm80_pair_interleaved). SwiGLU fusion is orthogonal to the + // weight layout: it only concerns the fc1 output-column (gate/value) order, so both variants + // fuse it identically. Takes precedence over the split-K path so the kernel matches the + // 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 (sm80_pair_interleaved || Fp4MoeGemvUseInterleaved()) { + auto launch_interleaved = [&](auto details_tag) { + using DetailsI = typename decltype(details_tag)::type; + 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); + } + }; + if (sm80_pair_interleaved) { + launch_interleaved(TypeTag>{}); } 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); + launch_interleaved(TypeTag>{}); } return; } @@ -248,18 +302,18 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( template void launch_moe_gemv_fp4_symmetric( const half*, const uint8_t*, const half*, const half*, half*, const int64_t*, const int*, int, - int64_t, int64_t, int64_t, int, int, MoeGemvConfig, cudaStream_t); + int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, cudaStream_t); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const half*, const uint8_t*, const half*, const half*, half*, const int64_t*, const int*, int, - int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, MoeGemvConfig, cudaStream_t); + int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, MoeGemvConfig, bool, cudaStream_t); #ifdef ENABLE_BF16 template void launch_moe_gemv_fp4_symmetric<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, - const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, cudaStream_t); + const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, MoeGemvConfig, bool, cudaStream_t); template void launch_moe_gemv_fp4_symmetric_interleaved_swiglu<__nv_bfloat16>( const __nv_bfloat16*, const uint8_t*, const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const int64_t*, const int*, int, int64_t, int64_t, int64_t, int, int, cutlass_kernels::ActivationParams, - MoeGemvConfig, cudaStream_t); + MoeGemvConfig, bool, cudaStream_t); #endif } // namespace moe_gemv 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 index 7b6e9d12adace..7239bd94f94f5 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h +++ b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h @@ -41,6 +41,17 @@ bool Fp4MoeGemvUseInterleaved(); 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); +// `sm80_pair_interleaved` = the weights were pre-packed in the SM80 grouped-GEMM layout (the +// interleaved layout plus the [e0,e2,e4,e6,e1,e3,e5,e7] nibble pair-interleave) and the GEMV will +// un-permute it in-register. It implies the interleaved shape rules (MXFP4 group_size 32, +// n % 16 == 0, k % 64 == 0) regardless of ORT_FP4_GEMV_INTERLEAVED. +bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int64_t k, int group_size, + MoeGemvConfig config, bool sm80_pair_interleaved); + +// Layout-only (row-count independent) form of the rules above, for PrePack: true when a GEMV can +// decode an [n, k] MXFP4 problem straight out of the SM80 grouped-GEMM pair-interleaved buffer, so +// no separate ColToRow decode copy of the expert weights has to be allocated. +bool is_moe_gemv_fp4_sm80_layout_supported(int64_t n, int64_t k, int group_size); // Launches the MXFP4 (e2m1) MoE GEMV in the non-interleaved ColumnMajor layout. // act: [expanded_num_rows, k] permuted activations (row-major), T = half/bf16 @@ -52,11 +63,14 @@ bool is_moe_gemv_fp4_supported(int sm, int64_t expanded_num_rows, int64_t n, int // 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). +// When `sm80_pair_interleaved` is true, `weight` is instead the SM80 grouped-GEMM buffer produced +// by QMoE::PrePackRepackFP4Weights(gemv_interleaved=true, sm80_pair_interleave=true) and the +// kernel inverts the nibble pair-interleave while decoding, so prefill and decode share one copy. template void launch_moe_gemv_fp4_symmetric( const T* act, const uint8_t* weight, const T* scales, const T* bias, T* 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, int sm, MoeGemvConfig config, cudaStream_t stream); + int64_t n, int64_t k, int group_size, int sm, MoeGemvConfig config, bool sm80_pair_interleaved, 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 @@ -66,7 +80,7 @@ void launch_moe_gemv_fp4_symmetric_interleaved_swiglu( const T* act, const uint8_t* weight, const T* scales, const T* bias, T* 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, int sm, cutlass_kernels::ActivationParams activation_params, - MoeGemvConfig config, cudaStream_t stream); + MoeGemvConfig config, bool sm80_pair_interleaved, cudaStream_t stream); } // namespace moe_gemv } // namespace kernels diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc index cbbe15ff88975..4fda47aafb6e7 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc @@ -42,6 +42,16 @@ void LogQMoESwigluFusionRemapOnce() { }); } +void LogQMoEFp4DualLayoutOnce() { + static std::once_flag log_info; + std::call_once(log_info, []() { + LOGS_DEFAULT(INFO) << "QMoE MXFP4: this expert shape is outside the interleaved GEMV rules, so " + "the decode GEMV needs its own copy of the e2m1 weights alongside the " + "SM80 grouped-GEMM layout. This doubles the persistent MoE weight " + "footprint; shapes with k % 64 == 0 and n % 16 == 0 avoid it."; + }); +} + // 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. @@ -287,6 +297,11 @@ QMoE::QMoE(const OpKernelInfo& op_kernel_info) : CudaKernel(op_kernel_info), MoE onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_SM80_GEMM", 1) == 1; if (enable_fp4_sm80_gemm_) { enable_fp4_gemv_ = true; + // In this regime every dispatch is served by a pre-packed buffer -- prefill by the SM80 + // interleaved weights, decode by the fused GEMV -- and the dequant fallback (the only + // consumer of the raw [E, K, N/2] e2m1 initializer) is unreachable. Let PrePack release the + // initializer instead of pinning ~9 GiB of dead device memory for a 20B-class MXFP4 MoE. + release_fp4_raw_weights_ = 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 @@ -505,6 +520,17 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { // to the runner. const bool int_weights_consumed_by_prepack = is_int && !weights_prepacked_ && packed_fc1_weights_ != nullptr && packed_fc2_weights_ != nullptr; + // Same idea for MXFP4 in the SM80 grouped-GEMM regime: PrePack laid the e2m1 weights out in the + // CUTLASS interleaved layout (and, unless low-memory mode is on, an extra ColToRow copy for the + // decode GEMV), nothing reads the raw [E, K, N/2] initializer, so it was released + // (``is_packed = true``). Gate on the packed buffers actually existing so a session with + // ``session.disable_prepacking`` (where PrePack never ran and the initializer is still live) + // keeps reading the raw tensors. + const bool fp4_weights_consumed_by_prepack = + is_fp4 && release_fp4_raw_weights_ && + gemv_fp4_fc1_weights_ != nullptr && gemv_fp4_fc2_weights_ != nullptr && + fc1_weights_shape_.NumDimensions() == 3 && fc2_weights_shape_.NumDimensions() == 3; + const bool weights_consumed_by_prepack = int_weights_consumed_by_prepack || fp4_weights_consumed_by_prepack; // When ``weights_prepacked == 0`` the raw ``[E, N, K/pack]`` int weights must be // converted to the CUTLASS fpA_intB layout by PrePack before the runner can consume // them. If PrePack never ran (e.g. ``session.disable_prepacking`` is set), the prepack @@ -517,10 +543,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { "buffers were not produced (is session.disable_prepacking set?). Provide " "CUTLASS-prepacked weights with weights_prepacked=1, or enable prepacking."); } - const Tensor* fc1_experts_weights = int_weights_consumed_by_prepack ? nullptr : context->Input(2); + const Tensor* fc1_experts_weights = weights_consumed_by_prepack ? nullptr : context->Input(2); const Tensor* fc1_scales = (is_int && !packed_fc1_scales_) ? context->Input(3) : nullptr; const Tensor* fc1_experts_bias_optional = context->Input(4); - const Tensor* fc2_experts_weights = int_weights_consumed_by_prepack ? nullptr : context->Input(5); + const Tensor* fc2_experts_weights = weights_consumed_by_prepack ? nullptr : context->Input(5); const Tensor* fc2_scales = (is_int && !packed_fc2_scales_) ? context->Input(6) : nullptr; const Tensor* fc2_experts_bias_optional = context->Input(7); // The CUTLASS MoE runner has no separate FC3 GEMM — gate and up projection weights must be @@ -554,10 +580,10 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { return Status::OK(); }; - // When PrePack consumed the int weight initializers, the dtype check + // When PrePack consumed the int/MXFP4 weight initializers, the dtype check // is no longer applicable (we know they were uint8 — that's what - // PrePackIntExpertWeights validated and consumed). - if (!int_weights_consumed_by_prepack) { + // PrePackIntExpertWeights / PrePackRepackFP4Weights validated and consumed). + if (!weights_consumed_by_prepack) { ORT_RETURN_IF_ERROR(check_weight_type(fc1_experts_weights, "fc1_experts_weights", is_fp8)); ORT_RETURN_IF_ERROR(check_weight_type(fc2_experts_weights, "fc2_experts_weights", is_fp8)); } @@ -595,8 +621,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { bool is_fused_swiglu = activation_type_ == onnxruntime::llm::kernels::cutlass_kernels::ActivationType::Swiglu; MoEParameters moe_params; // Prefer the cached shapes when PrePack consumed the source initializer. - const TensorShape& fc1_shape = int_weights_consumed_by_prepack ? fc1_weights_shape_ : fc1_experts_weights->Shape(); - const TensorShape& fc2_shape = int_weights_consumed_by_prepack ? fc2_weights_shape_ : fc2_experts_weights->Shape(); + const TensorShape& fc1_shape = weights_consumed_by_prepack ? fc1_weights_shape_ : fc1_experts_weights->Shape(); + const TensorShape& fc2_shape = weights_consumed_by_prepack ? fc2_weights_shape_ : fc2_experts_weights->Shape(); ORT_RETURN_IF_ERROR(onnxruntime::contrib::moe_helper::CheckInputs( moe_params, input, router_probs, &fc1_shape, fc1_experts_bias_optional, fc1_scales, fc1_zeros, @@ -756,6 +782,17 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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; + // Releasing the raw e2m1 initializers (fp4_weights_consumed_by_prepack) is only sound because + // every dispatch is served by a pre-packed buffer. If the SM80 prefill buffers were not fully + // built (e.g. the group scales never materialized because a global-scale input was missing), + // the dequant fallback below would be selected and would dereference the released initializer. + // Fail loudly rather than crash. + if (fp4_weights_consumed_by_prepack && !fp4_sm80_prefill) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "QMoE MXFP4: PrePack released the raw expert-weight initializers but the " + "SM80 grouped-GEMM buffers are incomplete (missing pre-packed weights or " + "group scales). Set ORT_FP4_SM80_GEMM=0 to keep the dequant fallback."); + } // 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. @@ -1332,9 +1369,14 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { use_fp4_dequant_fallback_ || ((enable_fp4_cutlass_gemm_ || enable_nvfp4_cutlass_gemm_) && fp4_prefill_min_tokens_ > 0 && moe_params.num_rows < fp4_prefill_min_tokens_); + // True when the decode GEMV reads the SM80 pair-interleaved prefill buffer directly (single + // weight copy) instead of a dedicated GEMV-native copy; see gemv_fp4_fc*_reads_sm80_layout_. + const bool fc1_gemv_sm80_layout = gemv_fp4_fc1_reads_sm80_layout_; + const bool fc2_gemv_sm80_layout = gemv_fp4_fc2_reads_sm80_layout_; const bool fp4_gemv_buffers_ready = (enable_fp4_sm80_gemm_ - ? (gemv_fp4_fc1_weights_decode_ != nullptr && gemv_fp4_fc2_weights_decode_ != nullptr) + ? ((fc1_gemv_sm80_layout ? gemv_fp4_fc1_weights_ != nullptr : gemv_fp4_fc1_weights_decode_ != nullptr) && + (fc2_gemv_sm80_layout ? gemv_fp4_fc2_weights_ != 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 && @@ -1349,16 +1391,21 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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); + const bool fc1_gemv_supported = gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, hidden, gemv_group_size, + gemv::MoeGemvConfig::kDefault, fc1_gemv_sm80_layout); + const bool fc2_gemv_supported = gemv::is_moe_gemv_fp4_supported(sm_, expanded, hidden, inter, gemv_group_size, + gemv::MoeGemvConfig::kDefault, fc2_gemv_sm80_layout); 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). + // interleaved layout, the decode GEMV either un-permutes that same buffer in-register + // (fc*_gemv_sm80_layout, the single-copy default) or reads the dedicated GEMV-native copy + // in gemv_fp4_fc*_weights_decode_. const uint8_t* gemv_fc1_weight = static_cast( - (enable_fp4_sm80_gemm_ ? gemv_fp4_fc1_weights_decode_ : gemv_fp4_fc1_weights_).get()); + ((enable_fp4_sm80_gemm_ && !fc1_gemv_sm80_layout) ? 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()); + ((enable_fp4_sm80_gemm_ && !fc2_gemv_sm80_layout) ? 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)); @@ -1425,7 +1472,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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); + p_efto, p_exp, num_experts, expanded, inter, hidden, gemv_group_size, sm_, act_params, cfg, + fc1_gemv_sm80_layout, stream); }; auto launch_fc2 = [&](MoeGemvConfig cfg) { gemv::launch_moe_gemv_fp4_symmetric( @@ -1433,7 +1481,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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); + p_efto, p_exp, num_experts, expanded, hidden, inter, gemv_group_size, sm_, cfg, + fc2_gemv_sm80_layout, stream); }; if (do_tune) { @@ -1469,7 +1518,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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)) { + if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, fc1_n, hidden, gemv_group_size, cfg, + fc1_gemv_sm80_layout)) { continue; } const float ms = time_launch([&] { launch_fc1(cfg); }); @@ -1487,7 +1537,8 @@ Status QMoE::ComputeInternal(OpKernelContext* context) const { 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)) { + if (!gemv::is_moe_gemv_fp4_supported(sm_, expanded, hidden, inter, gemv_group_size, cfg, + fc2_gemv_sm80_layout)) { continue; } const float ms = time_launch([&] { launch_fc2(cfg); }); @@ -1771,14 +1822,17 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, } 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 + // Fused FP4 GEMV: lay out fc1 weights as [E, 2*inter, hidden/2] row-major. Unless + // release_fp4_raw_weights_ is set, 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. + // so this buffer feeds the prefill grouped GEMM. The decode GEMV reads that same buffer and + // inverts the nibble pair-interleave in-register (Fp4KernelDetailsSm80Pair), so one copy of + // the e2m1 weights serves both regimes; only when the shape misses the interleaved GEMV rules + // do we pack a dedicated GEMV-native copy into gemv_fp4_fc1_weights_decode_. 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_); @@ -1786,12 +1840,27 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, 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_) { + if (sm80_pair) { + // Raw initializer is [E, K, N/2]; the GEMV problem is n = N, k = K. + const auto& raw_shape = tensor.Shape(); + gemv_fp4_fc1_reads_sm80_layout_ = onnxruntime::llm::kernels::moe_gemv::is_moe_gemv_fp4_sm80_layout_supported( + raw_shape[2] * 2, raw_shape[1], /*group_size=*/32); + if (!gemv_fp4_fc1_reads_sm80_layout_) { + LogQMoEFp4DualLayoutOnce(); + } + } + if (enable_fp4_sm80_gemm_ && !gemv_fp4_fc1_reads_sm80_layout_) { 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); } + if (release_fp4_raw_weights_ && local_packed) { + // Nothing reads the raw [E, K, N/2] e2m1 initializer any more (see release_fp4_raw_weights_); + // cache its shape for CheckInputs and let ORT free it. + fc1_weights_shape_ = tensor.Shape(); + is_packed = true; + } } else if (input_idx == 5 && (quant_type_ == "fp4" || quant_type_ == "nvfp4") && enable_fp4_gemv_) { const bool nvfp4 = (quant_type_ == "nvfp4"); const bool use_interleave = @@ -1800,12 +1869,24 @@ Status QMoE::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc, 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_) { + if (sm80_pair) { + const auto& raw_shape = tensor.Shape(); + gemv_fp4_fc2_reads_sm80_layout_ = onnxruntime::llm::kernels::moe_gemv::is_moe_gemv_fp4_sm80_layout_supported( + raw_shape[2] * 2, raw_shape[1], /*group_size=*/32); + if (!gemv_fp4_fc2_reads_sm80_layout_) { + LogQMoEFp4DualLayoutOnce(); + } + } + if (enable_fp4_sm80_gemm_ && !gemv_fp4_fc2_reads_sm80_layout_) { 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); } + if (release_fp4_raw_weights_ && local_packed) { + fc2_weights_shape_ = tensor.Shape(); + is_packed = true; + } } 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 @@ -2365,6 +2446,16 @@ void QMoE::TryBuildGemvFp4Scales(int fc, cudaStream_t stream, AllocatorPtr alloc } } CUDA_CALL_THROW(cudaStreamSynchronize(stream)); + + // In the SM80 grouped-GEMM regime the combined [E, k/32, n] activation-dtype scales built above + // are the only scales any dispatch reads (prefill via QuantParams::GroupWise, decode via the + // fused GEMV), and the dequant fallback that consumes the raw e8m0 block scales is unreachable. + // Drop the staged e8m0 copy -- for a 20B-class MXFP4 MoE that is a few hundred MiB of dead + // device memory. ``block`` is the buffer ``src_block`` aliases in this regime (raw_block is only + // populated by the native CUTLASS paths, which do not set release_fp4_raw_weights_). + if (release_fp4_raw_weights_ && raw_block == nullptr) { + block.reset(); + } } // --------------------------------------------------------------------------- diff --git a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h index e5d1c2fb3d340..17d212b67b989 100644 --- a/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h +++ b/onnxruntime/contrib_ops/cuda/moe/moe_quantization.h @@ -81,8 +81,8 @@ class QMoE final : public CudaKernel, public MoEBase { // 1 is reserved for a possible future Hopper-specific layout (e.g. W4A8). bool weights_prepacked_ = true; // Cached source weight shapes captured at PrePack time. When the - // PrePack hook consumed and released the original int4/int8 weight - // initializers (``is_packed = true``), ``context->Input(2)`` + // PrePack hook consumed and released the original int4/int8 (or MXFP4) + // weight initializers (``is_packed = true``), ``context->Input(2)`` // and ``(5)`` return nothing, so ``moe_helper::CheckInputs`` can no // longer read the shapes from the live tensors. We feed it these // cached shapes instead via the ``TensorShape*`` overload, matching @@ -156,6 +156,22 @@ class QMoE final : public CudaKernel, public MoEBase { // 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 true, PrePack reports ``is_packed = true`` for the e2m1 weight initializers (inputs 2/5) + // so ORT releases them. Safe only in the ``enable_fp4_sm80_gemm_`` regime, where prefill reads + // gemv_fp4_fc*_weights_ and decode reads either that same buffer (pair-interleaved un-permute) + // or gemv_fp4_fc*_weights_decode_, and the dequant fallback -- the only consumer of the raw + // [E, K, N/2] layout -- is unreachable. For a 20B-class MXFP4 MoE the retained initializers are + // ~9 GiB of otherwise dead device memory. NOTE: the memory is returned to the *device* only when + // initializers bypass the BFC arena, i.e. with the session option + // ``session.use_device_allocator_for_initializers = 1``; otherwise it is merely recycled inside + // the arena for later activation/KV allocations. + bool release_fp4_raw_weights_ = false; + // Set by PrePack (per weight tensor) when the decode GEMV will read the SM80 pair-interleaved + // buffer in gemv_fp4_fc*_weights_ instead of a dedicated gemv_fp4_fc*_weights_decode_ copy. + // False when SM80 GEMM is off (gemv_fp4_fc*_weights_ is already GEMV-native) or when the shape + // misses the interleaved rules. + bool gemv_fp4_fc1_reads_sm80_layout_ = false; + bool gemv_fp4_fc2_reads_sm80_layout_ = 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. @@ -175,10 +191,11 @@ class QMoE final : public CudaKernel, public MoEBase { 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. + // prefill (SM80 pair-interleaved layout), these hold a dedicated decode-GEMV copy of the e2m1 + // weights in the GEMV-native layout (ColToRow, or the interleaved layout's preprocessor steps + // 1-3). Only allocated when gemv_fp4_fc*_reads_sm80_layout_ is false, i.e. when the GEMV cannot + // un-permute the pair-interleaved buffer itself. 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 diff --git a/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py index 9e3db03c76d68..dd91b3772b8ce 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py @@ -13,6 +13,7 @@ # Requires SM90+ (Hopper or newer) and CUDA 12.8+ (ENABLE_FP4 build flag). # -------------------------------------------------------------------------- +import contextlib import math import os import unittest @@ -825,6 +826,242 @@ def run(enable_native): self.assertEqual(max_diff, 0.0) +# ============================================================================ +# SM80 grouped-GEMM regime: one weight copy, raw initializers released +# ============================================================================ + + +@contextlib.contextmanager +def _env_overrides(**overrides): + """Temporarily set environment variables; a value of None removes the variable. + + The QMoE FP4 dispatch knobs are read in the operator constructor, i.e. while the + InferenceSession is being built, so an override only has to be live for that call. + """ + previous = {name: os.environ.get(name) for name in overrides} + try: + for name, value in overrides.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + yield + finally: + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +def _quantize_random_mxfp4_experts(num_experts, fc1_n, fc1_k, fc2_n, fc2_k, block_size=32, seed=42): + """Quantize random per-expert FC1/FC2 weights to MXFP4. + + Returns ``(packed, block_scales, global_scales, dequantized)``; each is a dict keyed by + "fc1"/"fc2" holding the stacked ``[E, ...]`` QMoE initializer tensors, and ``dequantized`` + holds the ``[E, N, K]`` weights the PyTorch reference multiplies with. + """ + torch.manual_seed(seed) + packed, block_scales, global_scales, dequantized = {}, {}, {}, {} + for fc, n, k in (("fc1", fc1_n, fc1_k), ("fc2", fc2_n, fc2_k)): + p, b, g, d = [], [], [], [] + for _ in range(num_experts): + pe, be, ge, de = quantize_weight_to_mxfp4(torch.randn(n, k, device=device) * 0.1, block_size) + p.append(pe) + b.append(be) + g.append(torch.tensor(ge, dtype=torch.float32)) + d.append(de) + packed[fc] = torch.stack(p, dim=0) + block_scales[fc] = torch.stack(b, dim=0) + global_scales[fc] = torch.stack(g) + dequantized[fc] = torch.stack(d, dim=0) + return packed, block_scales, global_scales, dequantized + + +@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") +class TestQMoEFP4Sm80SingleWeightCopy(unittest.TestCase): + """MXFP4 SM80 grouped-GEMM regime (ORT_FP4_SM80_GEMM=1, the default for 80 <= SM < 120). + + PrePack packs one pair-interleaved e2m1 buffer that serves both prefill (the SM80 grouped + GEMM) and decode (the fused GEMV inverts the nibble pair-interleave in-register -- + ``gemv_fp4_fc*_reads_sm80_layout_``), and then reports ``is_packed = true`` so ORT releases + the raw ``[E, K, N/2]`` initializers (``release_fp4_raw_weights_``). + + Two properties have to hold for that to be safe, one test each: + + * every dispatch stays numerically equivalent to ORT_FP4_SM80_GEMM=0, which keeps the raw + initializers and serves decode from a GEMV-native weight copy and prefill from the + dequant fallback -- ``test_sm80_parity_vs_dequant_fallback``; + * the release actually returns the initializer bytes to the device when initializers bypass + the BFC arena -- ``test_sm80_release_frees_raw_weight_initializers``. + """ + + NUM_EXPERTS = 8 + TOP_K = 2 + + def _skip_unless_sm80_regime(self): + sm = _cuda_sm() + if sm < 90: + self.skipTest(f"FP4 requires SM90+, got SM{sm}") + if sm >= 120: + self.skipTest(f"The SM80 MXFP4 grouped GEMM is only used below SM120, got SM{sm}") + + def _build_model(self, hidden_size, inter_size, num_tokens, onnx_dtype, num_experts=None): + """SwiGLU MXFP4 MoE model. FC1 is [E, hidden, inter] packed (n = 2 * inter_size).""" + num_experts = self.NUM_EXPERTS if num_experts is None else num_experts + packed, block_scales, global_scales, dequantized = _quantize_random_mxfp4_experts( + num_experts, + fc1_n=2 * inter_size, + fc1_k=hidden_size, + fc2_n=hidden_size, + fc2_k=inter_size, + ) + model = create_fp4_moe_onnx_graph( + num_tokens=num_tokens, + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=self.TOP_K, + onnx_dtype=onnx_dtype, + fc1_weights=packed["fc1"], + fc2_weights=packed["fc2"], + fc1_block_scales=block_scales["fc1"], + fc1_global_scale=global_scales["fc1"], + fc2_block_scales=block_scales["fc2"], + fc2_global_scale=global_scales["fc2"], + use_swiglu=True, + ) + return model, dequantized + + def _run(self, model, sm80_mode, input_tensor, router_logits, onnx_dtype): + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + with _env_overrides(ORT_FP4_SM80_GEMM=sm80_mode): + session = onnxruntime.InferenceSession( + model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + output = torch.zeros_like(input_tensor) + 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.shape, output.data_ptr()) + iobinding.synchronize_inputs() + session.run_with_iobinding(iobinding) + iobinding.synchronize_outputs() + return output + + @parameterized.expand( + [ + ("fp16_decode_1_token", TensorProto.FLOAT16, 1), + ("fp16_decode_4_tokens", TensorProto.FLOAT16, 4), + ("bf16_decode_1_token", TensorProto.BFLOAT16, 1), + ("bf16_decode_4_tokens", TensorProto.BFLOAT16, 4), + ("fp16_prefill_64_tokens", TensorProto.FLOAT16, 64), + ("bf16_prefill_64_tokens", TensorProto.BFLOAT16, 64), + ("fp16_prefill_256_tokens", TensorProto.FLOAT16, 256), + ] + ) + def test_sm80_parity_vs_dequant_fallback(self, _name, onnx_dtype, num_tokens): + """Cross-check the released-weights regime against the one that keeps the raw weights. + + hidden = inter = 512 with SwiGLU gives fc1 (n=1024, k=512) and fc2 (n=512, k=512), both + of which satisfy ``is_moe_gemv_fp4_sm80_layout_supported`` (n % 16 == 0, k % 64 == 0), + so with SM80 GEMM on there is exactly one copy of the e2m1 weights and the raw + initializers are released. ``num_tokens * top_k <= 8`` picks the fused decode GEMV -- + which un-permutes the pair-interleaved buffer with SM80 on and reads the GEMV-native + buffer with SM80 off. The larger token counts exceed that bound, so they compare the + SM80 grouped GEMM against the dequant fallback instead. + """ + self._skip_unless_sm80_regime() + hidden_size, inter_size = 512, 512 + torch_dtype = torch.float16 if onnx_dtype == TensorProto.FLOAT16 else torch.bfloat16 + model, dequantized = self._build_model(hidden_size, inter_size, num_tokens, onnx_dtype) + + torch.manual_seed(7) + input_tensor = torch.randn(num_tokens, hidden_size, device=device, dtype=torch_dtype) + router_logits = torch.randn(num_tokens, self.NUM_EXPERTS, device=device, dtype=torch_dtype) + + sm80_output = self._run(model, "1", input_tensor, router_logits, onnx_dtype).float() + fallback_output = self._run(model, "0", input_tensor, router_logits, onnx_dtype).float() + reference = TestQMoEFP4._compute_reference( + input_tensor, + router_logits, + dequantized["fc1"], + dequantized["fc2"], + self.NUM_EXPERTS, + self.TOP_K, + True, + torch_dtype, + ).float() + + sm80_diff = (sm80_output - reference).abs().max().item() + fallback_diff = (fallback_output - reference).abs().max().item() + cross_diff = (sm80_output - fallback_output).abs().max().item() + print( + f"FP4 SM80 parity: {_name} tokens={num_tokens} " + f"sm80_vs_ref={sm80_diff:.6f} fallback_vs_ref={fallback_diff:.6f} cross={cross_diff:.6f}" + ) + + # FP4 quantization is lossy; same tolerance the other MXFP4 parity tests use. + atol = 0.15 if torch_dtype == torch.bfloat16 else 0.12 + self.assertLess(fallback_diff, atol, f"ORT_FP4_SM80_GEMM=0 vs reference: {fallback_diff:.6f}") + self.assertLess(sm80_diff, atol, f"ORT_FP4_SM80_GEMM=1 vs reference: {sm80_diff:.6f}") + self.assertLess(cross_diff, atol, f"SM80 grouped GEMM vs dequant fallback: {cross_diff:.6f}") + + def test_sm80_release_frees_raw_weight_initializers(self): + """The released raw e2m1 initializers must actually come back to the device. + + The saving is only observable when initializers bypass the BFC arena, i.e. with + ``session.use_device_allocator_for_initializers=1``; otherwise the arena merely recycles + the block for later activation allocations. The CUDA allocator is shared across sessions + in a process, so a throwaway ORT_FP4_SM80_GEMM=0 session (the configuration with the + larger pre-pack footprint) is created first to grow the arena; after that the pre-pack + buffers of the measured sessions are served from it and only the initializer + cudaMalloc/cudaFree moves free device memory. + """ + self._skip_unless_sm80_regime() + hidden_size, inter_size, num_experts = 2048, 1024, 16 + onnx_dtype = TensorProto.FLOAT16 + model, _ = self._build_model(hidden_size, inter_size, 1, onnx_dtype, num_experts=num_experts) + + # Raw e2m1 initializers PrePack can release: fc1 [E, hidden, inter] + fc2 [E, inter, hidden/2]. + raw_weight_bytes = num_experts * (hidden_size * inter_size + inter_size * (hidden_size // 2)) + + def device_bytes_for_session(sm80_mode): + opts = onnxruntime.SessionOptions() + opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL + opts.add_session_config_entry("session.use_device_allocator_for_initializers", "1") + torch.cuda.synchronize() + free_before = torch.cuda.mem_get_info(0)[0] + with _env_overrides(ORT_FP4_SM80_GEMM=sm80_mode): + session = onnxruntime.InferenceSession( + model, opts, providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + torch.cuda.synchronize() + used = free_before - torch.cuda.mem_get_info(0)[0] + del session + return used + + device_bytes_for_session("0") # warm the shared CUDA arena with the larger configuration + retained = device_bytes_for_session("0") + released = device_bytes_for_session("1") + saved = retained - released + mib = 1024 * 1024 + print( + f"FP4 SM80 initializer release: raw_weights={raw_weight_bytes / mib:.1f} MiB " + f"retained_session={retained / mib:.1f} MiB released_session={released / mib:.1f} MiB " + f"saved={saved / mib:.1f} MiB" + ) + self.assertGreater( + saved, + raw_weight_bytes // 2, + f"ORT_FP4_SM80_GEMM=1 only saved {saved / mib:.1f} MiB of the " + f"{raw_weight_bytes / mib:.1f} MiB of raw e2m1 initializers", + ) + + # ============================================================================ # Standalone packing utility tests # ============================================================================