[CUDA] Add NVFP4 support in QMoE

[tianleiwu]

Description

Adds NVFP4 support to the CUDA QMoE operator for W4A16 mixture-of-experts inference. The implementation accepts packed FP4 expert weights with E4M3 block scales and per-expert global scales, then dispatches to the appropriate GEMV/GEMM path while retaining the dequantization fallback where native block-scaled kernels are unavailable.

Key changes:

• Extend the QMoE schema and runtime validation with quant_type="nvfp4", including block-size, block-scale, and global-scale shape checks.
• Add CUDA NVFP4 preprocessing, dequantization, QMoE kernels, and FP4 GEMV support.
• Extend CUTLASS-based MoE dispatch and numeric-conversion utilities for the FP4 layouts and scale handling needed by NVFP4.
• Document the NVFP4 tensor layout, quantization parameters, and execution behavior.
• Add CUDA tests for FP16/BF16, SiLU/SwiGLU, GEMV decode and disabled-GEMV paths, varying token/expert/dimension configurations, and malformed NVFP4 inputs.
• Expand existing FP4 coverage for row-varying native CUTLASS scales and invalid hidden dimensions.

Motivation and Context

NVFP4 supplies a block-scaled FP4 weight format intended for low-memory MoE inference. Supporting it in QMoE enables models using this representation to run through the CUDA execution provider with explicit validation of its scale metadata and coverage for both prefill and decode-shaped workloads.

Testing

• git diff --check origin/main...HEAD
• python /home/tianlei/git/dev/scripts/h200_18/bench_qmoe_fp4_compare.py --iters 400 --warmup 40 --with-fallback

[Copilot]

Pull request overview

This PR extends the CUDA QMoE implementation to support NVFP4 (E2M1 FP4 weights with E4M3 block scales and per-expert FP32 global scales) for W4A16 MoE inference, including runtime validation, CUDA dequant/scale-combine kernels, GEMV/GEMM dispatch updates, tests, and documentation.

Changes:

• Add quant_type="nvfp4" to the QMoE schema/docs and implement CUDA kernels for NVFP4 dequantization + GEMV scale folding.
• Expand FP4 execution paths (GEMV decode + SM80/SM90 dispatch decisions, FP4 numeric conversion, and preprocessing/layout support).
• Add/extend Python tests covering NVFP4 parity, GEMV on/off behavior, and additional FP4 validation/regression cases.

Reviewed changes

Copilot reviewed 28 out of 29 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py	New CUDA test coverage for NVFP4 end-to-end parity + malformed input validation + GEMV on/off checks.
onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py	Adds FP4 validation and additional coverage for GEMV decode and native CUTLASS scale behavior.
onnxruntime/test/python/transformers/test_qmoe_cuda.py	Updates an INT quantization test to reflect earlier (prepack-time) rejection behavior and updated error expectations.
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Extends QMoE schema docs/type constraints to include nvfp4 and float8e4m3fn scales.
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h	Declares new CUDA helpers for FP4/NVFP4 scale combining, TMA scale packing, and NVFP4 dequantization.
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu	Implements FP4/NVFP4 scale folding kernels, FP4 TMA scale packing, and NVFP4 dequantization kernels.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Extends QMoE prepack and tuning state to support FP4 GEMV, SM80 routing, and additional packed buffers.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_tma_warp_specialized_traits.h	Adjusts Ampere specialization rules to allow FP4 weights with FP16/BF16 activations.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h	Adds runner API to push the “use SM80 FP4” decision from QMoE into the CUTLASS runner.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu	Updates TMA pointer math, alignment guards, finalize fusion gating, and profiler quant-param setup for FP4 routing changes.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h	New header for MXFP4 GEMV (including an opt-in interleaved layout path) used by QMoE decode.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu	New MXFP4/NVFP4 GEMV implementation and shape gating, including interleaved option and dtype-conditional accumulation.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh	New shared device-side GEMV kernels/dispatch helpers reused by FP4 GEMV launchers.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_template_dispatch.h	Enables SM80 FP4 path selection, adjusts dispatch behavior, and improves an error message cast.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_kernels.h	Adds internal knobs/state for SM80 FP4 routing and updates groupwise parameter metadata.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/launchers/moe_gemm_tma_ws_mixed_input_launcher.inl	Formatting cleanup and adjusts alpha scaling pointer selection logic for mixed-input grouped GEMM.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h	Adds FP4 weight details and a fast FP4 (E2M1) decoder used by GEMV paths.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h	Extends weight preprocessor API to optionally skip bias+interleave and/or do interleave-without-bias for FP4.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu	Implements interleave-without-bias kernel and wires new preprocessor options.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h	Adds fast converters for interleaved FP4 weights (e2m1) to half/bf16 for fused-dequant GEMM.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h	Adds CUTLASS DqMma specializations for FP16 activations + FP4 weights.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h	Adds CUTLASS DqMma specializations for BF16 activations + FP4 weights.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h	Permits float_e2m1_t as a supported dequantized-B element type.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h	Permits float_e2m1_t as a supported dequantized-B element type.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/moe_cutlass_kernel.h	Treats FP4 weights like other quantized-B cases w.r.t. scale requirements.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h	Defines the FP4 (e2m1) interleaved B layout for mixed GEMM on Turing+.
onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/detail/collective/mixed_input_utils.hpp	Fixes inline PTX syntax (lop3.b32).
docs/contrib_ops/cuda/moe_qmoe.md	Documents NVFP4 format/layout/dispatch and adds detailed FP4 runtime knob documentation.

