[CUDA] Add 2-bit MatMulNBits support

[Tianlei Wu (tianleiwu)]

Description

Add CUDA support for affine 2-bit MatMulNBits weights. The change adds 2-bit dequantization and specialized GEMV/batched kernels, dispatches fp16, fp32, and bf16 inputs, and makes packed weight and zero-point offsets depend on the configured bit width.

The exported Bonsai graph contract was verified directly: 401 MatMulNBits nodes use bits=2, block_size=128, fp16 activations, packed uint8 affine zero points, and no g_idx input. CUDA supports 2-bit block sizes from 16 through 256, matching MLAS, and rejects g_idx explicitly.

Validation on CUDA 13.0 / A100:

onnxruntime_provider_test --gtest_filter='MatMul2BitsCuda.*:*GatherBlockQuantized*Cuda*'
[==========] 24 tests from 2 test suites ran.
[  PASSED  ] 24 tests.

This includes all 17 MatMul2BitsCuda tests, typed and packed zero-point fallback coverage, forced chunked fallback with and without packed zero points, validation failures for unsupported g_idx and oversized blocks, and the production QKV dimensions at M=1 and the fallback boundary M=32, N=10240, K=5120, block_size=128.

Motivation and Context

Models with native ternary weights export four 2-bit values per byte through the existing MatMulNBits affine ABI. CUDA previously accepted this schema but lacked a correct 2-bit execution path. This change implements that path and rejects unsupported widths and input combinations instead of interpreting them with an incompatible packed stride.

[Copilot]

🟡 Changes recommended

The implementation has SM52 build compatibility and reorder correctness defects, while CPU fallback and missing chunked coverage weaken the CUDA tests.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds CUDA support for affine 2-bit MatMulNBits, including fused small-M kernels and generic dequantization fallback.

Changes:

• Adds FP16, FP32, and BF16 2-bit GEMV/batched kernels.
• Extends dequantization, packed offsets, and bit-width validation.
• Adds CUDA correctness, ternary-weight, and production-shape tests.

File summaries

File	Description
onnxruntime/test/contrib_ops/matmul_2bits_test.cc	Adds CUDA 2-bit tests.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h	Validates supported bit widths.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh	Dispatches 2-bit kernels.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc	Adds fallback dequantization and offset handling.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu	Implements top-level 2-bit dispatch.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1.cuh	Declares the single-row kernel API.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh	Implements single-row GEMV.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_half.cu	Instantiates FP16 GEMV.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_float.cu	Instantiates FP32 GEMV.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh	Adds shared unpacking and accumulation primitives.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_bfloat16.cu	Instantiates BF16 GEMV.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched.cuh	Declares batched kernel dispatch.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh	Implements small-M batched kernels.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_half.cu	Instantiates batched FP16 support.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_float.cu	Instantiates batched FP32 support.
onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_bfloat16.cu	Instantiates batched BF16 support.
onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh	Exposes 2-bit dequantization dispatch.
onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu	Implements generic 2-bit dequantization.

Review details

Suppressed comments (2)

onnxruntime/test/contrib_ops/matmul_2bits_test.cc:1838

• This direct CUDA test also leaves CPU fallback enabled, so it can validate the CPU implementation instead of the new CUDA path if CUDA does not claim the node. Disable CPU EP fallback in the session options before running it.

  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
  execution_providers.emplace_back(DefaultCudaExecutionProvider());
  test.ConfigEps(std::move(execution_providers));
  test.RunWithConfig();

onnxruntime/test/contrib_ops/matmul_2bits_test.cc:1894

• This negative test can pass through CPU fallback because the CPU kernel also rejects bits=3 with an error containing bits; it therefore does not prove that the new CUDA constructor guard ran. Disable CPU EP fallback before asserting the CUDA-specific failure.

  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
  execution_providers.emplace_back(DefaultCudaExecutionProvider());
  test.Config(OpTester::ExpectResult::kExpectFailure, "bits")
      .ConfigEps(std::move(execution_providers))
      .RunWithConfig();

• Files reviewed: 18/18 changed files
• Comments generated: 7
• Review effort level: Balanced

---

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

[Akshay Sonawane (apsonawane)]

Can we update matmul_nbits.md file as well

[Tianlei Wu (tianleiwu)]

Akshay Sonawane (@apsonawane) Updated docs/contrib_ops/cuda/matmul_nbits.md with the 2-bit CUDA support details, the MLAS-aligned 16–256 block-size range, unsupported g_idx, and the relevant implementation/test references.

[Copilot]

Copilot review overview

🟡 Changes recommended

The updated CUDA kernel test references an undeclared sm_ identifier and therefore fails to compile in fpA_intB builds.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1 · 1

Open (3)

• Use device_prop_ instead of undeclared sm_ · New
• Add CUDA coverage for same-type typed zero points · New
• Align documented block range with the 256-block limit · New

Resolved since last review (7)

• These unguarded sub.f16x2/__hfma2 instructions break the SM52 compilation pass used by the…
• When reorder_idx is present with typed zero points, ZeroT equals T, so this condition…
• CheckInputs accepts any power-of-two block size >=16, but for block_size > 16384 this integer…
• The host-side reference for these four shapes executes 41,995,468,800 scalar inner-loop iterations;…
• These CUDA-specific cases do not disable CPU EP fallback. ORT implicitly adds CPU EP, and CPU…
• No added test exercises this new 2-bit chunk offset/ZP-stride path: the largest new case has about…
• The validation command reported in the PR description targets onnxruntime_test_all, but…