From 2a2ce830127b6a21c4eb8e31643117ed48485728 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 19 Jun 2026 17:21:35 -0700 Subject: [PATCH 01/10] [CUDA] Fuse GPT-OSS router bias into MatMulNBits GEMV --- .../contrib_ops/cuda/qmoe_gemv_experiments.md | 203 ++++++++++++++++++ .../cuda/quantization/matmul_4bits.cu | 100 +++++++++ .../cuda/quantization/matmul_nbits.cc | 11 +- .../cuda/quantization/matmul_nbits.cuh | 7 +- .../core/optimizer/graph_transformer_utils.cc | 2 +- .../core/optimizer/matmul_nbits_fusion.cc | 47 ++++ .../cuda_kernels/fpA_intB_gemm_kernel_test.cc | 1 + .../test/contrib_ops/matmul_4bits_test.cc | 36 ++++ .../test/optimizer/graph_transform_test.cc | 152 ++++++++++++- 9 files changed, 541 insertions(+), 18 deletions(-) diff --git a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md index d2165640b0051..7eb9be55b679e 100644 --- a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md +++ b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md @@ -1031,3 +1031,206 @@ fp32 default and the fp16-accumulation experiment: - Make fp16 accumulation the default for fp16 QMoE GEMV. - Keep bf16 on fp32 accumulation. - Keep `ORT_MOE_GEMV_FP32_ACCUM=1` as the opt-in numerical fallback and A/B knob. + +## 2026-06-19: Exact-Shape GPT-OSS Router GEMV Specialization + +### Change Under Test + +- Scope: standalone `MatMulNBits` router projection used before each GPT-OSS + `QMoE` node. +- Exact shape: `M=1`, `N=32`, `K=2880`, `block_size=32`, 4-bit weights, no + zero points. +- Added CUDA provider kernel: `MatMulFloatInt4RouterKernel`. +- Dispatch: `TryMatMul4Bits` selects the specialized kernel only for the exact + shape above. For clean A/B measurement, set + `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1` to force the generic + `MatMulFloatInt4Kernel` route. +- Motivation: avoiding 32 router logits in global memory is small. The more + direct win is reducing overhead in the special `N=32`, `K=2880` int4 router + GEMV itself. + +### Build and Sync + +```bash +cd ~/onnxruntime +git diff --check +cmake --build build/cu130/Release --target onnxruntime_providers_cuda --parallel $(nproc) +cp build/cu130/Release/libonnxruntime_providers_cuda.so \ + build/cu130/Release/onnxruntime/capi/libonnxruntime_providers_cuda.so +cp build/cu130/Release/libonnxruntime_providers_cuda.so \ + .venv_cu130/lib/python3.14/site-packages/onnxruntime/capi/libonnxruntime_providers_cuda.so +``` + +### Nsight Route Check + +Profile command: + +```bash +cd ~/onnxruntime +PYTHONPATH=~/onnxruntime/build/cu130/Release \ +MODEL=models/gpt-oss-20b/variants/cuda_int4_int4_qmoe_rtn_matmul_only \ +GPU=0 PROMPT_LEN=512 GEN_LEN=8 REPS=1 WARMUP=0 CUDA_GRAPH=0 XQA=1 SYNC_LIB=0 \ +ORT_FORCE_DETERMINISTIC_MOE=1 \ +~/cuda13.0/bin/nsys profile --trace=cuda,nvtx --sample=none --cpuctxsw=none \ + --force-overwrite=true --export=sqlite -o /tmp/gptoss_routergemv_special_cg0 \ + bash scripts/h200_18/bench_gpt_oss_ort_decode.sh +``` + +Nsight artifacts: + +- `/tmp/gptoss_routergemv_special_cg0.nsys-rep` +- `/tmp/gptoss_routergemv_special_cg0.sqlite` +- `/tmp/gptoss_routergemv_special_stats_cuda_gpu_kern_sum.csv` + +Kernel summary rows from the no-CUDA-graph profile: + +| Kernel | Total ns | Instances | Avg ns | +|--------|---------:|----------:|-------:| +| generic `MatMulFloatInt4Kernel<__half, 32, false>` | 3318180 | 343 | 9674.0 | +| specialized `MatMulFloatInt4RouterKernel<__half>` | 655863 | 168 | 3903.9 | +| `SoftmaxTopKWarpBitonicKernel<__half, 8>` | 439868 | 192 | 2291.0 | + +The `168` specialized router calls correspond to decode router projections in +the measured `GEN_LEN=8` no-CUDA-graph run. The generic int4 GEMV row remains for +other `MatMulNBits` projections. + +### Model-Level Decode Benchmark With CUDA Graph + +Both comparisons used the same rebuilt provider, CUDA graph enabled, XQA enabled, +and deterministic MoE tactic selection: + +```bash +PYTHONPATH=~/onnxruntime/build/cu130/Release \ +MODEL=models/gpt-oss-20b/variants/cuda_int4_int4_qmoe_rtn_matmul_only \ +GPU=0 PROMPT_LEN=512 GEN_LEN=128 REPS=10 WARMUP=3 CUDA_GRAPH=1 XQA=1 SYNC_LIB=0 \ +ORT_FORCE_DETERMINISTIC_MOE=1 \ +bash scripts/h200_18/bench_gpt_oss_ort_decode.sh +``` + +The disabled runs additionally set +`ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1`. + +| Run order | Mode | Decode latency ms/token | Decode throughput tok/s | Output CSV | +|-----------|------|-------------------------:|-------------------------:|------------| +| off then on | disabled | 2.870915 | 348.321025 | `~/ort_gptoss_routergemv_special_off.csv` | +| off then on | enabled | 2.819178 | 354.713290 | `~/ort_gptoss_routergemv_special_on.csv` | +| on then off | enabled | 2.824963 | 353.986941 | `~/ort_gptoss_routergemv_special_on_2.csv` | +| on then off | disabled | 2.870436 | 348.379097 | `~/ort_gptoss_routergemv_special_off_2.csv` | + +### Decision + +- Keep this as the current positive router experiment. The two CUDA-graph A/B + pairs show about `+1.6%` to `+1.8%` model-level decode throughput. +- This is more useful than the earlier node-only router graph fusion because it + reduces actual router GEMV kernel time and is exercised by the active GenAI + benchmark through provider dispatch. +- The specialization is deliberately exact-shape gated. Broader routing should + wait for more model shapes and correctness coverage. +- The earlier fused-router QMoE graph rewrite remains a separate prototype. Its + main value is future launch reduction if GEMV, bias, softmax, and top-k are + fused into one kernel; avoiding only the 32-float logits write is not expected + to be the dominant win. + +### Quality Follow-Up + +- The dispatch logic now uses named GPT-OSS router shape constants and a helper + predicate instead of repeating raw `M/N/K/block_size` literals at the launch + site. +- Added C++ CUDA coverage: + `MatMulNBits.Fp16_Int4_GptOssRouterShapeNoZeroPoint`. The test runs the exact + `M=1`, `N=32`, `K=2880`, `block_size=32`, no-zero-point FP16 shape with the + specialization enabled and then with + `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1` to cover the generic fallback. + +Validation command: + +```bash +cd ~/onnxruntime/build/cu130/Release +./onnxruntime_provider_test --gtest_filter='MatMulNBits.Fp16_Int4_GptOssRouterShapeNoZeroPoint' +``` + +Result: `1 test passed`. + +## 2026-06-19: Router Bias Fusion Into Exact-Shape GEMV + +### Change Under Test + +- Scope: GPT-OSS router chain `MatMulNBits(K=2880, N=32, bits=4, block_size=32)` + followed by `Add([32] bias)` before `QMoE`. +- Reuses the existing `MatMulNBitsFusion` graph transformer to rewrite + `MatMulNBits + Add -> MatMulNBits` with bias input slot 5. +- CPU behavior remains unchanged. CUDA fusion is gated to the exact GPT-OSS + router shape with no zero-points or `g_idx` input. +- The exact router GEMV kernel now accepts an optional bias pointer and adds the + `[32]` bias before storing the router logits. +- `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` disables only this CUDA router bias + fusion. `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1` still disables both + the exact router GEMV specialization and this bias fusion. + +### Graph Validation + +Optimized model export used the rebuilt Python runtime directly from +`build/cu130/Release`: + +```bash +cd /tmp +source ~/onnxruntime/.venv_cu130/bin/activate +PYTHONPATH=~/onnxruntime/build/cu130/Release \ +LD_LIBRARY_PATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/build/cu130/Release/onnxruntime/capi:$LD_LIBRARY_PATH \ +python +``` + +| Mode | Total `Add` nodes | Router `MatMulNBits` nodes with bias | +|------|------------------:|--------------------------------------:| +| enabled | 48 | 24 | +| `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` | 72 | 0 | + +This confirms all 24 real GPT-OSS router bias Adds are fused and the opt-out +restores the unfused graph. + +### Model-Level Decode Benchmark With CUDA Graph + +Both runs used the rebuilt runtime, CUDA graph enabled, XQA enabled, and +deterministic MoE tactic selection: + +```bash +PYTHONPATH=~/onnxruntime/build/cu130/Release \ +LD_LIBRARY_PATH=~/onnxruntime/build/cu130/Release:~/onnxruntime/build/cu130/Release/onnxruntime/capi:$LD_LIBRARY_PATH \ +MODEL=models/gpt-oss-20b/variants/cuda_int4_int4_qmoe_rtn_matmul_only \ +GPU=0 PROMPT_LEN=512 GEN_LEN=128 REPS=10 WARMUP=3 CUDA_GRAPH=1 XQA=1 SYNC_LIB=0 \ +ORT_FORCE_DETERMINISTIC_MOE=1 \ +bash scripts/h200_18/bench_gpt_oss_ort_decode.sh +``` + +| Mode | Decode latency ms/token | Decode throughput tok/s | +|------|-------------------------:|-------------------------:| +| enabled | 2.822998 | 354.233302 | +| `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` | 2.829066 | 353.473500 | + +### Decision + +- Keep the change. The gain is intentionally small, about `+0.2%` in this + CUDA-graph run, but the implementation is exact-shape gated, tested, and + removes 24 router Add kernels from the optimized graph. +- This result also bounds the value of router-only graph fusion: after the exact + GEMV specialization, the remaining separate router bias Add is measurable but + much smaller than the GEMV and softmax/top-k work. +- The next router/QMoE task should fuse more of the router tail, especially + GEMV plus softmax/top-k, so the 32 logits can remain in registers or shared + memory and avoid a separate top-k launch. + +### Validation + +```bash +cmake --build ~/onnxruntime/build/cu130/Release \ + --target onnxruntime_test_all onnxruntime_pybind11_state onnxruntime_provider_test \ + --parallel $(nproc) + +cd ~/onnxruntime/build/cu130/Release +./onnxruntime_test_all \ + --gtest_filter='GraphTransformationTests.MatMulNBitsBiasFusion:GraphTransformationTests.MatMulNBitsBiasFusionDoesNotFuseUnsupportedCudaShape:GraphTransformationTests.MatMulNBitsBiasFusionCanDisableCudaRouterFusion' +./onnxruntime_provider_test \ + --gtest_filter='MatMulNBits.Fp16_Int4_GptOssRouterShapeNoZeroPoint' +``` + +Result: graph transformer tests `3 passed`; provider test `1 passed`. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index ac93777a58e41..8df13ce692e87 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -5,6 +5,9 @@ #include #include #include + +#include + #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/quantization/matmul_nbits.cuh" @@ -16,6 +19,21 @@ namespace onnxruntime { namespace contrib { namespace cuda { +constexpr int kGptOssRouterM = 1; +constexpr int kGptOssRouterN = 32; +constexpr int kGptOssRouterK = 2880; +constexpr int kGptOssRouterBlockSize = 32; + +static bool GptOssRouterGemvSpecializationDisabled() { + const char* value = std::getenv("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION"); + return value != nullptr && value[0] == '1'; +} + +static bool IsGptOssRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { + return zero_points == nullptr && m == kGptOssRouterM && n == kGptOssRouterN && + k == kGptOssRouterK && block_size == kGptOssRouterBlockSize; +} + template __device__ __forceinline__ T WarpUniform(T value) { struct { @@ -325,6 +343,70 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt } } // namespace cuda +template +__global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4RouterKernel( + T* output, + const T* a_data, + const uint8_t* b_data_quant, + const T* scales_data, + const T* bias_data) { + constexpr int kHiddenSize = kGptOssRouterK; + constexpr int kBlockSize = kGptOssRouterBlockSize; + constexpr int kBlocksPerK = kHiddenSize / kBlockSize; + constexpr int kPerIter = kWarpSize * kElementsPerThreadPerIteration; + + const int lane_id = threadIdx.x; + const int warp_id = WarpUniform(threadIdx.y); + const int n_id = blockIdx.x * kColsPerThreadBlock + warp_id; + + a_data += lane_id << 3; + b_data_quant += n_id * kBlocksPerK * (kBlockSize / 2) + lane_id * 4; + scales_data += n_id * kBlocksPerK; + + T sums[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + int k_id = 0; + int scale_id = lane_id * 8 / kBlockSize; + +#define RouterUnRollReduction(unroll_size) \ + do { \ + constexpr int kUnroll = unroll_size; \ + constexpr int kUnrollStep = kUnroll * kPerIter; \ + const int k_unroll_bound = kHiddenSize - kHiddenSize % kUnrollStep; \ + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ + _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ + uint32_t value = *(reinterpret_cast(b_data_quant + kPerIter / 2 * i)); \ + T scale = scales_data[scale_id + kPerIter / kBlockSize * i]; \ + AccumulateEightElements4b(value, scale, 8, a_data + k_id + i * kPerIter, sums); \ + } \ + b_data_quant += kPerIter / 2 * kUnroll; \ + scale_id += kPerIter / kBlockSize * kUnroll; \ + } \ + } while (false) + + RouterUnRollReduction(16); + RouterUnRollReduction(4); + RouterUnRollReduction(1); +#undef RouterUnRollReduction + + if (k_id + lane_id * 8 < kHiddenSize) { + uint32_t value = *(reinterpret_cast(b_data_quant)); + T scale = scales_data[scale_id]; + AccumulateEightElements4b(value, scale, 8, a_data + k_id, sums); + } + + float sum = static_cast(sums[0] + sums[1] + sums[2] + sums[3] + sums[4] + sums[5] + sums[6] + sums[7]); + for (int i = kWarpSize / 2; i > 0; i = i / 2) { + sum += WARP_SHFL_DOWN(sum, i); + } + + if (lane_id == 0) { + if (bias_data != nullptr) { + sum += static_cast(bias_data[n_id]); + } + output[n_id] = sum; + } +} + template bool TryMatMul4Bits( T* output, @@ -332,6 +414,7 @@ bool TryMatMul4Bits( const uint8_t* b_data_quant, const T* scales_data, const uint8_t* zero_points, + const T* bias_data, int m, int n, int k, @@ -341,6 +424,20 @@ bool TryMatMul4Bits( if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > 1) { return false; } + + if (IsGptOssRouterGemvShape(zero_points, m, n, k, block_size) && + !GptOssRouterGemvSpecializationDisabled()) { + dim3 blocks(kGptOssRouterN / kColsPerThreadBlock, 1); + dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + MatMulFloatInt4RouterKernel<<>>( + output, a_data, b_data_quant, scales_data, bias_data); + return true; + } + + if (bias_data != nullptr) { + return false; + } + dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, m); dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); int blocks_per_K = (k + block_size - 1) / block_size; @@ -382,6 +479,7 @@ template bool TryMatMul4Bits( const uint8_t* b_data_quant, const float* scales_data, const uint8_t* zero_points, + const float* bias_data, int m, int n, int k, @@ -395,6 +493,7 @@ template bool TryMatMul4Bits( const uint8_t* b_data_quant, const half* scales_data, const uint8_t* zero_points, + const half* bias_data, int m, int n, int k, @@ -408,6 +507,7 @@ template bool TryMatMul4Bits( const uint8_t* b_data_quant, const nv_bfloat16* scales_data, const uint8_t* zero_points, + const nv_bfloat16* bias_data, int m, int n, int k, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index cd5358c41e898..d1e9348f295b4 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -272,7 +272,6 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const uint8_t* blob_data = is_prepacked_weight_ ? nullptr : b->Data(); const auto* scales_data = is_prepacked_scale_ ? nullptr : scales->Data(); const auto* zero_points_data = (is_prepacked_zero_point_ || zero_points == nullptr) ? nullptr : zero_points->DataRaw(); - const auto* bias_data = bias == nullptr ? nullptr : bias->Data(); #else const Tensor* b = ctx->Input(1); const Tensor* scales = ctx->Input(2); @@ -281,10 +280,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { const auto* scales_data = scales->Data(); const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw(); #endif - - if (bias != nullptr) { - ORT_THROW("MatMulNBits does not support bias in CUDA kernel"); - } + const auto* bias_data = bias == nullptr ? nullptr : bias->Data(); ORT_RETURN_IF_ERROR(matmul_nbits_helper::CheckInputs( a, b, scales, zero_points, reorder_idx, bias, N_, K_, block_size_, nbits_)); @@ -383,6 +379,7 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { blob_data, reinterpret_cast(scales_data), static_cast(zero_points_data), + reinterpret_cast(bias_data), m, n, k, @@ -393,6 +390,10 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { } } + if (bias != nullptr) { + ORT_THROW("MatMulNBits CUDA kernel only supports bias for the exact GPT-OSS router GEMV path"); + } + int64_t K_padded = (K_ + block_size_ - 1) / block_size_ * block_size_; // Chunked dequant+GEMM trades peak scratch memory for repeated kernel launches. diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index e97990d884ae5..b6be71b1d5646 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -15,6 +15,7 @@ bool TryMatMul4Bits( const uint8_t* b_data_quant, const T* scales_data, const uint8_t* zero_points, + const T* bias_data, int m, int n, int k, @@ -44,6 +45,7 @@ bool TryMatMulNBits( const uint8_t* b_data_quant, const T* scales_data, const uint8_t* zero_points, + const T* bias_data, int m, int n, int k, @@ -51,12 +53,15 @@ bool TryMatMulNBits( size_t shared_mem_per_block, cudaStream_t stream) { if (bits == 8) { + if (bias_data != nullptr) { + return false; + } return TryMatMul8Bits(output, a_data, b_data_quant, scales_data, zero_points, m, n, k, block_size, shared_mem_per_block, stream); } if (bits == 4) { - return TryMatMul4Bits(output, a_data, b_data_quant, scales_data, zero_points, + return TryMatMul4Bits(output, a_data, b_data_quant, scales_data, zero_points, bias_data, m, n, k, block_size, shared_mem_per_block, stream); } diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index a998eeacda734..a7a836f93f2d8 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -451,7 +451,7 @@ InlinedVector> GenerateTransformers( } #endif - transformers.emplace_back(std::make_unique(cpu_ep)); + transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); bool has_matmul_nbits_mlp_kernel = false; diff --git a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc index 2739ff96b1366..540fc7f9f40d5 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc @@ -5,6 +5,7 @@ #include "core/common/common.h" #include "core/optimizer/selectors_actions/actions.h" +#include "core/platform/env.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/graph/graph_utils.h" @@ -18,6 +19,47 @@ namespace { #if !defined(ORT_MINIMAL_BUILD) +constexpr int64_t kGptOssRouterN = 32; +constexpr int64_t kGptOssRouterK = 2880; +constexpr int64_t kGptOssRouterBlockSize = 32; +constexpr int64_t kGptOssRouterBits = 4; + +bool GptOssRouterGemvSpecializationDisabled() { + return Env::Default().GetEnvironmentVar("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION") == "1"; +} + +bool GptOssRouterBiasFusionDisabled() { + return Env::Default().GetEnvironmentVar("ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION") == "1"; +} + +bool HasInput(const Node& node, size_t input_index) { + const auto& input_defs = node.InputDefs(); + return input_defs.size() > input_index && input_defs[input_index]->Exists(); +} + +std::optional GetIntAttribute(const Node& node, const std::string& attr_name) { + const auto* attr = graph_utils::GetNodeAttribute(node, attr_name); + if (attr == nullptr) { + return std::nullopt; + } + return attr->i(); +} + +bool IsSupportedCudaGptOssRouterMatMulNBits(const Node& node) { + if (GptOssRouterGemvSpecializationDisabled() || GptOssRouterBiasFusionDisabled()) { + return false; + } + + if (HasInput(node, 3) || HasInput(node, 4)) { + return false; + } + + return GetIntAttribute(node, "N") == kGptOssRouterN && + GetIntAttribute(node, "K") == kGptOssRouterK && + GetIntAttribute(node, "block_size") == kGptOssRouterBlockSize && + GetIntAttribute(node, "bits") == kGptOssRouterBits; +} + namespace selectors { class BiasFusion : public NodeSelector { @@ -45,6 +87,11 @@ class BiasFusion : public NodeSelector { return std::nullopt; } + if (node.GetExecutionProviderType() == kCudaExecutionProvider && + !IsSupportedCudaGptOssRouterMatMulNBits(node)) { + return std::nullopt; + } + // check shape of other Add input // at this time, we only support adding a bias with shape [N] diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc index 1652d16f5cb66..9b6f90e1214bd 100644 --- a/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc +++ b/onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc @@ -478,6 +478,7 @@ class KernelTestFixture : public ::testing::Test { reinterpret_cast(d_weight_->data()), reinterpret_cast(d_scales_->data()), static_cast(d_uint8_zeros.data()), + nullptr, m_, n_, k_, block_size_, device_prop_.sharedMemPerBlock, s_); }, warmup_, repeats_, s_); diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index aadbbab1c135b..0ffa37b3bbdea 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -805,6 +805,42 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint) { } } +TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { + constexpr float abs_error = 0.1f; + + TestOptions opts{}; + opts.M = 1, opts.N = 32, opts.K = 2880; + opts.block_size = 32; + opts.has_zero_point = false; + opts.zp_is_4bit = true; + opts.output_abs_error = abs_error; + opts.output_rel_error = 0.02f; + + { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } + + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", "1"}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } + + { + opts.has_bias = true; + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } +} + TEST(MatMulNBits, BFloat16_Int4_NoZeroPoint) { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "Skipping BFloat16 8-bit MatMul tests on CUDA < 8.0"; diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index cbb060fe4a14a..923ea0c2e4988 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -95,6 +95,7 @@ #include "test/util/include/asserts.h" #include "test/util/include/default_providers.h" #include "test/util/include/inference_session_wrapper.h" +#include "test/util/include/scoped_env_vars.h" #include "test/util/include/temp_dir.h" #include "test/util/include/test_utils.h" #ifdef ENABLE_TRAINING @@ -10789,20 +10790,30 @@ TEST_F(GraphTransformationTests, GatherToSliceFusion) { #if !defined(DISABLE_CONTRIB_OPS) TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}, + {"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", optional{}}}}; + struct TestOptions { bool bias_is_first_add_input{false}; bool add_produces_graph_output{false}; + bool use_cuda_ep{false}; + bool use_gpt_oss_router_shape{false}; }; auto run_test = [&logger = *logger_](const TestOptions& opts) { SCOPED_TRACE(MakeString("bias_is_first_add_input:", opts.bias_is_first_add_input, - ", add_produces_graph_output:", opts.add_produces_graph_output)); + ", add_produces_graph_output:", opts.add_produces_graph_output, + ", use_cuda_ep:", opts.use_cuda_ep, + ", use_gpt_oss_router_shape:", opts.use_gpt_oss_router_shape)); auto build_test_case = [&](ModelTestBuilder& builder) { constexpr size_t qbits = 4; constexpr size_t block_size = 32; - constexpr int64_t M = 2, K = 4, N = 8; + const int64_t M = opts.use_gpt_oss_router_shape ? 1 : 2; + const int64_t K = opts.use_gpt_oss_router_shape ? 2880 : 4; + const int64_t N = opts.use_gpt_oss_router_shape ? 32 : 8; int q_rows, q_cols; MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, @@ -10820,8 +10831,10 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { uint8_t{0}, uint8_t{255}); auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, 1.0f, 2.0f); - auto* B_zero_points = builder.MakeInitializer({static_cast(q_zp_size_in_bytes)}, - uint8_t{0}, uint8_t{255}); + NodeArg* B_zero_points = opts.use_gpt_oss_router_shape + ? builder.MakeEmptyInput() + : builder.MakeInitializer({static_cast(q_zp_size_in_bytes)}, + uint8_t{0}, uint8_t{255}); auto* matmul_output = builder.MakeIntermediate(); @@ -10833,6 +10846,9 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { matmul.AddAttribute("K", K); matmul.AddAttribute("block_size", static_cast(block_size)); matmul.AddAttribute("bits", static_cast(qbits)); + if (opts.use_cuda_ep) { + matmul.SetExecutionProviderType(kCudaExecutionProvider); + } auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); @@ -10840,15 +10856,21 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { auto* add_output = opts.add_produces_graph_output ? graph_output : builder.MakeIntermediate(); - builder.AddNode("Add", - {opts.bias_is_first_add_input ? Bias : matmul_output, - opts.bias_is_first_add_input ? matmul_output : Bias}, - {add_output}); + auto& add = builder.AddNode("Add", + {opts.bias_is_first_add_input ? Bias : matmul_output, + opts.bias_is_first_add_input ? matmul_output : Bias}, + {add_output}); + if (opts.use_cuda_ep) { + add.SetExecutionProviderType(kCudaExecutionProvider); + } if (!opts.add_produces_graph_output) { - builder.AddNode("Identity", - {add_output}, - {graph_output}); + auto& identity = builder.AddNode("Identity", + {add_output}, + {graph_output}); + if (opts.use_cuda_ep) { + identity.SetExecutionProviderType(kCudaExecutionProvider); + } } }; @@ -10874,10 +10896,118 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { opts.bias_is_first_add_input = bias_is_first_add_input; opts.add_produces_graph_output = add_produces_graph_output; run_test(opts); + + opts.use_cuda_ep = true; + opts.use_gpt_oss_router_shape = true; + run_test(opts); } } } +TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseUnsupportedCudaShape) { + auto build_test_case = [](ModelTestBuilder& builder) { + constexpr size_t qbits = 4; + constexpr size_t block_size = 32; + + constexpr int64_t M = 2, K = 4, N = 8; + + int q_rows, q_cols; + MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, + K, N, + q_rows, q_cols); + + size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; + MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, + K, N, + q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); + + auto* A = builder.MakeInput(std::vector{M, K}, "A"); + auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, + uint8_t{0}, uint8_t{255}); + auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, + 1.0f, 2.0f); + auto* B_zero_points = builder.MakeInitializer({static_cast(q_zp_size_in_bytes)}, + uint8_t{0}, uint8_t{255}); + + auto* matmul_output = builder.MakeIntermediate(); + auto& matmul = builder.AddNode("MatMulNBits", + {A, B_data, B_scales, B_zero_points}, + {matmul_output}, + kMSDomain); + matmul.AddAttribute("N", N); + matmul.AddAttribute("K", K); + matmul.AddAttribute("block_size", static_cast(block_size)); + matmul.AddAttribute("bits", static_cast(qbits)); + matmul.SetExecutionProviderType(kCudaExecutionProvider); + + auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); + auto* add_output = builder.MakeOutput(); + auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); + add.SetExecutionProviderType(kCudaExecutionProvider); + }; + + auto post_graph_checker = [](Graph& graph) { + auto op_count = CountOpsInGraph(graph); + EXPECT_EQ(op_count["Add"], 1); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), + TransformerLevel::Level2, 0, nullptr, post_graph_checker)); +} + +TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionCanDisableCudaRouterFusion) { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", "1"}}}; + + auto build_test_case = [](ModelTestBuilder& builder) { + constexpr size_t qbits = 4; + constexpr size_t block_size = 32; + + constexpr int64_t M = 1, K = 2880, N = 32; + + int q_rows, q_cols; + MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, + K, N, + q_rows, q_cols); + + size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; + MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, + K, N, + q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); + + auto* A = builder.MakeInput(std::vector{M, K}, "A"); + auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, + uint8_t{0}, uint8_t{255}); + auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, + 1.0f, 2.0f); + + auto* matmul_output = builder.MakeIntermediate(); + auto& matmul = builder.AddNode("MatMulNBits", + {A, B_data, B_scales, builder.MakeEmptyInput()}, + {matmul_output}, + kMSDomain); + matmul.AddAttribute("N", N); + matmul.AddAttribute("K", K); + matmul.AddAttribute("block_size", static_cast(block_size)); + matmul.AddAttribute("bits", static_cast(qbits)); + matmul.SetExecutionProviderType(kCudaExecutionProvider); + + auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); + auto* add_output = builder.MakeOutput(); + auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); + add.SetExecutionProviderType(kCudaExecutionProvider); + }; + + auto post_graph_checker = [](Graph& graph) { + auto op_count = CountOpsInGraph(graph); + EXPECT_EQ(op_count["Add"], 1); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), + TransformerLevel::Level2, 0, nullptr, post_graph_checker)); +} + #endif // !defined(DISABLE_CONTRIB_OPS) // Tests for graceful handling of zero-element initializers in optimizer passes. From 13e6ffb43f37ffab8c096530202b7cfd14a455c5 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 12:12:37 -0700 Subject: [PATCH 02/10] Require static M==1 for CUDA router bias fusion The CUDA MatMulNBits kernel only accepts a fused bias on the exact router GEMV fast path (M==1). The fusion gate previously checked only N/K/bits/ block_size attributes, so a MatMulNBits+Add with matching attributes but M>1 (or dynamic M) could be fused and then throw at runtime. Gate CUDA bias fusion on A having a statically-known M==1, and add a regression test. --- .../core/optimizer/matmul_nbits_fusion.cc | 29 ++++++++++ .../test/optimizer/graph_transform_test.cc | 58 +++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc index 540fc7f9f40d5..c9d12f9649abb 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc @@ -45,6 +45,31 @@ std::optional GetIntAttribute(const Node& node, const std::string& attr return attr->i(); } +// The CUDA MatMulNBits kernel only accepts a fused bias on the exact router GEMV fast path, +// which requires M == 1 (M is the product of all A dimensions except the last, K). If bias +// is fused when M > 1 (or M is dynamic/unknown), TryMatMul4Bits rejects the bias path and +// MatMulNBits::ComputeInternal throws at runtime. Only allow fusion when A's M is statically 1. +bool MatMulNBitsHasStaticGemvShape(const Node& node) { + const auto& input_defs = node.InputDefs(); + if (input_defs.empty() || !input_defs[0]->Exists()) { + return false; + } + + const auto* a_shape = input_defs[0]->Shape(); + if (a_shape == nullptr || a_shape->dim_size() < 1) { + return false; + } + + // Every leading dimension (everything except the last K dim) must be statically 1. + for (int i = 0; i < a_shape->dim_size() - 1; ++i) { + if (!utils::HasDimValue(a_shape->dim(i)) || a_shape->dim(i).dim_value() != 1) { + return false; + } + } + + return true; +} + bool IsSupportedCudaGptOssRouterMatMulNBits(const Node& node) { if (GptOssRouterGemvSpecializationDisabled() || GptOssRouterBiasFusionDisabled()) { return false; @@ -54,6 +79,10 @@ bool IsSupportedCudaGptOssRouterMatMulNBits(const Node& node) { return false; } + if (!MatMulNBitsHasStaticGemvShape(node)) { + return false; + } + return GetIntAttribute(node, "N") == kGptOssRouterN && GetIntAttribute(node, "K") == kGptOssRouterK && GetIntAttribute(node, "block_size") == kGptOssRouterBlockSize && diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 923ea0c2e4988..15d1834696a58 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -10956,6 +10956,64 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseUnsupportedCuda TransformerLevel::Level2, 0, nullptr, post_graph_checker)); } +// The CUDA bias fusion fast path only handles the router GEMV (M == 1). Even when +// N/K/block_size/bits match the router specialization, a non-GEMV A input (M > 1) must +// not be fused, otherwise the CUDA kernel would throw at runtime. +TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseCudaRouterShapeWhenMGreaterThanOne) { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}, + {"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", optional{}}}}; + + auto build_test_case = [](ModelTestBuilder& builder) { + constexpr size_t qbits = 4; + constexpr size_t block_size = 32; + + // Exact GPT-OSS router N/K but with M > 1 (not a GEMV). + constexpr int64_t M = 2, K = 2880, N = 32; + + int q_rows, q_cols; + MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, + K, N, + q_rows, q_cols); + + size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; + MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, + K, N, + q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); + + auto* A = builder.MakeInput(std::vector{M, K}, "A"); + auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, + uint8_t{0}, uint8_t{255}); + auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, + 1.0f, 2.0f); + + auto* matmul_output = builder.MakeIntermediate(); + auto& matmul = builder.AddNode("MatMulNBits", + {A, B_data, B_scales, builder.MakeEmptyInput()}, + {matmul_output}, + kMSDomain); + matmul.AddAttribute("N", N); + matmul.AddAttribute("K", K); + matmul.AddAttribute("block_size", static_cast(block_size)); + matmul.AddAttribute("bits", static_cast(qbits)); + matmul.SetExecutionProviderType(kCudaExecutionProvider); + + auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); + auto* add_output = builder.MakeOutput(); + auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); + add.SetExecutionProviderType(kCudaExecutionProvider); + }; + + auto post_graph_checker = [](Graph& graph) { + auto op_count = CountOpsInGraph(graph); + EXPECT_EQ(op_count["Add"], 1); + return Status::OK(); + }; + + ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), + TransformerLevel::Level2, 0, nullptr, post_graph_checker)); +} + TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionCanDisableCudaRouterFusion) { ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", "1"}}}; From f69cb5fd4fedbfc1d50b05463a2a6a5dc6fa0f51 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 13:25:41 -0700 Subject: [PATCH 03/10] Fix -Wshorten-64-to-32 error in MatMulNBitsBiasFusion test K and N became runtime const int64_t (router shape selection), so passing them to MlasBlockwiseQuantizedShape/BufferSizes (int params) now triggers -Werror,-Wshorten-64-to-32 on macOS/arm64 and Android. Cast to int. --- onnxruntime/test/optimizer/graph_transform_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 15d1834696a58..57aa10602eccd 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -10817,12 +10817,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { int q_rows, q_cols; MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_rows, q_cols); size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); auto* A = builder.MakeInput(std::vector{M, K}, "A"); From 5acf8a99b4365af884af04412b2669dae6d43204 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 14:59:24 -0700 Subject: [PATCH 04/10] Address review: qualify std::optional and add explicit int casts in tests - Qualify optional as std::optional in matmul_4bits_test.cc (no using namespace std there). - Add explicit static_cast for K/N passed to MlasBlockwiseQuantizedShape and MlasBlockwiseQuantizedBufferSizes to avoid int64->int narrowing warnings-as-errors. --- onnxruntime/test/contrib_ops/matmul_4bits_test.cc | 4 ++-- onnxruntime/test/optimizer/graph_transform_test.cc | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 0ffa37b3bbdea..812fb16291550 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -818,7 +818,7 @@ TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { { ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}}}; + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; std::vector> eps; eps.push_back(DefaultCudaExecutionProvider()); RunTest(opts, std::move(eps)); @@ -834,7 +834,7 @@ TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { { opts.has_bias = true; ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}}}; + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; std::vector> eps; eps.push_back(DefaultCudaExecutionProvider()); RunTest(opts, std::move(eps)); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 57aa10602eccd..fac158876e8b7 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -10913,12 +10913,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseUnsupportedCuda int q_rows, q_cols; MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_rows, q_cols); size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); auto* A = builder.MakeInput(std::vector{M, K}, "A"); @@ -10973,12 +10973,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseCudaRouterShape int q_rows, q_cols; MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_rows, q_cols); size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); auto* A = builder.MakeInput(std::vector{M, K}, "A"); @@ -11025,12 +11025,12 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionCanDisableCudaRouterFusion int q_rows, q_cols; MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_rows, q_cols); size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - K, N, + static_cast(K), static_cast(N), q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); auto* A = builder.MakeInput(std::vector{M, K}, "A"); From b0a4d86e4a857b6c075e275bfbc033fc8658d5f7 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 16:49:58 -0700 Subject: [PATCH 05/10] add fallback --- .../cuda/quantization/matmul_4bits.cu | 43 ++++- .../cuda/quantization/matmul_nbits.cc | 45 ++++- .../cuda/quantization/matmul_nbits.cuh | 10 ++ .../core/optimizer/matmul_nbits_fusion.cc | 76 -------- .../test/optimizer/graph_transform_test.cc | 170 +----------------- 5 files changed, 95 insertions(+), 249 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index 8df13ce692e87..bab9031ef07a2 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -6,12 +6,15 @@ #include #include -#include - #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" #include "contrib_ops/cuda/quantization/matmul_nbits.cuh" +// Include env_var_utils.h after cuda_common.h: the latter transitively pulls in provider_api.h, +// which defines SHARED_PROVIDER. That guard suppresses env_var_utils.h's own logging.h include and +// avoids redefining CREATE_MESSAGE/LOGS_CATEGORY in this provider-bridge translation unit. +#include "core/platform/env_var_utils.h" + using namespace onnxruntime::cuda; using namespace cub; @@ -25,8 +28,8 @@ constexpr int kGptOssRouterK = 2880; constexpr int kGptOssRouterBlockSize = 32; static bool GptOssRouterGemvSpecializationDisabled() { - const char* value = std::getenv("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION"); - return value != nullptr && value[0] == '1'; + // Use ORT's cross-platform env var helper instead of std::getenv, which is unsafe on Windows. + return ParseEnvironmentVariableWithDefault("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", false); } static bool IsGptOssRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { @@ -515,6 +518,38 @@ template bool TryMatMul4Bits( size_t shared_mem_per_block, cudaStream_t stream); +template +__global__ void MatMulNBitsBiasAddKernel(T* output, const T* bias_data, int n, int64_t total) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + const int col = static_cast(idx % n); + // Accumulate in float to stay accurate for half/bfloat16. + output[idx] = static_cast(static_cast(output[idx]) + static_cast(bias_data[col])); +} + +template +void LaunchMatMulNBitsBiasAdd(T* output, const T* bias_data, int m, int n, cudaStream_t stream) { + const int64_t total = static_cast(m) * static_cast(n); + if (total == 0) { + return; + } + constexpr int kThreadsPerBlock = 256; + const int64_t num_blocks = (total + kThreadsPerBlock - 1) / kThreadsPerBlock; + MatMulNBitsBiasAddKernel<<(num_blocks), kThreadsPerBlock, 0, stream>>>( + output, bias_data, n, total); +} + +template void LaunchMatMulNBitsBiasAdd( + float* output, const float* bias_data, int m, int n, cudaStream_t stream); + +template void LaunchMatMulNBitsBiasAdd( + half* output, const half* bias_data, int m, int n, cudaStream_t stream); + +template void LaunchMatMulNBitsBiasAdd( + nv_bfloat16* output, const nv_bfloat16* bias_data, int m, int n, cudaStream_t stream); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc index d1e9348f295b4..896291ce01b5d 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc @@ -372,6 +372,8 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { #endif if ((reorder_idx_data == nullptr) && (!zero_points || !zero_points->IsDataType())) { + // First, try the fused fast path. It handles bias only for the GPT-OSS router GEMV + // specialization; for other shapes it fails when bias is present. if (TryMatMulNBits( static_cast(nbits_), reinterpret_cast(Y->MutableData()), @@ -388,12 +390,39 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { stream)) { return Status::OK(); } - } - if (bias != nullptr) { - ORT_THROW("MatMulNBits CUDA kernel only supports bias for the exact GPT-OSS router GEMV path"); + // If bias prevented the fused path from running, retry the fast path without bias and + // add the bias with a separate kernel. This keeps the lightweight GEMV path for cases + // (e.g. 8-bit, or 4-bit non-router shapes) where only the bias was unsupported. + if (bias_data != nullptr && + TryMatMulNBits( + static_cast(nbits_), + reinterpret_cast(Y->MutableData()), + reinterpret_cast(a_data), + blob_data, + reinterpret_cast(scales_data), + static_cast(zero_points_data), + /*bias_data*/ static_cast(nullptr), + m, + n, + k, + SafeInt(block_size_), + GetDeviceProp().sharedMemPerBlock, + stream)) { + LaunchMatMulNBitsBiasAdd( + reinterpret_cast(Y->MutableData()), + reinterpret_cast(bias_data), + m, + n, + stream); + return Status::OK(); + } } + // When bias is present but the fused router GEMV specialization above did not apply, + // fall back to the generic dequantize + GEMM path (which ignores bias) and add the + // bias with a separate kernel after the GEMM completes. + int64_t K_padded = (K_ + block_size_ - 1) / block_size_ * block_size_; // Chunked dequant+GEMM trades peak scratch memory for repeated kernel launches. @@ -578,6 +607,16 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { GetDeviceProp(), UseTF32())); } + + // Fallback bias handling: the generic dequant + GEMM path above ignores bias, so add it here. + if (bias_data != nullptr) { + LaunchMatMulNBitsBiasAdd( + reinterpret_cast(Y->MutableData()), + reinterpret_cast(bias_data), + m, + n, + stream); + } } return Status::OK(); diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh index b6be71b1d5646..61eb62fef4a2b 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh @@ -68,6 +68,16 @@ bool TryMatMulNBits( return false; } +// Adds a per-column bias of shape [n] to the output of shape [m, n] (row-major). +// Used as a fallback when the fused bias GEMV specialization does not apply. +template +void LaunchMatMulNBitsBiasAdd( + T* output, + const T* bias_data, + int m, + int n, + cudaStream_t stream); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc index c9d12f9649abb..2739ff96b1366 100644 --- a/onnxruntime/core/optimizer/matmul_nbits_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_nbits_fusion.cc @@ -5,7 +5,6 @@ #include "core/common/common.h" #include "core/optimizer/selectors_actions/actions.h" -#include "core/platform/env.h" #if !defined(ORT_MINIMAL_BUILD) #include "core/graph/graph_utils.h" @@ -19,76 +18,6 @@ namespace { #if !defined(ORT_MINIMAL_BUILD) -constexpr int64_t kGptOssRouterN = 32; -constexpr int64_t kGptOssRouterK = 2880; -constexpr int64_t kGptOssRouterBlockSize = 32; -constexpr int64_t kGptOssRouterBits = 4; - -bool GptOssRouterGemvSpecializationDisabled() { - return Env::Default().GetEnvironmentVar("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION") == "1"; -} - -bool GptOssRouterBiasFusionDisabled() { - return Env::Default().GetEnvironmentVar("ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION") == "1"; -} - -bool HasInput(const Node& node, size_t input_index) { - const auto& input_defs = node.InputDefs(); - return input_defs.size() > input_index && input_defs[input_index]->Exists(); -} - -std::optional GetIntAttribute(const Node& node, const std::string& attr_name) { - const auto* attr = graph_utils::GetNodeAttribute(node, attr_name); - if (attr == nullptr) { - return std::nullopt; - } - return attr->i(); -} - -// The CUDA MatMulNBits kernel only accepts a fused bias on the exact router GEMV fast path, -// which requires M == 1 (M is the product of all A dimensions except the last, K). If bias -// is fused when M > 1 (or M is dynamic/unknown), TryMatMul4Bits rejects the bias path and -// MatMulNBits::ComputeInternal throws at runtime. Only allow fusion when A's M is statically 1. -bool MatMulNBitsHasStaticGemvShape(const Node& node) { - const auto& input_defs = node.InputDefs(); - if (input_defs.empty() || !input_defs[0]->Exists()) { - return false; - } - - const auto* a_shape = input_defs[0]->Shape(); - if (a_shape == nullptr || a_shape->dim_size() < 1) { - return false; - } - - // Every leading dimension (everything except the last K dim) must be statically 1. - for (int i = 0; i < a_shape->dim_size() - 1; ++i) { - if (!utils::HasDimValue(a_shape->dim(i)) || a_shape->dim(i).dim_value() != 1) { - return false; - } - } - - return true; -} - -bool IsSupportedCudaGptOssRouterMatMulNBits(const Node& node) { - if (GptOssRouterGemvSpecializationDisabled() || GptOssRouterBiasFusionDisabled()) { - return false; - } - - if (HasInput(node, 3) || HasInput(node, 4)) { - return false; - } - - if (!MatMulNBitsHasStaticGemvShape(node)) { - return false; - } - - return GetIntAttribute(node, "N") == kGptOssRouterN && - GetIntAttribute(node, "K") == kGptOssRouterK && - GetIntAttribute(node, "block_size") == kGptOssRouterBlockSize && - GetIntAttribute(node, "bits") == kGptOssRouterBits; -} - namespace selectors { class BiasFusion : public NodeSelector { @@ -116,11 +45,6 @@ class BiasFusion : public NodeSelector { return std::nullopt; } - if (node.GetExecutionProviderType() == kCudaExecutionProvider && - !IsSupportedCudaGptOssRouterMatMulNBits(node)) { - return std::nullopt; - } - // check shape of other Add input // at this time, we only support adding a bias with shape [N] diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index fac158876e8b7..1f07bf3204894 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -10790,10 +10790,6 @@ TEST_F(GraphTransformationTests, GatherToSliceFusion) { #if !defined(DISABLE_CONTRIB_OPS) TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { - ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}, - {"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", optional{}}}}; - struct TestOptions { bool bias_is_first_add_input{false}; bool add_produces_graph_output{false}; @@ -10897,175 +10893,17 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { opts.add_produces_graph_output = add_produces_graph_output; run_test(opts); + // CUDA now fuses bias generically for any shape (the kernel adds the bias with a + // separate kernel when the fused GEMV fast path does not apply). opts.use_cuda_ep = true; + run_test(opts); + opts.use_gpt_oss_router_shape = true; run_test(opts); } } } -TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseUnsupportedCudaShape) { - auto build_test_case = [](ModelTestBuilder& builder) { - constexpr size_t qbits = 4; - constexpr size_t block_size = 32; - - constexpr int64_t M = 2, K = 4, N = 8; - - int q_rows, q_cols; - MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_rows, q_cols); - - size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; - MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); - - auto* A = builder.MakeInput(std::vector{M, K}, "A"); - auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, - uint8_t{0}, uint8_t{255}); - auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, - 1.0f, 2.0f); - auto* B_zero_points = builder.MakeInitializer({static_cast(q_zp_size_in_bytes)}, - uint8_t{0}, uint8_t{255}); - - auto* matmul_output = builder.MakeIntermediate(); - auto& matmul = builder.AddNode("MatMulNBits", - {A, B_data, B_scales, B_zero_points}, - {matmul_output}, - kMSDomain); - matmul.AddAttribute("N", N); - matmul.AddAttribute("K", K); - matmul.AddAttribute("block_size", static_cast(block_size)); - matmul.AddAttribute("bits", static_cast(qbits)); - matmul.SetExecutionProviderType(kCudaExecutionProvider); - - auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); - auto* add_output = builder.MakeOutput(); - auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); - add.SetExecutionProviderType(kCudaExecutionProvider); - }; - - auto post_graph_checker = [](Graph& graph) { - auto op_count = CountOpsInGraph(graph); - EXPECT_EQ(op_count["Add"], 1); - return Status::OK(); - }; - - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), - TransformerLevel::Level2, 0, nullptr, post_graph_checker)); -} - -// The CUDA bias fusion fast path only handles the router GEMV (M == 1). Even when -// N/K/block_size/bits match the router specialization, a non-GEMV A input (M > 1) must -// not be fused, otherwise the CUDA kernel would throw at runtime. -TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionDoesNotFuseCudaRouterShapeWhenMGreaterThanOne) { - ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", optional{}}, - {"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", optional{}}}}; - - auto build_test_case = [](ModelTestBuilder& builder) { - constexpr size_t qbits = 4; - constexpr size_t block_size = 32; - - // Exact GPT-OSS router N/K but with M > 1 (not a GEMV). - constexpr int64_t M = 2, K = 2880, N = 32; - - int q_rows, q_cols; - MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_rows, q_cols); - - size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; - MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); - - auto* A = builder.MakeInput(std::vector{M, K}, "A"); - auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, - uint8_t{0}, uint8_t{255}); - auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, - 1.0f, 2.0f); - - auto* matmul_output = builder.MakeIntermediate(); - auto& matmul = builder.AddNode("MatMulNBits", - {A, B_data, B_scales, builder.MakeEmptyInput()}, - {matmul_output}, - kMSDomain); - matmul.AddAttribute("N", N); - matmul.AddAttribute("K", K); - matmul.AddAttribute("block_size", static_cast(block_size)); - matmul.AddAttribute("bits", static_cast(qbits)); - matmul.SetExecutionProviderType(kCudaExecutionProvider); - - auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); - auto* add_output = builder.MakeOutput(); - auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); - add.SetExecutionProviderType(kCudaExecutionProvider); - }; - - auto post_graph_checker = [](Graph& graph) { - auto op_count = CountOpsInGraph(graph); - EXPECT_EQ(op_count["Add"], 1); - return Status::OK(); - }; - - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), - TransformerLevel::Level2, 0, nullptr, post_graph_checker)); -} - -TEST_F(GraphTransformationTests, MatMulNBitsBiasFusionCanDisableCudaRouterFusion) { - ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION", "1"}}}; - - auto build_test_case = [](ModelTestBuilder& builder) { - constexpr size_t qbits = 4; - constexpr size_t block_size = 32; - - constexpr int64_t M = 1, K = 2880, N = 32; - - int q_rows, q_cols; - MlasBlockwiseQuantizedShape(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_rows, q_cols); - - size_t q_data_size_in_bytes, q_scale_size, q_zp_size_in_bytes; - MlasBlockwiseQuantizedBufferSizes(block_size, /* columnwise */ true, - static_cast(K), static_cast(N), - q_data_size_in_bytes, q_scale_size, &q_zp_size_in_bytes); - - auto* A = builder.MakeInput(std::vector{M, K}, "A"); - auto* B_data = builder.MakeInitializer({int64_t{q_rows}, int64_t{q_cols}}, - uint8_t{0}, uint8_t{255}); - auto* B_scales = builder.MakeInitializer({static_cast(q_scale_size)}, - 1.0f, 2.0f); - - auto* matmul_output = builder.MakeIntermediate(); - auto& matmul = builder.AddNode("MatMulNBits", - {A, B_data, B_scales, builder.MakeEmptyInput()}, - {matmul_output}, - kMSDomain); - matmul.AddAttribute("N", N); - matmul.AddAttribute("K", K); - matmul.AddAttribute("block_size", static_cast(block_size)); - matmul.AddAttribute("bits", static_cast(qbits)); - matmul.SetExecutionProviderType(kCudaExecutionProvider); - - auto* Bias = builder.MakeInput(std::vector{N}, "Bias"); - auto* add_output = builder.MakeOutput(); - auto& add = builder.AddNode("Add", {matmul_output, Bias}, {add_output}); - add.SetExecutionProviderType(kCudaExecutionProvider); - }; - - auto post_graph_checker = [](Graph& graph) { - auto op_count = CountOpsInGraph(graph); - EXPECT_EQ(op_count["Add"], 1); - return Status::OK(); - }; - - ASSERT_STATUS_OK(TestGraphTransformer(build_test_case, 21, *logger_, std::make_unique(), - TransformerLevel::Level2, 0, nullptr, post_graph_checker)); -} - #endif // !defined(DISABLE_CONTRIB_OPS) // Tests for graceful handling of zero-element initializers in optimizer passes. From 25789d98f4cdbb0b36271721da944c17fe024a9a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 17:29:35 -0700 Subject: [PATCH 06/10] extend to qwen and gemma --- .../cuda/quantization/matmul_4bits.cu | 76 +++++++++++++------ 1 file changed, 53 insertions(+), 23 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index bab9031ef07a2..848d78925b09c 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -22,19 +22,37 @@ namespace onnxruntime { namespace contrib { namespace cuda { -constexpr int kGptOssRouterM = 1; +// The specialized router GEMV kernel only handles M=1, or batch size 1. +constexpr int kRouterM = 1; + +// MoE router shapes (N = number of experts, K = hidden size) constexpr int kGptOssRouterN = 32; constexpr int kGptOssRouterK = 2880; -constexpr int kGptOssRouterBlockSize = 32; +constexpr int kQwen3RouterN = 256; +constexpr int kQwen3RouterK = 2048; +constexpr int kGemmaRouterN = 128; +constexpr int kGemmaRouterK = 2112; -static bool GptOssRouterGemvSpecializationDisabled() { +static bool IsRouterGemvSpecializationDisabled() { // Use ORT's cross-platform env var helper instead of std::getenv, which is unsafe on Windows. return ParseEnvironmentVariableWithDefault("ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", false); } -static bool IsGptOssRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { - return zero_points == nullptr && m == kGptOssRouterM && n == kGptOssRouterN && - k == kGptOssRouterK && block_size == kGptOssRouterBlockSize; +// The router GEMV kernel handles any symmetric (no zero point) M=1 shape with an int4 group size of +// 32 or 64 (whichever quantizes best) and N divisible by kColsPerThreadBlock. We gate on the known +// MoE router shapes to avoid changing the dispatch for general MatMulNBits cases. K must be a multiple +// of the group size (always true for a router) and N a multiple of kColsPerThreadBlock (checked in +// TryMatMul4Bits). Note kPerIter (256) is divisible by both 32 and 64, so the scale stride is exact. +static bool IsSupportedRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { + if (zero_points != nullptr || m != kRouterM || (block_size != 32 && block_size != 64)) { + return false; + } + if (k % block_size != 0) { + return false; + } + return (n == kGptOssRouterN && k == kGptOssRouterK) || // gpt-oss-20b + (n == kQwen3RouterN && k == kQwen3RouterK) || // Qwen3.6-35B-A3B + (n == kGemmaRouterN && k == kGemmaRouterK); // gemma-4-26B-A4B } template @@ -346,43 +364,50 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt } } // namespace cuda -template +// GEMV specialization for MoE routers: output(1, N) = a(1, K) x dequant(B(N, K)) [+ bias(N)]. +// B is 4-bit block-quantized (symmetric, no zero point) with group size BlockSize (32 or 64). One warp +// computes one expert column; the thread block holds kColsPerThreadBlock warps. N is passed via the +// grid and K at runtime, so a single instantiation per (T, BlockSize) serves every router shape. +// Requirements (satisfied by the dispatch in TryMatMul4Bits): N % kColsPerThreadBlock == 0, +// K % BlockSize == 0, and kPerIter % BlockSize == 0 so the per-iteration scale stride is exact. +template __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt4RouterKernel( T* output, const T* a_data, const uint8_t* b_data_quant, const T* scales_data, - const T* bias_data) { - constexpr int kHiddenSize = kGptOssRouterK; - constexpr int kBlockSize = kGptOssRouterBlockSize; - constexpr int kBlocksPerK = kHiddenSize / kBlockSize; + const T* bias_data, + int n, + int k) { constexpr int kPerIter = kWarpSize * kElementsPerThreadPerIteration; + static_assert(kPerIter % BlockSize == 0, "kPerIter must be a multiple of BlockSize for exact scale stride"); const int lane_id = threadIdx.x; const int warp_id = WarpUniform(threadIdx.y); const int n_id = blockIdx.x * kColsPerThreadBlock + warp_id; + // Each column occupies k/2 bytes of packed 4-bit weights and k/BlockSize scales. a_data += lane_id << 3; - b_data_quant += n_id * kBlocksPerK * (kBlockSize / 2) + lane_id * 4; - scales_data += n_id * kBlocksPerK; + b_data_quant += static_cast(n_id) * (k / 2) + lane_id * 4; + scales_data += static_cast(n_id) * (k / BlockSize); T sums[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; int k_id = 0; - int scale_id = lane_id * 8 / kBlockSize; + int scale_id = lane_id * 8 / BlockSize; #define RouterUnRollReduction(unroll_size) \ do { \ constexpr int kUnroll = unroll_size; \ constexpr int kUnrollStep = kUnroll * kPerIter; \ - const int k_unroll_bound = kHiddenSize - kHiddenSize % kUnrollStep; \ + const int k_unroll_bound = k - k % kUnrollStep; \ for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ uint32_t value = *(reinterpret_cast(b_data_quant + kPerIter / 2 * i)); \ - T scale = scales_data[scale_id + kPerIter / kBlockSize * i]; \ + T scale = scales_data[scale_id + kPerIter / BlockSize * i]; \ AccumulateEightElements4b(value, scale, 8, a_data + k_id + i * kPerIter, sums); \ } \ b_data_quant += kPerIter / 2 * kUnroll; \ - scale_id += kPerIter / kBlockSize * kUnroll; \ + scale_id += kPerIter / BlockSize * kUnroll; \ } \ } while (false) @@ -391,7 +416,7 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt RouterUnRollReduction(1); #undef RouterUnRollReduction - if (k_id + lane_id * 8 < kHiddenSize) { + if (k_id + lane_id * 8 < k) { uint32_t value = *(reinterpret_cast(b_data_quant)); T scale = scales_data[scale_id]; AccumulateEightElements4b(value, scale, 8, a_data + k_id, sums); @@ -428,12 +453,17 @@ bool TryMatMul4Bits( return false; } - if (IsGptOssRouterGemvShape(zero_points, m, n, k, block_size) && - !GptOssRouterGemvSpecializationDisabled()) { - dim3 blocks(kGptOssRouterN / kColsPerThreadBlock, 1); + if (IsSupportedRouterGemvShape(zero_points, m, n, k, block_size) && + !IsRouterGemvSpecializationDisabled()) { + dim3 blocks(n / kColsPerThreadBlock, 1); dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); - MatMulFloatInt4RouterKernel<<>>( - output, a_data, b_data_quant, scales_data, bias_data); + if (block_size == 32) { + MatMulFloatInt4RouterKernel<<>>( + output, a_data, b_data_quant, scales_data, bias_data, n, k); + } else { + MatMulFloatInt4RouterKernel<<>>( + output, a_data, b_data_quant, scales_data, bias_data, n, k); + } return true; } From 4f64dc597ab583de3a67e0b8e2d358e568867266 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 20 Jun 2026 17:42:36 -0700 Subject: [PATCH 07/10] add doc --- docs/contrib_ops/cuda/matmul_nbits.md | 258 ++++++++++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 docs/contrib_ops/cuda/matmul_nbits.md diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md new file mode 100644 index 0000000000000..de3056bb36cf1 --- /dev/null +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -0,0 +1,258 @@ +# MatMulNBits — CUDA Operator Documentation + +This document describes the CUDA execution-provider implementation of the +**MatMulNBits** (`com.microsoft::MatMulNBits`) operator: its kernel dispatch +chain, the fast / fallback / specialized kernels, the weight format they expect, +and the environment variables that control routing. + +MatMulNBits computes `Y = A · dequant(B)ᵀ (+ bias)` where `B` is an `N × K` +weight matrix quantized to 4 or 8 bits with block-wise (group) scales and +optional zero points. It is the building block for weight-only quantized linear +layers (including MoE routers and LM heads). + +Source files: + +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc) — operator, `ComputeInternal`, dispatch chain, dequant+GEMM fallback. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h) — kernel class, constructor-time configuration, environment-variable parsing. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh) — `TryMatMulNBits` fast-path entry and bias-add launcher. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu](onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu) — 4-bit fast GEMV kernels (generic + router specialization). + +--- + +## Table of Contents + +1. [Operator Schema](#1-operator-schema) +2. [Weight Format](#2-weight-format) +3. [Dispatch Chain](#3-dispatch-chain) +4. [Fast Path — Fused GEMV](#4-fast-path--fused-gemv) + - [4.1 Generic 4-bit GEMV kernel](#41-generic-4-bit-gemv-kernel) + - [4.2 Router GEMV specialization](#42-router-gemv-specialization) +5. [Fallback Path — Dequantize + GEMM](#5-fallback-path--dequantize--gemm) +6. [fpA_intB_gemm Path (CUTLASS weight-only)](#6-fpa_intb_gemm-path-cutlass-weight-only) +7. [Bias Handling](#7-bias-handling) +8. [Environment Variables](#8-environment-variables) +9. [Testing](#9-testing) + +--- + +## 1. Operator Schema + +| Attribute | Meaning | +|-----------|---------| +| `K` | Input feature dimension (columns of `A`, columns of the logical `B`). | +| `N` | Output feature dimension (rows of the logical `B`). | +| `bits` | Quantization bit width: `4` or `8`. | +| `block_size` | Quantization group size along `K` (16 / 32 / 64 / 128). One scale (and optional zero point) per group. | + +| Input | Index | Notes | +|-------|-------|-------| +| `A` | 0 | Activations, FP16 / BF16 / FP32. Shape `[M, K]`. | +| `B` | 1 | Packed 4/8-bit weights. | +| `scales` | 2 | Per-group scales, same element type as `A`. | +| `zero_points` | 3 | Optional. Packed integer (symmetric default) **or** same type as `A`. | +| `g_idx` / `reorder_idx` | 4 | Optional group/reorder index (act-order). | +| `bias` | 5 | Optional `[N]` bias added to the output. | + +`M` is the (flattened) token count: `M = 1` is the decode / GEMV case that the +fast kernels target. + +--- + +## 2. Weight Format + +For 4-bit, `B` is stored as `[N, ceil(K / block_size), block_size / 2]` bytes: +each expert/output row `n` is a contiguous run of `K/2` bytes (two 4-bit weights +per byte), preceded conceptually by `K / block_size` scales in the `scales` +tensor. Quantization is **column-wise block** by default +(`column_wise_quant_blk_ = true`); row-wise layouts interleave `K` blocks across +`N` and cannot be sliced along `N` (this disables the chunked fallback). + +Symmetric 4-bit weights store values `0..15` that dequantize to `(q − 8) · +scale`; the fast kernels hard-code the zero point of `8` when no `zero_points` +input is present. + +--- + +## 3. Dispatch Chain + +`MatMulNBits::ComputeInternal` tries the cheapest applicable path first and +falls through to progressively more general ones: + +```mermaid +flowchart TD + A[ComputeInternal] --> F{has_fpA_intB_gemm_?
FP16/BF16, prepacked,
block 64/128, sm>=75} + F -- yes --> FP[fpA_intB CUDA GEMV
or CUTLASS grouped GEMM] --> R[return] + F -- no --> G{reorder_idx == null
and zero_points not typed-T?} + G -- no --> DQ + G -- yes --> T1[TryMatMulNBits
fused fast GEMV] + T1 -- success --> R + T1 -- fail due to bias --> T2[TryMatMulNBits without bias] + T2 -- success --> BIAS[MatMulNBitsBiasAdd] --> R + T2 -- fail --> DQ + DQ[Dequantize blockwise 4b/8b] --> GEMM[cuBLAS GEMM
full-N or chunked along N] + GEMM --> BIAS2[optional bias add] --> R +``` + +The fast fused path is only attempted when there is **no** `reorder_idx` and the +`zero_points` (if any) are packed integers (not the same element type as `A`). + +--- + +## 4. Fast Path — Fused GEMV + +`TryMatMulNBits` ([matmul_nbits.cuh](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh)) +dispatches by bit width: + +- `bits == 8` → `TryMatMul8Bits` (no bias support; returns `false` if bias set). +- `bits == 4` → `TryMatMul4Bits`. + +`TryMatMul4Bits` ([matmul_4bits.cu](onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu)) +first applies a guard common to all fused kernels: + +``` +n % kColsPerThreadBlock (8) == 0 and k % 8 == 0 and m <= 1 +``` + +i.e. the fast path is a **GEMV** (single token). If the guard passes it then +chooses between the router specialization (§4.2) and the generic kernel (§4.1). + +### 4.1 Generic 4-bit GEMV kernel + +`MatMulFloatInt4Kernel`: + +- **Launch:** grid `(ceil(N / 8), M)`, block `(warpSize, 8)` — one **warp per + output column** (`kColsPerThreadBlock = 8` warps per block). +- **Scales / zero points** are staged into shared memory once per thread block + (hence the `shared_mem_size > sharedMemPerBlock` bail-out in + `TryMatMul4Bits`). +- **Inner loop:** each lane consumes `kElementsPerThreadPerIteration = 8` + weights per step; a warp covers `warpSize × 8 = 256` `K` elements per + iteration via `AccumulateEightElements4b` (a `prmt` + `lop3`-based int4→half2 + conversion). A 3-tier macro unroll (`×16`, `×4`, `×1`) plus a scalar remainder + step walks all of `K`, followed by a warp-shuffle reduction. +- **Supported `block_size`:** 16 / 32 / 64 / 128 (others throw). +- **Bias:** not supported — the kernel returns `false` to `TryMatMul4Bits` when + `bias != nullptr` (see §7 for how bias is then handled). + +### 4.2 Router GEMV specialization + +`MatMulFloatInt4RouterKernel` is a specialization for MoE-router +GEMVs (`output(1, N) = A(1, K) · dequant(B(N, K)) + bias(N)`). It is selected by +`IsSupportedRouterGemvShape` when: + +- `zero_points == nullptr` (symmetric), `M == 1`, +- `block_size ∈ {32, 64}` and `K % block_size == 0`, +- the `(N, K)` pair matches a known router shape: + +| Model | `N` (experts) | `K` (hidden size) | +|-------|---------------|-------------------| +| gpt-oss-20b | 32 | 2880 | +| Qwen3.6-35B-A3B | 256 | 2048 | +| gemma-4-26B-A4B | 128 | 2112 | + +Design notes: + +- **One warp per expert column**, `kColsPerThreadBlock = 8` warps per block; + grid is `(N / 8, 1)`. `N` is supplied at runtime via the grid and `K` at + runtime as a kernel argument, so a single instantiation per `(T, BlockSize)` + serves every router shape. Only `BlockSize` is a template parameter because it + drives the scale stride. +- **No shared memory:** scales are read directly from global memory (and are + L2-resident at these tiny sizes), avoiding the staging `__syncthreads` of the + generic kernel. +- **Bias is fused:** lane 0 adds `bias[n_id]` before the store. +- **Group size 32 vs 64:** both are supported because `kPerIter = 256` is + divisible by each, keeping the per-iteration scale stride exact. 64 halves the + number of scale loads relative to 32; pick whichever quantization granularity + gives the best accuracy/latency trade-off for the model. (Per-row / per-column + quantization, i.e. one scale per expert, is **not** a supported MatMulNBits + layout and is intentionally excluded.) +- The same correctness invariants as the generic kernel apply (3-tier unroll + + remainder, warp-shuffle reduction). + +To add another router, extend `IsSupportedRouterGemvShape` with its `(N, K)`; +no kernel change is required as long as `N % 8 == 0` and `K % block_size == 0`. + +--- + +## 5. Fallback Path — Dequantize + GEMM + +When neither the fpA_intB path nor the fused GEMV applies (e.g. `M > 1`, +`reorder_idx` present, typed zero points, or an unsupported shape), the operator +dequantizes `B` into a scratch buffer and runs a dense cuBLAS GEMM: + +- `DequantizeBlockwise4b` / `DequantizeBlockwise8b` (or the column-wise helper) + expands packed weights to `T` into a `N × K_padded` scratch buffer, where + `K_padded = ceil(K / block_size) · block_size`. +- A single cuBLAS `GEMM` (`transb = true`) then produces `Y`. + +**Chunked variant.** For large `N`, materializing the full `N × K_padded` +dequantized matrix can dominate device memory. The implementation slices the +dequant+GEMM along `N` into chunks when: + +- column-wise quantization and no `reorder_idx`, **and** +- `force_chunked_` is set, **or** scratch `> 256 MB` and `N > 2 × + chunk_target_rows`. + +`chunk_target_rows` defaults to 32768 (configurable, see §8) — chosen so each +tile saturates the SMs while keeping scratch ≲128 MB. + +--- + +## 6. fpA_intB_gemm Path (CUTLASS weight-only) + +When built with `USE_FPA_INTB_GEMM` and enabled via `ORT_FPA_INTB_GEMM`, FP16/BF16 +MatMulNBits can use the TensorRT-LLM-derived CUTLASS weight-only kernels. The +constructor sets `has_fpA_intB_gemm_` only when: + +- dtype is FP16 or BF16, `bits ∈ {4, 8}`, `block_size ∈ {64, 128}`, +- no `g_idx`, no `bias`, `N % (bits==8 ? 32 : 64) == 0`, `K % block_size == 0`, +- `sm_ >= 75`, and weight/scale/zero-point inputs are **prepacked** initializers. + +At run time a profiler picks the best tactic; small `M` may use a dedicated CUDA +GEMV kernel (`bestTactic->enableCudaKernel`), otherwise a CUTLASS grouped GEMM. +This path takes precedence over everything in §3 when active. + +--- + +## 7. Bias Handling + +Only the router specialization (§4.2) fuses bias inside the GEMV. For every other +fast-path shape, `TryMatMul4Bits` / `TryMatMul8Bits` return `false` when bias is +present. `ComputeInternal` then: + +1. Retries `TryMatMulNBits` with `bias = nullptr`; on success it adds the bias + with a separate `MatMulNBitsBiasAdd` kernel (`LaunchMatMulNBitsBiasAdd`, + accumulating in float for half/bfloat16 accuracy). +2. If the fast path still does not apply, falls through to the dequant+GEMM + fallback (§5), which ignores bias, followed by the same bias-add kernel. + +--- + +## 8. Environment Variables + +| Variable | Type / default | Effect | +|----------|----------------|--------| +| `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION` | bool, `0` | Disable the router GEMV specialization (§4.2); shapes fall back to the generic GEMV / dequant path. Useful for A/B benchmarking. | +| `ORT_FPA_INTB_GEMM` | int bitmask, `0` | Enable the CUTLASS weight-only path (§6). `0x01` = all, `0x02` = CUDA GEMV, `0x04` = int4, `0x08` = int8. `0` disables it. | +| `ORT_MATMULNBITS_FORCE_CHUNKED` | int, `0` | Force the chunked dequant+GEMM fallback (§5) regardless of the size heuristic. | +| `ORT_MATMULNBITS_CHUNK_SIZE` | int64, `32768` | Target rows per chunk in the chunked fallback. Values `< 1` reset to the default. | + +> Environment variables are read with ORT's cross-platform +> `ParseEnvironmentVariableWithDefault` helper (safe on Windows), not +> `std::getenv`. + +--- + +## 9. Testing + +- Python operator tests: `onnxruntime/test/python/transformers` (see the QMoE / + GEMV profiling helpers, e.g. `profile_qmoe_gemv.sh`). +- GEMV profiling baselines and methodology are recorded in + [qmoe_gemv_experiments.md](qmoe_gemv_experiments.md). +- To compare the router specialization against the generic path, run the same + model with and without `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1`. + +After editing any `.cu` kernel, rebuild the CUDA provider +(`ninja onnxruntime_providers_cuda`) and re-run the relevant tests; note the +nvcc incremental-build caveats in the repository build notes. From fbff53a3f6d07952de8f740baa9a883ca03091f3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 22 Jun 2026 17:17:40 -0700 Subject: [PATCH 08/10] Address review: restrict router GEMV to GPT-OSS, cap bias-add grid, refactor unroll - IsSupportedRouterGemvShape now matches only the GPT-OSS-20B router shape (M=1, N=32, K=2880), matching the PR's stated exact-shape scope; remove untested Qwen3/Gemma shapes. - LaunchMatMulNBitsBiasAdd caps the grid at the CUDA gridDim.x limit and the kernel uses a grid-stride loop, avoiding silent truncation when casting to unsigned int. - Replace the RouterUnRollReduction macro with a templated __device__ RouterUnrollReduction helper. - Mark router dispatch dim3 blocks/threads const. --- .../cuda/quantization/matmul_4bits.cu | 93 +++++++++++-------- 1 file changed, 52 insertions(+), 41 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu index 848d78925b09c..b24e445401b99 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -25,13 +25,10 @@ namespace cuda { // The specialized router GEMV kernel only handles M=1, or batch size 1. constexpr int kRouterM = 1; -// MoE router shapes (N = number of experts, K = hidden size) +// MoE router shape (N = number of experts, K = hidden size). The specialization is exact-shape +// gated to the GPT-OSS-20B router projection to keep the dispatch change conservative. constexpr int kGptOssRouterN = 32; constexpr int kGptOssRouterK = 2880; -constexpr int kQwen3RouterN = 256; -constexpr int kQwen3RouterK = 2048; -constexpr int kGemmaRouterN = 128; -constexpr int kGemmaRouterK = 2112; static bool IsRouterGemvSpecializationDisabled() { // Use ORT's cross-platform env var helper instead of std::getenv, which is unsafe on Windows. @@ -39,10 +36,11 @@ static bool IsRouterGemvSpecializationDisabled() { } // The router GEMV kernel handles any symmetric (no zero point) M=1 shape with an int4 group size of -// 32 or 64 (whichever quantizes best) and N divisible by kColsPerThreadBlock. We gate on the known -// MoE router shapes to avoid changing the dispatch for general MatMulNBits cases. K must be a multiple -// of the group size (always true for a router) and N a multiple of kColsPerThreadBlock (checked in -// TryMatMul4Bits). Note kPerIter (256) is divisible by both 32 and 64, so the scale stride is exact. +// 32 or 64 (whichever quantizes best) and N divisible by kColsPerThreadBlock. We gate on the exact +// GPT-OSS-20B router shape to avoid changing the dispatch for general MatMulNBits cases. K must be a +// multiple of the group size (always true for a router) and N a multiple of kColsPerThreadBlock +// (checked in TryMatMul4Bits). Note kPerIter (256) is divisible by both 32 and 64, so the scale +// stride is exact. static bool IsSupportedRouterGemvShape(const uint8_t* zero_points, int m, int n, int k, int block_size) { if (zero_points != nullptr || m != kRouterM || (block_size != 32 && block_size != 64)) { return false; @@ -50,9 +48,7 @@ static bool IsSupportedRouterGemvShape(const uint8_t* zero_points, int m, int n, if (k % block_size != 0) { return false; } - return (n == kGptOssRouterN && k == kGptOssRouterK) || // gpt-oss-20b - (n == kQwen3RouterN && k == kQwen3RouterK) || // Qwen3.6-35B-A3B - (n == kGemmaRouterN && k == kGemmaRouterK); // gemma-4-26B-A4B + return n == kGptOssRouterN && k == kGptOssRouterK; // gpt-oss-20b } template @@ -364,6 +360,33 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt } } // namespace cuda +// Reduces kUnroll groups of kPerIter elements per step, advancing the packed-weight pointer, scale +// index and k position in lockstep. Factored out of MatMulFloatInt4RouterKernel so the three unroll +// factors (16, 4, 1) share one implementation instead of a macro. +template +__device__ __forceinline__ void RouterUnrollReduction( + const uint8_t*& b_data_quant, + const T* scales_data, + const T* a_data, + int k, + int& k_id, + int& scale_id, + T (&sums)[8]) { + constexpr int kPerIter = kWarpSize * kElementsPerThreadPerIteration; + constexpr int kUnrollStep = kUnroll * kPerIter; + const int k_unroll_bound = k - k % kUnrollStep; + for (; k_id < k_unroll_bound; k_id += kUnrollStep) { +#pragma unroll + for (int i = 0; i < kUnroll; i++) { + uint32_t value = *(reinterpret_cast(b_data_quant + kPerIter / 2 * i)); + T scale = scales_data[scale_id + kPerIter / BlockSize * i]; + AccumulateEightElements4b(value, scale, 8, a_data + k_id + i * kPerIter, sums); + } + b_data_quant += kPerIter / 2 * kUnroll; + scale_id += kPerIter / BlockSize * kUnroll; + } +} + // GEMV specialization for MoE routers: output(1, N) = a(1, K) x dequant(B(N, K)) [+ bias(N)]. // B is 4-bit block-quantized (symmetric, no zero point) with group size BlockSize (32 or 64). One warp // computes one expert column; the thread block holds kColsPerThreadBlock warps. N is passed via the @@ -395,26 +418,9 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloatInt int k_id = 0; int scale_id = lane_id * 8 / BlockSize; -#define RouterUnRollReduction(unroll_size) \ - do { \ - constexpr int kUnroll = unroll_size; \ - constexpr int kUnrollStep = kUnroll * kPerIter; \ - const int k_unroll_bound = k - k % kUnrollStep; \ - for (; k_id < k_unroll_bound; k_id += kUnrollStep) { \ - _Pragma("unroll") for (int i = 0; i < kUnroll; i++) { \ - uint32_t value = *(reinterpret_cast(b_data_quant + kPerIter / 2 * i)); \ - T scale = scales_data[scale_id + kPerIter / BlockSize * i]; \ - AccumulateEightElements4b(value, scale, 8, a_data + k_id + i * kPerIter, sums); \ - } \ - b_data_quant += kPerIter / 2 * kUnroll; \ - scale_id += kPerIter / BlockSize * kUnroll; \ - } \ - } while (false) - - RouterUnRollReduction(16); - RouterUnRollReduction(4); - RouterUnRollReduction(1); -#undef RouterUnRollReduction + RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); + RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); + RouterUnrollReduction(b_data_quant, scales_data, a_data, k, k_id, scale_id, sums); if (k_id + lane_id * 8 < k) { uint32_t value = *(reinterpret_cast(b_data_quant)); @@ -455,8 +461,8 @@ bool TryMatMul4Bits( if (IsSupportedRouterGemvShape(zero_points, m, n, k, block_size) && !IsRouterGemvSpecializationDisabled()) { - dim3 blocks(n / kColsPerThreadBlock, 1); - dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + const dim3 blocks(n / kColsPerThreadBlock, 1); + const dim3 threads(GPU_WARP_SIZE_HOST, kColsPerThreadBlock); if (block_size == 32) { MatMulFloatInt4RouterKernel<<>>( output, a_data, b_data_quant, scales_data, bias_data, n, k); @@ -550,13 +556,14 @@ template bool TryMatMul4Bits( template __global__ void MatMulNBitsBiasAddKernel(T* output, const T* bias_data, int n, int64_t total) { - const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx >= total) { - return; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < total; + idx += stride) { + const int col = static_cast(idx % n); + // Accumulate in float to stay accurate for half/bfloat16. + output[idx] = static_cast(static_cast(output[idx]) + static_cast(bias_data[col])); } - const int col = static_cast(idx % n); - // Accumulate in float to stay accurate for half/bfloat16. - output[idx] = static_cast(static_cast(output[idx]) + static_cast(bias_data[col])); } template @@ -566,7 +573,11 @@ void LaunchMatMulNBitsBiasAdd(T* output, const T* bias_data, int m, int n, cudaS return; } constexpr int kThreadsPerBlock = 256; - const int64_t num_blocks = (total + kThreadsPerBlock - 1) / kThreadsPerBlock; + // Cap the grid at the CUDA gridDim.x limit (2^31 - 1); the kernel uses a grid-stride loop, so a + // capped grid still covers every element without truncating the launch size. + constexpr int64_t kMaxGridBlocks = 2147483647; + const int64_t blocks = (total + kThreadsPerBlock - 1) / kThreadsPerBlock; + const int64_t num_blocks = blocks < kMaxGridBlocks ? blocks : kMaxGridBlocks; MatMulNBitsBiasAddKernel<<(num_blocks), kThreadsPerBlock, 0, stream>>>( output, bias_data, n, total); } From 4d2bc950dc83c7ecdaf41a70697d6338ae8556e2 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 22 Jun 2026 23:57:33 -0700 Subject: [PATCH 09/10] docs: fix relative source links in matmul_nbits.md --- docs/contrib_ops/cuda/matmul_nbits.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index de3056bb36cf1..7e87869f12a70 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -12,10 +12,10 @@ layers (including MoE routers and LM heads). Source files: -- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc) — operator, `ComputeInternal`, dispatch chain, dequant+GEMM fallback. -- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h) — kernel class, constructor-time configuration, environment-variable parsing. -- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh) — `TryMatMulNBits` fast-path entry and bias-add launcher. -- [onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu](onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu) — 4-bit fast GEMV kernels (generic + router specialization). +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc) — operator, `ComputeInternal`, dispatch chain, dequant+GEMM fallback. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h) — kernel class, constructor-time configuration, environment-variable parsing. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh) — `TryMatMulNBits` fast-path entry and bias-add launcher. +- [onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu) — 4-bit fast GEMV kernels (generic + router specialization). --- @@ -100,13 +100,13 @@ The fast fused path is only attempted when there is **no** `reorder_idx` and the ## 4. Fast Path — Fused GEMV -`TryMatMulNBits` ([matmul_nbits.cuh](onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh)) +`TryMatMulNBits` ([matmul_nbits.cuh](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh)) dispatches by bit width: - `bits == 8` → `TryMatMul8Bits` (no bias support; returns `false` if bias set). - `bits == 4` → `TryMatMul4Bits`. -`TryMatMul4Bits` ([matmul_4bits.cu](onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu)) +`TryMatMul4Bits` ([matmul_4bits.cu](../../../onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu)) first applies a guard common to all fused kernels: ``` From 9b23007d598970462d25bf5dd8d7e78e752b6bc3 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 23 Jun 2026 00:02:39 -0700 Subject: [PATCH 10/10] docs+test: align router GEMV scope with code; cover block_size 64 --- docs/contrib_ops/cuda/matmul_nbits.md | 4 +- .../contrib_ops/cuda/qmoe_gemv_experiments.md | 16 ++--- .../test/contrib_ops/matmul_4bits_test.cc | 58 ++++++++++--------- 3 files changed, 41 insertions(+), 37 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md index 7e87869f12a70..51b8fa55fcad9 100644 --- a/docs/contrib_ops/cuda/matmul_nbits.md +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -142,13 +142,11 @@ GEMVs (`output(1, N) = A(1, K) · dequant(B(N, K)) + bias(N)`). It is selected b - `zero_points == nullptr` (symmetric), `M == 1`, - `block_size ∈ {32, 64}` and `K % block_size == 0`, -- the `(N, K)` pair matches a known router shape: +- the `(N, K)` pair matches the exact-gated router shape: | Model | `N` (experts) | `K` (hidden size) | |-------|---------------|-------------------| | gpt-oss-20b | 32 | 2880 | -| Qwen3.6-35B-A3B | 256 | 2048 | -| gemma-4-26B-A4B | 128 | 2112 | Design notes: diff --git a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md index 7eb9be55b679e..2e570c9c4fc1a 100644 --- a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md +++ b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md @@ -1163,9 +1163,11 @@ Result: `1 test passed`. router shape with no zero-points or `g_idx` input. - The exact router GEMV kernel now accepts an optional bias pointer and adds the `[32]` bias before storing the router logits. -- `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` disables only this CUDA router bias - fusion. `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1` still disables both - the exact router GEMV specialization and this bias fusion. +- The only runtime control is `ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION=1`, + which disables the exact router GEMV specialization; the fused bias then falls + back to a separate bias-add in the generic path. The graph-level + `MatMulNBits + Add` rewrite itself is performed by the `MatMulNBitsFusion` + transformer and has no dedicated opt-out env var. ### Graph Validation @@ -1183,10 +1185,10 @@ python | Mode | Total `Add` nodes | Router `MatMulNBits` nodes with bias | |------|------------------:|--------------------------------------:| | enabled | 48 | 24 | -| `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` | 72 | 0 | +| router bias fusion disabled | 72 | 0 | -This confirms all 24 real GPT-OSS router bias Adds are fused and the opt-out -restores the unfused graph. +This confirms all 24 real GPT-OSS router bias Adds are fused and disabling the +fusion restores the unfused graph. ### Model-Level Decode Benchmark With CUDA Graph @@ -1205,7 +1207,7 @@ bash scripts/h200_18/bench_gpt_oss_ort_decode.sh | Mode | Decode latency ms/token | Decode throughput tok/s | |------|-------------------------:|-------------------------:| | enabled | 2.822998 | 354.233302 | -| `ORT_DISABLE_QMOE_ROUTER_BIAS_FUSION=1` | 2.829066 | 353.473500 | +| router bias fusion disabled | 2.829066 | 353.473500 | ### Decision diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 812fb16291550..07b275b813aa7 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -808,36 +808,40 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint) { TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { constexpr float abs_error = 0.1f; - TestOptions opts{}; - opts.M = 1, opts.N = 32, opts.K = 2880; - opts.block_size = 32; - opts.has_zero_point = false; - opts.zp_is_4bit = true; - opts.output_abs_error = abs_error; - opts.output_rel_error = 0.02f; + // IsSupportedRouterGemvShape enables the specialization for block_size 32 and 64, so exercise both + // to keep the MatMulFloatInt4RouterKernel and instantiations covered. + for (auto block_size : {32, 64}) { + TestOptions opts{}; + opts.M = 1, opts.N = 32, opts.K = 2880; + opts.block_size = block_size; + opts.has_zero_point = false; + opts.zp_is_4bit = true; + opts.output_abs_error = abs_error; + opts.output_rel_error = 0.02f; - { - ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; - std::vector> eps; - eps.push_back(DefaultCudaExecutionProvider()); - RunTest(opts, std::move(eps)); - } + { + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } - { - ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", "1"}}}; - std::vector> eps; - eps.push_back(DefaultCudaExecutionProvider()); - RunTest(opts, std::move(eps)); - } + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", "1"}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } - { - opts.has_bias = true; - ScopedEnvironmentVariables scoped_env_vars{ - EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; - std::vector> eps; - eps.push_back(DefaultCudaExecutionProvider()); - RunTest(opts, std::move(eps)); + { + opts.has_bias = true; + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{"ORT_DISABLE_QMOE_ROUTER_GEMV_SPECIALIZATION", std::optional{}}}}; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } } }