[CUDA] QMoE GEMV fast path for batch-1 decode

[tianleiwu]

Summary

Adds a symmetric weight-only MoE GEMV fast path for single-token (batch-1) decode of quantized MoE models such as GPT-OSS-20B, replacing the CUTLASS grouped-GEMM path when the expanded row count is small. The fast path now covers INT4 and INT8 weights, per-column and block-wise (group size 32/64/128) scales, and FP16 and BF16 activations. On an H200 (SM90) this improves end-to-end GPT-OSS-20B INT4 token-generation throughput by roughly 15% over the CUTLASS grouped-GEMM baseline and about 8% over the FasterTransformer kernel used in ORT 1.26 across prompt lengths 128/1024/2048, while producing bit-faithful output.

Motivation

At batch-1 decode each token expands to top_k rows (4 for GPT-OSS-20B), so the MoE FC1/FC2 GEMMs are extremely skinny. The CUTLASS grouped GEMM is built for throughput at larger M and leaves the decode path memory-bound and underutilized. A dedicated weight-only INT GEMV with per-expert dispatch is a better fit for this regime and closes the gap to (and surpasses) ORT 1.26 for GPT-OSS-20B.

This also adds block_size=32 support as requested in #29035.

Key Changes

New MoE GEMV kernel

File	Change
contrib_ops/cuda/llm/moe_gemm/moe_gemv.h	Public interface: is_moe_gemv_supported dispatch predicate and the symmetric INT launchers launch_moe_gemv_int_symmetric<T, WeightType> / launch_moe_gemv_int_symmetric_interleaved_swiglu<T, WeightType> (plus the original INT4 per-channel launchers).
contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu	Symmetric INT MoE GEMV (INT4 uint4b_t / INT8 uint8_t, FP16/BF16). One CTA per expanded row × N-tile; per-expert weight/scale/bias offsets via a direct row-to-expert map (prefix-offset scan fallback). Supports per-column (group_size <= 0) and block-wise (group_size 64/128) scales. FC1 has an interleaved SwiGLU-fused epilogue; a static top-k one-row finalize specialization is used for FC2 routing.

Runner wiring and profiler

File	Change
contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu / .h	Branch FC1 and FC2 to the GEMV fast path when supported (carrying group_size through QuantParams); fall back to grouped GEMM otherwise. ORT_DISABLE_MOE_GEMV=1 forces the grouped-GEMM path for A/B testing.
contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc / .h, moe_util_kernels.h	Use M (token count) in the MoE GEMM profiler so tactic selection reflects the decode shape.

Backward-compatible SwiGLU fusion

File	Change
contrib_ops/cuda/moe/moe.cc, moe_quantization.cc (+ .h)	Treat swiglu_fusion == 0 with no separate FC3 as interleaved fusion (== 1). The published GPT-OSS-20B model (and any model exported by ORT < 1.27) hard-coded the interleaved layout and emits no swiglu_fusion attribute, so it defaults to 0; those weights are pre-fused into FC1 and must be treated as fusion mode 1.

Tests, profiling, and docs

File	Change
test/python/transformers/profile_qmoe_gemv.py / .sh	Standalone QMoE GEMV profiling harness (NVTX-ranged, GEMV-vs-grouped-GEMM kernel comparison).
test/python/transformers/test_qmoe_cuda.py, test_moe_cuda.py	QMoE GEMV decode-latency/parity coverage across INT4/INT8, per-column/block-wise, FP16/BF16, plus import tidy-ups.
docs/contrib_ops/cuda/qmoe_gemv_experiments.md	Full experiment log: kernel-level sweeps, dispatch-gate tuning, block-wise and BF16 enablement, INT8 per-column root-cause + fix, and the GenAI end-to-end throughput comparison (CUTLASS baseline / GEMV final / FT baseline).
docs/contrib_ops/cuda/moe_qmoe.md	QMoE doc updates describing the GEMV fast path and its dispatch gate.

Results

GPT-OSS-20B INT4 end-to-end (H200/SM90, batch 1)

Token-generation throughput (tps, higher is better):

Prompt length	CUTLASS baseline (gemm)	GEMV (final)	FT (ORT 1.26)	GEMV vs cutlass	GEMV vs FT
128	248.9	288.0	265.2	+15.7%	+8.6%
1024	237.8	272.2	252.9	+14.5%	+7.6%
2048	231.3	265.0	245.6	+14.6%	+7.9%

Decode microbenchmark coverage (H200/SM90, batch 1)

Benchmark-loop latency in milliseconds, lower is better. Enabled is the default GEMV build; Fallback sets ORT_DISABLE_MOE_GEMV=1 (grouped GEMM). Every case reported has_invalid_output=false.

Case	Quant	DType	Enabled ms	Fallback ms	Speedup
int8_per_column_m1_top2_1024x4096_e8	INT8 per-column	FP16	0.0566	0.0816	1.44x
int8_per_column_m1_top2_1024x4096_e8	INT8 per-column	BF16	0.0578	0.0862	1.49x
gpt_oss_20b_m1_top4_int8_2880x2880_e32	INT8 per-column	FP16	0.0785	0.0947	1.21x
gpt_oss_20b_m1_top4_int8_2880x2880_e32	INT8 per-column	BF16	0.0785	0.0989	1.26x