[kunal-vaishnavi]

From Copilot:

Merge gate checklist

✅ Must-pass correctness gates

• All existing CUDA QMoE tests pass with quant_type != "nvfp4" (regression guard for pre-existing paths).
• All new NVFP4 tests pass for:
  • FP16 and BF16
  • SiLU and SwiGLU
  • GEMV decode-shaped workloads
  • GEMV-disabled / GEMM fallback workloads
  • varying token/expert/hidden-dimension configurations
• Malformed NVFP4 input tests pass (shape/metadata validation):
  • invalid block size
  • invalid block-scale shape
  • invalid global-scale shape
  • invalid hidden-dimension / layout constraints
• Deterministic parity checks pass (fixed seeds) between:
  • native block-scaled execution path
  • dequantization fallback path
    with documented tolerances per dtype (FP16/BF16).

✅ Must-pass dispatch and validation gates

• quant_type="nvfp4" is validated consistently at schema + runtime boundaries.
• Native-kernel eligibility predicate is explicit and covered by tests for each disqualifier.
• Unsupported NVFP4 combinations fail fast with actionable error messages (no silent fallback unless intended and documented).
• Fallback behavior is explicitly tested where native block-scaled kernels are unavailable.

✅ Must-pass performance gates

• Benchmark run completed for representative decode and prefill shapes (as described in PR).
• NVFP4 path shows expected perf/memory benefit vs baseline for target workloads.
• No unacceptable perf regression on non-NVFP4 QMoE paths.
• No major perf cliff when routing between GEMV/GEMM/native/fallback for nearby shapes.

✅ Must-pass numerical stability gates

• Max/mean error metrics vs reference are within agreed thresholds for all tested shapes.
• No outlier instability on extreme but valid scale ranges.
• Row-varying scale handling remains correct for existing FP4 coverage after this change.

✅ Code quality / maintainability gates

• Validation logic for NVFP4 invariants is centralized or clearly non-duplicative.
• Scale semantics (block scale vs global scale order-of-application) are clearly documented in code/comments/docs.
• New docs accurately describe tensor layout, quantization params, and execution behavior.
• CI is green and no unresolved blocking review comments remain.

✅ Sign-off criteria

• Numerics sign-off: parity and tolerance results reviewed and approved.
• Performance sign-off: benchmark evidence reviewed and approved.
• Runtime safety sign-off: malformed input/unsupported path behavior reviewed and approved.
• Final reviewer recommendation: Approve / Request changes recorded with rationale.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 29 out of 30 changed files in this pull request and generated 5 comments.

[Copilot]

Pull request overview

Copilot reviewed 29 out of 30 changed files in this pull request and generated no new comments.

[tianleiwu]

Merge gate checklist — status update (head 62afe12468)

Updating the checklist from #29697 (comment) against the current head. Items are ticked only where there is concrete evidence in this PR (tests/validation/docs/benchmark). Items needing live CI, benchmark-evidence review, or human sign-off are left unchecked with a short note.

