diff --git a/docs/contrib_ops/cuda/matmul_nbits.md b/docs/contrib_ops/cuda/matmul_nbits.md new file mode 100644 index 0000000000000..51b8fa55fcad9 --- /dev/null +++ b/docs/contrib_ops/cuda/matmul_nbits.md @@ -0,0 +1,256 @@ +# 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 the exact-gated router shape: + +| Model | `N` (experts) | `K` (hidden size) | +|-------|---------------|-------------------| +| gpt-oss-20b | 32 | 2880 | + +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. diff --git a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md index d2165640b0051..2e570c9c4fc1a 100644 --- a/docs/contrib_ops/cuda/qmoe_gemv_experiments.md +++ b/docs/contrib_ops/cuda/qmoe_gemv_experiments.md @@ -1031,3 +1031,208 @@ 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. +- 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 + +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 | +| router bias fusion disabled | 72 | 0 | + +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 + +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 | +| router bias fusion disabled | 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..b24e445401b99 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits.cu @@ -5,10 +5,16 @@ #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; @@ -16,6 +22,35 @@ namespace onnxruntime { namespace contrib { namespace cuda { +// The specialized router GEMV kernel only handles M=1, or batch size 1. +constexpr int kRouterM = 1; + +// 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; + +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); +} + +// 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 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; + } + if (k % block_size != 0) { + return false; + } + return n == kGptOssRouterN && k == kGptOssRouterK; // gpt-oss-20b +} + template __device__ __forceinline__ T WarpUniform(T value) { struct { @@ -325,6 +360,87 @@ __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 +// 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, + 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 += 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 / BlockSize; + + 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)); + 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 +448,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 +458,25 @@ bool TryMatMul4Bits( if (n % kColsPerThreadBlock != 0 || k % 8 != 0 || m > 1) { return false; } + + if (IsSupportedRouterGemvShape(zero_points, m, n, k, block_size) && + !IsRouterGemvSpecializationDisabled()) { + 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); + } else { + MatMulFloatInt4RouterKernel<<>>( + output, a_data, b_data_quant, scales_data, bias_data, n, k); + } + 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 +518,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 +532,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 +546,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, @@ -415,6 +554,43 @@ 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 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])); + } +} + +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; + // 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); +} + +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 cd5358c41e898..896291ce01b5d 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_)); @@ -376,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()), @@ -383,16 +381,48 @@ Status MatMulNBits::ComputeInternal(OpKernelContext* ctx) const { blob_data, reinterpret_cast(scales_data), static_cast(zero_points_data), + reinterpret_cast(bias_data), + m, + n, + k, + SafeInt(block_size_), + GetDeviceProp().sharedMemPerBlock, + stream)) { + return Status::OK(); + } + + // 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. @@ -577,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 e97990d884ae5..61eb62fef4a2b 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,18 +53,31 @@ 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); } 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/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/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..07b275b813aa7 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -805,6 +805,46 @@ TEST(MatMulNBits, Fp16_Int4_NoZeroPoint) { } } +TEST(MatMulNBits, Fp16_Int4_GptOssRouterShapeNoZeroPoint) { + constexpr float abs_error = 0.1f; + + // 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", "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)); + } + } +} + 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..1f07bf3204894 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 @@ -10792,26 +10793,32 @@ TEST_F(GraphTransformationTests, MatMulNBitsBiasFusion) { 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, - 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"); @@ -10820,8 +10827,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 +10842,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 +10852,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,6 +10892,14 @@ 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); + + // 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); } } }