[CUDA] Fuse MoE router bias into MatMulNBits GEMV

[tianleiwu]

Description

Fuse the MoE router MatMulNBits + Add([32] bias) pattern into the CUDA MatMulNBits router GEMV path.

This PR keeps the public surface conservative:

• no QMoE op schema change;
• no new router/top-k QMoE inputs;
• the optimized path is exact-shape gated to the GPT-OSS router projection: M=1, N=32, K=2880, 4-bit weights, block_size=32, no zero points;
• all other MatMulNBits shapes continue to use the existing generic path;
• ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1 disables the exact router GEMV specialization;
• ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1 disables only the graph rewrite that folds the router bias into MatMulNBits.

Motivation and Context

GPT-OSS-20B decode runs a tiny router projection before each QMoE node. The router projection is an exact-shape int4 MatMulNBits, followed by a [32] bias add before QMoE consumes the router logits.

The existing generic int4 GEMV works, but this router shape is small enough that specializing it reduces router GEMV overhead. Once that specialization is active, folding the [32] bias into the same kernel removes the remaining router-side Add launch without changing the QMoE op contract.

Key Changes

• Adds an exact-shape CUDA router GEMV specialization in MatMulFloatInt4RouterKernel.
• Extends the CUDA MatMulNBits path to pass an optional bias pointer to the router specialization.
• Extends MatMulNBitsFusion to rewrite the exact GPT-OSS router MatMulNBits + Add chain into biased MatMulNBits.
• Keeps the transformer registration compatible with the current origin/main WebGPU kernel-gated MatMulNBits fusion logic.
• Adds graph transformer and MatMul4Bits coverage for the specialization, fallback, and bias-fusion opt-out behavior.
• Records the router GEMV and router bias fusion measurements in the QMoE GEMV experiment log.

Validation

Completed locally on the clean PR branch:

• lintrunner -a docs/contrib_ops/cuda/qmoe_gemv_experiments.md onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh onnxruntime/core/optimizer/graph_transformer_utils.cc onnxruntime/core/optimizer/matmul_nbits_fusion.cc onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc onnxruntime/test/contrib_ops/matmul_4bits_test.cc onnxruntime/test/optimizer/graph_transform_test.cc
• git diff --check
• git diff --cached --check

Previously collected on the experiment branch before preparing this PR branch:

• Graph transformer tests for router GEMV/bias fusion passed.
• MatMul4Bits provider coverage for router GEMV specialization/fallback passed.
• Nsight confirmed the exact router specialization dispatches for GPT-OSS decode router projections.
• CUDA-graph GPT-OSS decode A/B showed the router GEMV specialization at about +1.6% to +1.8% throughput.
• Router bias fusion removed all 24 real GPT-OSS router bias Add nodes and measured about +0.2% throughput after the router GEMV specialization.

Compiled C++ tests were not rerun from this new worktree because it does not have a configured build directory; CI should provide the full compiled validation matrix.

[Copilot]

Pull request overview

This PR adds an exact-shape CUDA fast path for the GPT-OSS router MatMulNBits GEMV (M=1, N=32, K=2880, int4, block_size=32, no zero-points) and extends it to optionally apply the [32] bias in-kernel, plus a graph rewrite to fuse the corresponding Add into MatMulNBits under tight shape/EP gating and env-var opt-outs.

Changes:

• Adds a CUDA router GEMV specialization (MatMulFloatInt4RouterKernel) and threads an optional bias_data pointer through the TryMatMulNBits/TryMatMul4Bits dispatch.
• Updates MatMulNBitsFusion + transformer registration to enable CUDA-side bias fusion only for the exact router projection shape and when not opted out via env vars.
• Adds/extends unit tests and experiment documentation to cover specialization dispatch, fallback, and bias-fusion opt-out behavior.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/optimizer/graph_transform_test.cc	Extends transformer tests to cover CUDA EP + exact router shape, and adds opt-out / unsupported-shape cases.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc	Adds CUDA provider coverage for the exact router shape (specialized vs fallback) and biased GEMV.
onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc	Updates a MatMulNBits CUDA dispatch call site for the new bias pointer parameter.
onnxruntime/core/optimizer/matmul_nbits_fusion.cc	Adds CUDA-only gating for bias fusion and env-var opt-outs for the exact GPT-OSS router shape.
onnxruntime/core/optimizer/graph_transformer_utils.cc	Enables MatMulNBitsFusion for CUDA EP by expanding compatible EPs.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh	Extends CUDA TryMatMul* helpers to accept an optional bias pointer; disallows bias for int8.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc	Wires bias input through to CUDA TryMatMulNBits and tightens the unsupported-bias error path.
onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu	Implements the exact-shape router GEMV kernel and bias-in-kernel support; dispatches under env-var control.
docs/contrib_ops/cuda/qmoe_gemv_experiments.md	Records benchmarking/validation details for the specialization and bias fusion.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 8 comments.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.