✅ Must-pass correctness gates

• All existing CUDA QMoE tests pass with quant_type != "nvfp4" — non-nvfp4 dispatch is untouched; test_qmoe_cuda.py / test_qmoe_fp4_cuda.py retained and extended.
• All new NVFP4 tests pass for:
  • FP16 and BF16
  • SiLU and SwiGLU
  • GEMV decode-shaped workloads (test_nvfp4_fp16/bf16_gemv_decode_swiglu)
  • GEMV-disabled / GEMM fallback workloads (test_nvfp4_fp16_gemv_disabled_swiglu)
  • varying token/expert/hidden-dimension configurations (test_nvfp4_fp16_token_counts, _more_experts, _larger_dims)
• Malformed NVFP4 input tests pass (shape/metadata validation):
  • invalid block size (test_nvfp4_rejects_wrong_block_size)
  • invalid block-scale shape (test_nvfp4_rejects_malformed_prepacked_block_scales)
  • invalid global-scale shape (test_nvfp4_rejects_malformed_prepacked_global_scale)
  • invalid hidden-dimension / layout constraints — covered for FP4 (test_fp4_rejects_non_32_multiple_hidden_size) but no dedicated NVFP4 hidden-dim/layout rejection test yet.
• Deterministic parity checks pass (fixed seed 42) between:
  • native block-scaled execution path — N/A for NVFP4: no native block-scaled CUTLASS path yet (Blackwell-only); NVFP4 runs dequant-to-A16 fallback + fused GEMV decode only.
  • dequantization fallback path — asserted vs the dequantized reference with documented tolerances (atol 0.12 FP16 / 0.15 BF16).

✅ Must-pass dispatch and validation gates

• quant_type="nvfp4" is validated consistently at schema + runtime boundaries (contrib_defs.cc schema; moe_quantization.cc block-size / block-scale / global-scale checks).
• Native-kernel eligibility predicate is explicit and covered by tests for each disqualifier — is_moe_gemv_fp4_supported is explicit and GEMV on/off is exercised, but each disqualifier (SM bound, shape bounds, group_size) is not individually asserted.
• Unsupported NVFP4 combinations fail fast with actionable error messages (reject tests + mode-aware validation messages).
• Fallback behavior is explicitly tested where native block-scaled kernels are unavailable (all base NVFP4 tests + the gemv-disabled test exercise the dequant fallback).

✅ Must-pass performance gates

• Benchmark run completed for representative decode and prefill shapes (PR Testing: bench_qmoe_fp4_compare.py --iters 400 --warmup 40 --with-fallback).
• NVFP4 path shows expected perf/memory benefit vs baseline — benchmarked; awaiting reviewer sign-off on the evidence.
• No unacceptable perf regression on non-NVFP4 QMoE paths — awaiting reviewer sign-off.
• No major perf cliff when routing between GEMV/GEMM/native/fallback for nearby shapes — awaiting reviewer sign-off.

✅ Must-pass numerical stability gates

• Max/mean error metrics vs reference are within agreed thresholds for all tested shapes (per-dtype atol asserted in the parity tests).
• No outlier instability on extreme but valid scale ranges — not specifically exercised by a dedicated test.
• Row-varying scale handling remains correct for existing FP4 coverage after this change (test_fp4_native_cutlass_row_varying_scales).

✅ Code quality / maintainability gates

• Validation logic for NVFP4 invariants is centralized / clearly non-duplicative (is_nvfp4 / is_fp4_family predicates in moe_quantization.cc).
• Scale semantics (block scale vs global scale order-of-application) are clearly documented in code/comments/docs (moe_qmoe.md §9b).
• New docs accurately describe tensor layout, quantization params, and execution behavior (moe_qmoe.md, ContribOperators.md).
• CI is green and no unresolved blocking review comments remain — review comments are all resolved (0 open), but CI is currently pending on the latest push.

✅ Sign-off criteria

• Numerics sign-off: parity and tolerance results reviewed and approved.
• Performance sign-off: benchmark evidence reviewed and approved.
• Runtime safety sign-off: malformed input/unsupported path behavior reviewed and approved.
• Final reviewer recommendation: Approve / Request changes recorded with rationale.