[CUDA] QMoE: release raw MXFP4 weight initializers after PrePack

[tianleiwu]

Description

On SM80–SM119 the MXFP4 QMoE node used to keep up to three persistent copies of the expert
weights: the raw [E, K, N/2] initializers (only ever read by the dequant fallback), the SM80
pair-interleaved buffer consumed by the grouped-GEMM prefill, and a GEMV-native buffer consumed
by the fused decode GEMV. For a 20B-class MXFP4 MoE each copy is ~9 GiB, which put post-load
VRAM out of reach of 24 GB consumer cards.

This PR collapses that to a single copy in the default ORT_FP4_SM80_GEMM=1 regime:

1. The decode GEMV learns to read the prefill 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 (Fp4I2FConverter<AType, PairInterleaved>), so there are no
   extra branches, registers, or memory traffic. gemv_fp4_fc{1,2}_reads_sm80_layout_ records
   per-FC whether the dedicated gemv_fp4_fc*_weights_decode_ copy can be skipped.
2. PrePack releases the raw initializers. With both prefill and decode served from
   pre-packed buffers, the dequant fallback — the only consumer of the raw layout — is
   unreachable, so PrePack reports is_packed = true for inputs 2/5 and caches
   fc*_weights_shape_ so moe_helper::CheckInputs can still validate shapes. The staged e8m0
   block-scale copy is dropped in TryBuildGemvFp4Scales for the same reason.

A defensive guard in ComputeInternal returns a descriptive error (pointing at
ORT_FP4_SM80_GEMM=0) if the raw weights were released but the SM80 buffers are somehow
incomplete, rather than dereferencing a null initializer.

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 (this PR, default)	13996 MiB

Released initializers only shrink the process's device footprint when initializers bypass the
BFC arena, i.e. with the session option session.use_device_allocator_for_initializers = 1.
Otherwise the freed bytes are recycled inside the arena for later activation/KV allocations.
This is documented in docs/contrib_ops/cuda/moe_qmoe.md §9.11.

Shape gate and fallback

Reusing the interleaved GEMV (kInterleave=4, kStepK=32) inherits its 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 keep the
dedicated decode-layout copy (logged at INFO). NVFP4 (block 16) always uses the plain ColToRow
layout and is unaffected. gpt-oss-20b (k=2880, n=5760/2880) clears the gate.

Setting ORT_FP4_SM80_GEMM=0 restores the previous behavior: raw initializers retained, prefill
on the dequant fallback, decode on the GEMV-native copy.

Motivation and Context

Makes 20B-class MXFP4 MoE models loadable on 24 GB consumer GPUs, and removes ~9 GiB of dead
device memory on every SM80–SM119 deployment.

Testing

New TestQMoEFP4Sm80SingleWeightCopy in
onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py (gated on a build with
--use_fp4_qmoe and on 80 <= SM < 120):

• test_sm80_parity_vs_dequant_fallback — fp16/bf16 × 1/4/64/256 tokens, hidden = inter = 512, 8 experts, top-2, SwiGLU. Every case is run twice, with ORT_FP4_SM80_GEMM=1 and =0,
  and both are compared against a PyTorch reference built from the dequantized weights. The
  1/4-token cases exercise the fused decode GEMV un-permuting the pair-interleaved buffer; the
  64/256-token cases exercise the SM80 grouped GEMM against the dequant fallback.
• test_sm80_release_frees_raw_weight_initializers — 16 experts, hidden = 2048,
  inter = 1024, session.use_device_allocator_for_initializers=1. A throwaway session
  pre-warms the shared CUDA arena, then session-creation device-memory deltas are compared
  between the retaining and releasing configurations.

Results on H200 (SM90), CUDA 13.0 — 8 passed:

case	SM80-on vs ref	SM80-off vs ref	cross
fp16 decode 1 / 4 tok	0.031 / 0.033	0.018 / 0.018	0.031 / 0.037
bf16 decode 1 / 4 tok	0.125	0.125	0.000
fp16 prefill 64 / 256 tok	0.016	0.016	0.016
bf16 prefill 64 tok	0.125	0.125	0.125

Memory test: raw e2m1 initializers 48.0 MiB; retaining session 172.0 MiB vs releasing session
120.0 MiB → 52.0 MiB returned to the device.

Notes for reviewers

• onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.{cu,h} also change in [CUDA] Address feedback for block-scaled MatMul and MoE GEMV #29910; expect
  a conflict depending on merge order.

[Copilot]

Pull request overview

This PR reduces persistent device-memory usage for MXFP4 QMoE on SM80–SM119 by enabling both prefill (SM80 grouped GEMM) and decode (fused GEMV) to share the same SM80 pair-interleaved prepacked weight buffers, and then allowing ORT to release the original raw FP4 weight initializers after PrePack.

Changes:

• Update FP4 GEMV decode to optionally read the SM80 pair-interleaved (prefill) weight layout by in-register un-permutation, avoiding a second decode-layout weight copy when shapes satisfy the interleaved rules.
• Enable releasing raw MXFP4 weight initializers after PrePack in the SM80 grouped-GEMM regime, with a defensive runtime guard to fail fast if buffers are unexpectedly incomplete.
• Add Python CUDA tests validating numerical parity vs. the dequant fallback and validating that raw initializer release returns device memory (when initializers bypass the arena), plus documentation for the single-copy regime.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py	Adds SM80-regime tests for parity vs fallback and for device-memory reduction when raw initializers are released.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Adds state/flags tracking raw-initializer release and whether GEMV reads the SM80 layout.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Implements raw-initializer release via PrePack, routes decode GEMV to shared weights when possible, and adds a safety guard if fallback would need released tensors.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h	Extends GEMV APIs to support an SM80 pair-interleaved weight mode and adds a layout-only shape gate helper for PrePack.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu	Adds the SM80-pair GEMV dispatch variant and the layout-only support predicate; threads the new flag through launch paths.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h	Extends FP4 decode to optionally invert the SM80 pair-interleave at compile time via a templated converter.
docs/contrib_ops/cuda/moe_qmoe.md	Documents the new SM80 single-copy weights behavior, shape gate, and the initializer-release interaction with the device allocator / arena.