🔥 Jadwal & Hasil: microsoft/onnxruntime/pull/

Tianlei Wu (tianleiwu) · 2026-09-18T19:05:04Z

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) · 2026-09-18T20:30:37Z

Can we update matmul_nbits.md file as well

Tianlei Wu (tianleiwu) · 2026-09-19T00:05:11Z

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…

Copilot

Copilot review overview

🔵 Needs a closer look

The new architecture-sensitive CUDA kernels warrant final human review despite broad test coverage.

Review effort: Balanced
Findings: None

Resolved since last review (3)

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

Previously missed (1)

In code that hasn't changed since last review

Fix indentation of entries in section 4 nested list

docs/contrib_ops/cuda/matmul_nbits.md:33

These three entries are indented by only two spaces, unlike the adjacent 4.1 entry, so Markdown renders them outside the section 4 nested list. Align their indentation to keep the table of contents hierarchy intact.

## Summary - Add canonical 2-bit `MatMulNBits` weight packing to the model builder. - Allow DFlash 2 body weights to use `int2` and select FP16 or BF16 compute through the structured `drafter_options.quant_config` schema introduced by #2592. - Support either offline SM80 prepacking through `format.matmulnbits_weights_prepacked` or CUDA runtime prepacking through `runtime_config.model.dflash2.session_options["ep.cuda.fpa_intb_gemm"]`. - Keep the target model and adopted embedding/LM-head tensors at their independently configured precision and block size. This PR is stacked on #2592 and should be reviewed as `171e3ea5...HEAD`. ## Configuration ```json { "builder_config_version": 2, "target_options": { "attention": {"implementation": "paged"} }, "drafter_options": { "drafter_type": "dflash2", "path": "Qwen3.8-27B-DFlash2", "num_draft_tokens": 7, "quant_config": { "io_dtype": "bf16", "weights": {"type": "int2", "block_size": 64} } }, "runtime_config": { "model": { "dflash2": { "session_options": {"ep.cuda.fpa_intb_gemm": "1"} } }, "speculative": {"max_draft_tokens": 7} } } ``` The DFlash body can use the standard raw QOperator layout for runtime prepacking or an offline SM80-prepacked layout. INT2 prepacking requires block size 64 or 128; BF16 x INT2 also requires a full fpA-intB kernel build. ## Runtime Dependencies - microsoft/onnxruntime#32693: base 2-bit `MatMulNBits` support. - microsoft/onnxruntime#32699: INT2 fpA-intB GEMM/GEMV support, including BF16 in the full kernel build. ## Validation - Configuration/precision: 180 passed. - DFlash/quantization: 190 passed, 3 skipped. - Full structured-config export of Qwen3.8-27B with 401 target INT4 block-32 nodes, 46 DFlash INT2 block-64 nodes, and one adopted INT4 block-32 LM head. - Debug-only dispatch check: all 230 observed INT2 calls used `runtime-prepacked(SM80 layout)` with `packing_sm=80`; no fallback/error lines. - Clean timing runs explicitly had `ORT_FPA_INTB_DEBUG` absent. | Context | INT2 block-64 fpA | Same artifact, fpA off | fpA gain | INT4 control | | ---: | ---: | ---: | ---: | ---: | | 2,048 | 289.8 tok/s | 226.4 tok/s | +28.0% | 298.1 tok/s | | 8,192 | 279.3 tok/s | 218.4 tok/s | +27.9% | 294.2 tok/s | fpA reduces acceptance-independent target-step time by about 21.8%. The remaining 2.8-5.1% throughput gap versus INT4 is primarily lower block-64 draft acceptance. MMLU-Pro 800 in thinking mode: INT2 block-64 fpA scored 652/800 (81.50%) versus 653/800 (81.625%) for the byte-identical-target INT4 control. The paired comparison had 14 INT2-only and 15 INT4-only correct answers (McNemar exact p=1.0), with zero unextracted answers. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

[CUDA] Add 2-bit MatMulNBits support

ce56dcc

David Fan (jiafatom) requested a balanced review from Copilot September 18, 2026 19:18

Copilot started reviewing on behalf of David Fan (jiafatom) September 18, 2026 19:18

View session

Copilot AI reviewed Sep 18, 2026

View reviewed changes

Akshay Sonawane (apsonawane) reviewed Sep 18, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc Resolved

[CUDA] Address 2-bit MatMulNBits review feedback

76377a4

Akshay Sonawane (apsonawane) requested a balanced review from Copilot September 19, 2026 00:06

Copilot started reviewing on behalf of Akshay Sonawane (apsonawane) September 19, 2026 00:06

View session

Copilot AI reviewed Sep 19, 2026

View reviewed changes

Comment thread onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc Resolved

Comment thread onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc Resolved

Comment thread onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh Resolved

Tianlei Wu (tianleiwu) marked this pull request as ready for review September 19, 2026 00:28

[CUDA] Add typed zero-point test coverage

c719cd4

This was referenced Sep 19, 2026

[CUDA] Add 2-bit fpA_intB GEMM and GEMV support tianleiwu/onnxruntime#1

Closed

[CUDA] Add 2-bit fpA_intB GEMM and GEMV support #32699

Merged

feat(builder): support INT2 DFlash 2 drafters microsoft/onnxruntime-genai#2600

Merged

Akshay Sonawane (apsonawane) requested a balanced review from Copilot September 21, 2026 21:23

Copilot started reviewing on behalf of Akshay Sonawane (apsonawane) September 21, 2026 21:24

View session

Copilot AI reviewed Sep 21, 2026

View reviewed changes

Akshay Sonawane (apsonawane) approved these changes Sep 21, 2026

View reviewed changes

Tianlei Wu (tianleiwu) merged commit 34d838f into microsoft:main Sep 21, 2026
90 checks passed

Dmitrii Fedorov (dmitrii-f-t27) mentioned this pull request Sep 22, 2026

[EPIC] Stage 1 — t27 Ternary Check: bit-exact conformance of ternary weight-packing formats on real models dmitrii-f-t27/trinity-memory#27

Closed

9 tasks done

Navigation Menu

Repository navigation

[CUDA] Add 2-bit MatMulNBits support - #32693

[CUDA] Add 2-bit MatMulNBits support#32693
Tianlei Wu (tianleiwu) merged 3 commits into
microsoft:mainfrom
tianleiwu:tlwu/matmul-nbits-2bit

Tianlei Wu (tianleiwu) commented Sep 18, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Akshay Sonawane (apsonawane) commented Sep 18, 2026

Uh oh!

Tianlei Wu (tianleiwu) commented Sep 19, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Tianlei Wu (tianleiwu) commented Sep 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (2)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Akshay Sonawane (apsonawane) commented Sep 18, 2026

Uh oh!

Tianlei Wu (tianleiwu) commented Sep 19, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🔵 Needs a closer look

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Tianlei Wu (tianleiwu) commented Sep 18, 2026 •

edited

Loading