You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[CUDA] Add fused decode routing, single-pass finalize, and PDL overlap for NVFP4 QMoE - #33366
Stacked on #33354 (and #33094). The base branch is tlwu/20261013/nvfp4-qmoe-decode-gemv, so the diff below is this change only. Once #33094 and #33354 merge, this PR will be retargeted to main. Opened as a draft until then.
Description
Optimizes the non-GEMV kernels and inter-kernel launch serialization along the offline K-packed (weights_prepacked=1) NVFP4 QMoE decode path on SM90+:
Single-CTA fused decode routing (FusedSoftmaxTopKMapsKernel): Replaces LaunchSoftmaxTopK + fusedBuildExpertMapsSortFirstToken in a single CTA launch for num_rows <= 4, num_experts == 512, and k in {1, 2, 4, 6, 8, 10}. One warp per row reproduces the SoftmaxTopKMergeKernel<T, 128, 4> reduction trees operation for operation and emits the (expert, token)-sorted expert maps bit-identically (ORT_FP4_GEMV_FUSED_ROUTE=0 disables).
Single-pass finalizeMoeRoutingKernel: Sizes the block to one thread per 128-bit column vector (up to 1024 threads instead of a fixed 256) so hidden dimensions up to 4096 (such as Qwen3.8's hidden=2560) reduce in a single pass, and issues batched loads of up to 16 expert rows together while preserving top-k accumulation order (7.7 us -> 4.0 us per layer on H200).
Programmatic dependent launch (PDL) overlap of FC1 -> FC2 -> finalize: Launches FC2 and finalizeMoeRoutingKernel with cudaLaunchAttributeProgrammaticStreamSerialization (ORT_FP4_GEMV_PDL=0 disables; ORT_FP4_GEMV_FC1_PDL=1 optionally enables PDL for FC1 behind fused routing). FC1 triggers dependents at start, FC2 prefetches its first batch of weights/scales before waiting on FC1's output, and finalizeMoeRoutingKernel stages routing metadata into shared memory before waiting on FC2's output (~2.2 us saved per layer on H200, preserved under CUDA Graph capture).
All three optimizations preserve exact floating-point evaluation order and produce bit-identical outputs.
griddepcontrol.wait / griddepcontrol.launch_dependents in MoeGemvFp4DecodeKernel and programmatic_dependent_launch attribute in LaunchMoeGemvFp4RawNPacked (ORT_FP4_GEMV_PDL)
Bit-exactness: engine_gate.py --prefill steps passes 6/6 with 100% identical token sequences (256/256 tokens on text2048 and prose2048) compared against the baseline without fused routing / PDL.
End-to-end decode speed (bench_engine_mtp.py, 256 generated tokens, median of 3 runs):
…p for NVFP4 QMoE
For the offline K-packed (weights_prepacked=1) NVFP4 QMoE decode path on SM90+:
- Fused softmax + top-k + expert-map kernel (FusedSoftmaxTopKMapsKernel) for E=512, num_rows <= 4, k in {1, 2, 4, 6, 8, 10}, reproducing LaunchSoftmaxTopK + fusedBuildExpertMapsSortFirstToken bit-exactly in a single CTA launch (ORT_FP4_GEMV_FUSED_ROUTE=0 disables).
- Single-pass finalizeMoeRoutingKernel with one thread per 128-bit column vector (up to 1024 threads) and batched loads of up to 16 expert rows while preserving top-k accumulation order.
- Programmatic dependent launch (PDL) overlap of FC1 -> FC2 -> finalize so FC2 prefetches weights/scales and finalize loads routing metadata before waiting on the producer grid (ORT_FP4_GEMV_PDL=0 disables; ORT_FP4_GEMV_FC1_PDL=1 optionally enables PDL for FC1 behind fused routing).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Guard FusedSoftmaxTopKMapsKernel inline PTX instructions only for __CUDA_ARCH__ >= 900 (with a warp-shuffle fallback for __reduce_max_sync below SM80) and verify compiled SM90 PTX/SASS via cudaFuncGetAttributes in LaunchFusedRouteMapsT.
- Gate fc2_finalize_pdl on sm_ >= 90 (and disable PDL inside the autotune timing loop), and verify SM90 PTX/SASS in finalizeMoeRoutingKernelLauncher before enabling programmaticStreamSerializationAllowed.
- Clamp finalizeMoeRoutingKernel blockDim to the compiled specialization's register-limited maxThreadsPerBlock and check cudaLaunchKernelEx return status.
- Extend test_nvfp4_offline_row_major_decode_gemv_bit_exact to cover 512-expert fused routing, PDL opt-out, and wide hidden_size=5120 finalize.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Explicitly guard griddepcontrol PTX instructions inside MoeGemvFp4DecodeKernel with __CUDA_ARCH__ >= 900 and add memory clobber to finalizeMoeRoutingKernel griddepcontrol.wait.
- Check SM90+ device capability, CUDA driver >= 12.0, and ptxVersion >= 90 && binaryVersion >= 90 in QuerySm90DecodeKernelCaps, IsFusedRouteKernelSm90Available, and getFinalizeKernelLaunchCaps while clearing sticky CUDA errors on query failure.
- Avoid attaching cudaLaunchAttributeProgrammaticStreamSerialization when PDL is disabled, fall back to a regular launch if cudaLaunchKernelEx rejects the attribute, and gate FC2/finalize PDL on fc1_used_decode/fc2_used_decode.
- Move Fp4MoeGemvUsePdl outside the standalone NVFP4 decode device-code slice and extend test_benchmark_qwen38_mtp_kernels.py.
- Fix FP16 native-scale FC1 dispatch expectation for hidden_size <= 1024 and add 512-expert FC1 PDL + CUDA Graph test coverage in test_qmoe_nvfp4_cuda.py.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…uild dir
Resolve onnxruntime source root across parent directories, CI workspace environment variables (GITHUB_WORKSPACE/BUILD_SOURCESDIRECTORY), and cwd ancestors when benchmark_qwen38_mtp_kernels.py is copied into the CMake build tree.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Please add direct fused-versus-unfused routing parity tests. The NVFP4 model builder hardcodes normalize_routing_weights=1 (test_qmoe_nvfp4_cuda.py:221), and the standalone routing tests call only LaunchSoftmaxTopK. Thus normalize_scales=false is uncovered, and the end-to-end FP16/BF16 comparisons can hide differences in FP32 routing scales through final rounding. Extend softmax_topk_kernel_test.cc to compare FP32 scales bitwise, selected indices, both row maps, expert IDs, and expert offsets against LaunchSoftmaxTopK plus the existing map builder. Cover both normalization modes, FP16/BF16 inputs, supported row/k combinations (including four rows), and tied logits.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Optimizes the non-GEMV kernels and inter-kernel launch serialization along the offline K-packed (
weights_prepacked=1) NVFP4QMoEdecode path on SM90+:FusedSoftmaxTopKMapsKernel): ReplacesLaunchSoftmaxTopK+fusedBuildExpertMapsSortFirstTokenin a single CTA launch fornum_rows <= 4,num_experts == 512, andk in {1, 2, 4, 6, 8, 10}. One warp per row reproduces theSoftmaxTopKMergeKernel<T, 128, 4>reduction trees operation for operation and emits the(expert, token)-sorted expert maps bit-identically (ORT_FP4_GEMV_FUSED_ROUTE=0disables).finalizeMoeRoutingKernel: Sizes the block to one thread per 128-bit column vector (up to 1024 threads instead of a fixed 256) so hidden dimensions up to 4096 (such as Qwen3.8'shidden=2560) reduce in a single pass, and issues batched loads of up to 16 expert rows together while preserving top-k accumulation order (7.7 us -> 4.0 usper layer on H200).PDL) overlap ofFC1 -> FC2 -> finalize: Launches FC2 andfinalizeMoeRoutingKernelwithcudaLaunchAttributeProgrammaticStreamSerialization(ORT_FP4_GEMV_PDL=0disables;ORT_FP4_GEMV_FC1_PDL=1optionally enables PDL for FC1 behind fused routing). FC1 triggers dependents at start, FC2 prefetches its first batch of weights/scales before waiting on FC1's output, andfinalizeMoeRoutingKernelstages routing metadata into shared memory before waiting on FC2's output (~2.2 ussaved per layer on H200, preserved under CUDA Graph capture).All three optimizations preserve exact floating-point evaluation order and produce bit-identical outputs.
Summary of Changes
onnxruntime/contrib_ops/cuda/moe/qmoe_fused_route.cuhFusedSoftmaxTopKMapsKernelsingle-CTA fused softmax + top-k + expert-map kernel (E=512,rows <= 4,k in {1,2,4,6,8,10})onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h/.cuLaunchFusedRouteMapsdeclarations and launcheronnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu/moe_util_kernels.hfinalizeMoeRoutingKernel(blockDim.xup to 1024,FINALIZE_LOAD_BATCH=16) andearly_routing_loadsPDL supportonnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h/.cugriddepcontrol.wait/griddepcontrol.launch_dependentsinMoeGemvFp4DecodeKernelandprogrammatic_dependent_launchattribute inLaunchMoeGemvFp4RawNPacked(ORT_FP4_GEMV_PDL)onnxruntime/contrib_ops/cuda/moe/moe_quantization.ccLaunchFusedRouteMapsandfc2_finalize_pdlintoQMoE::ComputeInternaldocs/contrib_ops/cuda/moe_qmoe.mdORT_FP4_GEMV_FUSED_ROUTE,ORT_FP4_GEMV_PDL, andORT_FP4_GEMV_FC1_PDLTesting
Validated on NVIDIA H200 (
sm_90a, CUDA 13.0):python -m pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py: 178 passed, 8 skipped, 0 failed (includingtest_nvfp4_offline_row_major_decode_gemv_bit_exactfor FP16 and BF16).v12b3-session CUDA-graph + MTP engine, 48 MoE layers + 1 MTP layer):engine_gate.py --prefill stepspasses 6/6 with 100% identical token sequences (256/256tokens ontext2048andprose2048) compared against the baseline without fused routing / PDL.bench_engine_mtp.py, 256 generated tokens, median of 3 runs):#33354only)ms/cyclems/cycletok/s)lintrunneron all touched files: no issues.