[CUDA] Add decode GEMV kernel for K-packed NVFP4 QMoE

[Tianlei Wu (tianleiwu)]

Stacked on #33094. The base branch is kvaishnavi/qwen38-nvfp4-offline-packing, so the diff below is this change only. Once #33094 merges, this PR will be retargeted to main. Opened as a draft until then.

Description

Adds MoeGemvFp4DecodeKernel, a faster decode kernel for the offline K-packed (weights_prepacked=1, row-major [E, N, K/2]) NVFP4 QMoE GEMV that #33094 introduced. It reproduces MoeGemvFp4RawKPackedKernel's fp32 summation order bit-exactly, so it is a drop-in replacement: QMoE outputs are bitwise identical with or without it. On H200 it makes the Qwen3.8 FC1+FC2 decode launches 1.47x faster (BF16, 10 expanded rows) and the QMoE node 11% (1 token) to 17% (2 tokens) faster end to end.

Summary of Changes

Kernel

File	Change
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu	New MoeGemvFp4DecodeKernel and its launcher dispatch inside LaunchMoeGemvFp4RawNPacked (weights_row_major branch); the launcher now receives sm
docs/contrib_ops/cuda/moe_qmoe.md	New "Decode kernel for offline K-packed storage" section and ORT_FP4_GEMV_DECODE_KERNEL table row
onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py	New bit-exactness test (see Testing)

How it matches MoeGemvFp4RawKPackedKernel:

• Same K-lane count (KLanes = 4 for K <= 1024, else 8, copied from [CUDA] Support offline NVFP4 QMoE weights and fused block-FP8 execution #33094's launcher), same lane assignment (lane = block(16) % KLanes), same separate even/odd fp32 accumulators per lane, same acc_even + acc_odd followed by the xor tree KLanes/2 .. 1, same StoreMoeGemvFp4Pair epilogue (bias, fused SwiGLU).
• One thread owns two adjacent columns and the lane pair (2j, 2j+1): one 128-bit load covers two blocks, the xor-tree steps at lane distance KLanes/2 .. 2 are shuffles between the KLanes/2 threads of a column, and the last step (distance 1) is an in-thread add.
• Speed comes from 128-bit streaming weight loads with 2 K tiles in flight, table-free E2M1 decode (fp4 * 2^-126 bf16 / fp4 * 2^-14 fp16 bit patterns times a rescaled block scale, a single rounding identical to T(fp4 * scale)), E4M3 scale conversion with cvt.rn.f16x2.e4m3x2, and the activation row staged once per CTA as fp32 in shared memory.

Dispatch (everything else falls back to MoeGemvFp4RawKPackedKernel unchanged):

• weights_row_major, sm >= 90, FP16 or BF16, at most 80 expanded rows, K % 32 == 0, K <= 8192, even N, 16/2/16-byte aligned weights/scales/activations.
• Steps aside when ORT_NVFP4_FC1_EXPERT_REUSE=1 selects the two-row reuse variant.
• Opt-out: ORT_FP4_GEMV_DECODE_KERNEL=0 (latched once per process, same convention as ORT_FP4_GEMV_DEFAULT_TILING).

Testing

Standalone harness (links moe_gemv_fp4.cu directly, calls the public launchers vs. MoeGemvFp4RawKPackedKernel launched exactly like #33094's launcher), H200, CUDA 13.0:

• Bit-exact: 374/374 cases, bf16 and fp16, 0 mismatching elements. Real Qwen3.8 FC1/FC2 weights, random weights, bias on/off, fused SwiGLU on/off, invalid expert ids (-1, E), offsets-only expert map, rows 1..80 (every row count), K = 64 .. 8192 including K not a multiple of 1024 (partial last tile: 98 and 34 blocks), big/tiny/zero global scales (generic decode path), and shapes the kernel declines (fallback).

Kernel timing (us/launch, CUDA-graph average over 300 launches, 512-expert L2-cold layout, real FC1 N=1280 K=2560 fused SwiGLU and FC2 N=2560 K=640 weights, H200):

dtype / rows	#33094 FC1	#33094 FC2	this PR FC1	this PR FC2	FC1+FC2 speedup
bf16, top-10	18.51	10.27	12.13	7.49	1.47x
bf16, top-20	36.05	17.23	20.79	12.20	1.62x
fp16, top-10	17.55	9.44	13.77	8.41	1.22x
fp16, top-20	34.01	16.80	24.71	14.28	1.30x

ORT-level QMoE node (Qwen3.8 shape hidden=2560 inter=640 top_k=10, 512 experts, BF16, weights_prepacked=1, routing rotated over 64 inputs; GEMV kernel time per node from CUPTI, wall time per run_with_iobinding):

tokens	GEMV us/node #33094	GEMV us/node this PR	node wall us #33094	node wall us this PR
1	28.97	19.44	88.5	78.5
2	52.18	33.24	113.3	94.5

The node output is bitwise identical with ORT_FP4_GEMV_DECODE_KERNEL=0 and =1 in this benchmark too.

Python tests, CUDA plugin build of this branch (sm_90 only, H200, CUDA_VISIBLE_DEVICES=2, ORT_TEST_CUDA_PLUGIN_EP=1):

• New: test_nvfp4_offline_row_major_decode_gemv_bit_exact (FP16 and BF16). Runs five weights_prepacked=1 QMoE cases (Qwen3.8 shape at 1/2/6 tokens, i.e. 10/20/60 expanded rows, with and without bias; a partial-K-tile shape; a K % 32 != 0 shape that the decode kernel declines), each checked against the PyTorch reference, and requires the output to be bitwise equal to a fresh process with ORT_FP4_GEMV_DECODE_KERNEL=0. It also uses torch.profiler to assert that the decode kernel runs in the default process and MoeGemvFp4RawKPackedKernel runs with the opt-out.
• Full python -m pytest onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py: 178 passed, 8 skipped (opt-in large-memory tests), 0 failed (including all offline_row_major and kpacked tests from [CUDA] Support offline NVFP4 QMoE weights and fused block-FP8 execution #33094).
• lintrunner (CLANGFORMAT, RUFF, RUFF-FORMAT) on the touched files: no issues.

Known limitations:

• The table-free decode needs every block scale below 4 (|global_scale| * 448 < 3.98; Qwen3.8 has at most 0.09). A per-expert fallback to a generic Fp4I2FConverter decode keeps results bit-identical for larger scales, but it is about 30% slower than MoeGemvFp4RawKPackedKernel (BF16 top-10: 24.0/13.3 us vs 18.6/10.3 us). Such checkpoints should set ORT_FP4_GEMV_DECODE_KERNEL=0.
• FP16 gains less than BF16 (1.22x vs 1.47x at top-10).

Motivation and Context

#33094 added the K-packed row-major layout and its native MoeGemvFp4RawKPackedKernel. At Qwen3.8 decode shapes that kernel issues 8-byte weight loads and re-reads/decodes the activation row per 16-element block. This PR keeps its numerics and replaces only the memory/issue pattern.

Related: #33323 / #33324 touch the same launcher (opt-in native E4M3 scale decoding) and are based on an earlier #33094 head, so they need a rebase onto #33094 regardless of this PR. The two are semantically compatible: the native-scale option only changes how E4M3 scales are converted (exactly), and the decode kernel already uses the hardware E4M3 conversion. The only overlap is a textual conflict next to the launcher dispatch (#33323 adds a parameter to the launch_k_packed lambda right below the new decode dispatch).

Checklist

• Tests added/updated
• No breaking changes (default behavior is bitwise identical; ORT_FP4_GEMV_DECODE_KERNEL=0 restores the previous kernel)
• Documentation updated (docs/contrib_ops/cuda/moe_qmoe.md)

[Copilot]

🟡 Changes recommended

Dispatch assertions reject supported PTX builds, and the generic-scale decoder needs committed regression coverage.

2 open findings

• Add committed FP16/BF16 parity regression tests for decoder branches · New
• Gate decode dispatch assertion on compiled PTX image availability · New

What changed in this PR

Adds optimized CUDA decode for offline K-packed NVFP4 QMoE, building on #33094, and preserves FP32 intermediates for ungated block-FP8 SiLU.

Changes:

• Adds decode GEMV dispatch with fallback and an opt-out.
• Retains FP32 FC1 accumulators through ungated FP8 bias and activation.
• Adds parity and precision tests, plus documentation.

File	Description
onnxruntime/​test/​python/​transformers/​test_qmoe_nvfp4_cuda.py	Adds decode parity and dispatch tests.
onnxruntime/​test/​python/​transformers/​test_qmoe_fp8_cuda.py	Updates reference precision and adds regression tests.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.h	Supports separate projection and activation types.
onnxruntime/​contrib_ops/​cuda/​moe/​qmoe_kernels.cu	Implements FP32 projection intermediates.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_quantization.cc	Allocates and routes FP32 SiLU projections.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemv_fp4.cu	Adds optimized decode kernel and guarded dispatch.
docs/​contrib_ops/​cuda/​moe_qmoe.md	Documents decode eligibility and FP8 precision.

🧠 Review effort: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.