[Tianlei Wu (tianleiwu)]

Description

Optimizes the offline K-packed (weights_prepacked=1) NVFP4 QMoE decode path on SM90+ for 256-expert models with hidden_size = 3072 and 8-token speculative verification (64 expanded rows at top_k = 8, such as Ornith-1.5-122B):

1. Extend default MoeGemvFp4DecodeKernel envelope (ORT_FP4_GEMV_DECODE_KERNEL=1):
   • Raises kDecodeDefaultMaxN = 3072 (was 2560), kDecodeDefaultMaxK = 3072 (was 2560), and kDecodeDefaultMaxExpandedRows = 64 (was 20).
   • Enables the table-free register-decoded MoeGemvFp4DecodeKernel (with shared-memory FP32 activation staging and PDL overlap of FC1 -> FC2 -> finalize) by default for hidden_size = 3072, inter_size = 1024 (FC1 N=2048, K=3072; FC2 N=3072, K=1024) up to 8-token speculative verification (64 expanded rows), instead of falling back to MoeGemvFp4RawKPackedKernel.
2. Add 256-expert single-CTA fused decode routing (FusedSoftmaxTopKMaps256Kernel):
   • Extends LaunchFusedRouteMaps (ORT_FP4_GEMV_FUSED_ROUTE=1) to support num_experts == 256, 1..8 rows, and k in {1, 2, 4, 6, 8} (num_rows * k <= 64).
   • Uses 8 warps (256 threads, one warp per row) and reproduces SoftmaxTopKWarpTopKKernel<T, 8, 4>'s WarpReduceMax / WarpReduceSum / WarpBitonicTopN<8> evaluation order operation-for-operation while emitting (expert, token)-sorted expert maps in the same CTA with programmatic dependent launch (PDL).
   • Produces bit-identical outputs to LaunchSoftmaxTopK + fusedBuildExpertMapsSortFirstToken.

Summary of Changes

File	Change
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu	Raise kDecodeDefaultMaxN = 3072, kDecodeDefaultMaxK = 3072, kDecodeDefaultMaxExpandedRows = 64
onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuh	Add FusedSoftmaxTopKMaps256Kernel<T, K> (kFusedRoute256Experts = 256, kFusedRoute256MaxRows = 8, kFusedRoute256MaxK = 8, kFusedRoute256Threads = 256)
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h / .cu	Dispatch num_experts == 256 (1..8 rows, k in {1, 2, 4, 6, 8}) in LaunchFusedRouteMaps
docs/contrib_ops/cuda/moe_qmoe.md	Update §9.9 (ORT_FP4_GEMV_DECODE_KERNEL, ORT_FP4_GEMV_FUSED_ROUTE) and §9b.2 documentation
onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py	Add N=2048/3072, K=3072/1024, expanded=64 decode GEMV dispatch/bit-exactness cases and 256-expert 8-row fused routing + PDL tests

Testing & Performance

Validated on NVIDIA H200 (sm_90a, CUDA 13.0):

• pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py -k "test_nvfp4_offline_row_major_decode_gemv": 4/4 passed (FP16 and BF16 bit-exactness across ORT_FP4_GEMV_DECODE_KERNEL={0, 1, 2} and ORT_FP4_GEMV_FUSED_ROUTE={0, 1} with and without CUDA Graph capture).
• lintrunner: no issues.
• End-to-end 8-token DFlash2 speculative verification on Ornith-1.5-122B (48 NVFP4 QMoE layers, num_experts=256, top_k=8, hidden_size=3072, inter_size=1024, 64 expanded rows/step):

48-Layer QMoE Stage (ms/step at 2K Context)	Before (MoeGemvFp4RawKPackedKernel + 2-kernel routing)	With this PR (MoeGemvFp4DecodeKernel + FusedSoftmaxTopKMaps256Kernel)	Improvement
FC1 + FC2 NVFP4 GEMV (96 launches)	11.33 ms	6.10 ms	-5.23 ms (-46.2%, 1.86x faster)
MoE Routing + Expert Maps (96 -> 48 launches)	0.67 ms	0.24 ms	-0.43 ms (-64.2%, 2.79x faster)
Total 48-layer QMoE decode step	12.00 ms	6.34 ms	-5.66 ms/step

[Copilot]

🔵 Needs a closer look

The default-enabled CUDA routing and dependent-launch behavior need broader regression coverage and final human review of GPU correctness.

1 open issue

• Expand 256-expert tests across supported k values and row counts · New

What changed in this PR

This PR expands ONNX Runtime’s NVFP4 QMoE decode fast path for larger projections and adds fused routing for 256-expert models.

Changes:

• Raises default decode GEMV limits to 64 expanded rows and N/K dimensions of 3072.
• Adds single-block fused routing for 256 experts and up to eight tokens.
• Updates dispatch documentation and adds bitwise-parity, CUDA Graph, and programmatic dependent launch tests.

File	Description
onnxruntime/​test/​python/​transformers/​test_qmoe_nvfp4_cuda.py	Adds larger decode shapes and 256-expert routing tests.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Documents the expanded fused-routing limits.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Dispatches eligible workloads to the 256-expert kernel.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_fused_route.cuh	Implements fused 256-expert routing and expert-map generation.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_fp4.cu	Expands default decode GEMV limits.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Updates documented decode and routing limits.

_{Review effort: Balanced}

[Baiju Meswani (baijumeswani)]

This PR expands the default NVFP4 decode kernel range and combines routing and expert-map creation into one launch for 256-expert models. The change is focused, keeps the existing fallback and opt-out paths, and adds useful tests for the new shapes, top-k values and CUDA Graphs.[/+]

My independent checks passed 1,600 FP16/BF16 routing cases, with matching scales, selected experts and offsets, and valid forward/inverse maps. Memory, race and synchronization checks were clean. Routing compilation passed SM80/86/90 with CUDA 13.0 on Linux. These checks do not cover Hopper dependent launches or full-model parity. The reported H200 performance numbers are consistent, but I could not independently reproduce them on the available A100 hardware; they describe the QMoE stage, not the whole model.

I would hold approval until the Windows CUDA 12.8 build regression is fixed. The matching builds pass on the exact base commit and fail on this PR head. The blocking issue and one small comment correction are noted inline.