Block-wise INT4 (block_size=64, 1024x4096, e8) routes to GEMV for both FP16 and BF16 with FC1/FC2 kernel times within noise across dtypes (FP16 4.53/6.95 us, BF16 4.62/7.01 us). See qmoe_gemv_experiments.md for the full sweep.

Correctness

• GEMV-enabled and ORT_DISABLE_MOE_GEMV=1 (grouped GEMM) produce identical, correct output for the GPT-OSS-20B sanity prompt; Nsight traces confirm moe_gemv_kernel (FC2) and moe_gemv_interleaved_swiglu_kernel (FC1) run for the decode shapes.
• INT4/INT8 SwiGLU parity cases pass for FP16 and BF16 (max absolute difference ~1e-3 against the reference). Regression: pytest -k "TestSwigluQMoE or TestQMoEIntPrePackSmoke" → 34 passed, 4 skipped.
• Test on GPT-OSS-20b with block_size=32, and generated results looks good.

Testing Notes

• Build: bash .env/cuda_130.sh --build (CUDA 13.0, CMAKE_CUDA_ARCHITECTURES=89;90). Use --clean_moe after dispatch-code edits to avoid stale moe_kernels.cu.o.
• Kernel profiling: onnxruntime/test/python/transformers/profile_qmoe_gemv.sh.
• Parity/latency: pytest onnxruntime/test/python/transformers/test_qmoe_cuda.py.
• End-to-end: GenAI benchmark_e2e.py on GPT-OSS-20B INT4, batch 1, prompt lengths 128/1024/2048. Run on an idle GPU — shared-GPU contention corrupts the decode-latency measurement.
• A/B check: compare default vs ORT_DISABLE_MOE_GEMV=1 for both throughput and output parity.

[Copilot]

Pull request overview

This PR adds a CUDA INT4 per-channel MoE GEMV fast path intended for batch-1 decode (small expanded row counts), wiring it into the existing CUTLASS MoE runner with an environment-variable kill switch, plus profiling/tactic-selection improvements and test+doc updates.

Changes:

• Add an fpA_intB_gemv-based batched MoE GEMV kernel (including an interleaved SwiGLU-fused FC1 variant) and route FC1/FC2 to GEMV when supported, otherwise fall back to grouped GEMM (ORT_DISABLE_MOE_GEMV=1 forces fallback).
• Update MoE GEMM profiling/tactic caching to be M-bucket aware so decode vs prefill shapes can select different best tiles.
• Improve backward compatibility for legacy SwiGLU models (treat swiglu_fusion==0 + no separate FC3 weights as interleaved fusion) and add benchmark harness/docs.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_qmoe_cuda.py	Adds opt-in benchmark plumbing (NVTX ranges/env knobs), decode-sized parity coverage, and a helper benchmark runner.
onnxruntime/test/python/transformers/test_moe_cuda.py	Updates INT weight packing comments and tweaks test-case commentary to reflect GEMV routing.
onnxruntime/test/python/transformers/profile_qmoe_gemv.sh	New Nsight Systems profiling wrapper to compare GEMV vs grouped-GEMM fallback.
onnxruntime/test/python/transformers/profile_qmoe_gemv.py	New Python entrypoint to run the opt-in benchmark cases and emit structured results.
onnxruntime/contrib_ops/cuda/moe/moe.h	Removes per-instance cached GemmId members (now handled via profiler cache).
onnxruntime/contrib_ops/cuda/moe/moe.cc	Adds legacy SwiGLU fusion compatibility; updates profiler API usage and per-forward profiling.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Removes per-instance cached GemmId members (now handled via profiler cache).
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Adds legacy SwiGLU fusion compatibility; updates profiler API usage and per-forward profiling.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_util_kernels.h	Extends fused prologue API to materialize a permuted-row→expert map.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h	Plumbs the new permuted-row→expert map through GEMM entrypoints (for GEMV fast path).
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu	Implements GEMV dispatch + env kill switch, writes the row→expert map in prologue, adds one-row finalize specialization, and wires GEMV into FC1/FC2.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h	New public GEMV support predicate + launchers (incl. interleaved SwiGLU fused FC1).
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu	New batched MoE GEMV CUDA kernels and launch logic based on fpA_intB_gemv building blocks.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.h	Updates profiler API and documents M-bucket-based caching.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemm_profiler.cc	Implements M bucketing, per-bucket caching, and nearest-bucket fallback.
docs/contrib_ops/cuda/qmoe_gemv_experiments.md	Adds a detailed experiment/profiling log to support future tuning/regression checks.
docs/contrib_ops/cuda/moe_qmoe.md	Documents the GEMV path, dispatch gates, and improvement plan; adds profiling workflow notes.

[Copilot]

Pull request overview

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

[github-actions]

You can commit the suggested changes from lintrunner.