[CUDA] Add 2-bit fpA_intB GEMM and GEMV support by tianleiwu · Pull Request #32699 · microsoft/onnxruntime

Tianlei Wu (tianleiwu) · 2026-09-19T21:22:21Z

Description

Adds 2-bit weight support to the CUDA CUTLASS fpA_intB GEMM and GEMV paths used by MatMulNBits. Eligible FP16 and BF16 bits=2 workloads can use fused weight-only kernels instead of dequantize+cuBLAS while preserving existing 4-bit and 8-bit code generation.

The branch is rebased onto 48795e0281, which enables BF16 in the compact fpA_intB build by default. Compact mode now includes FP16 and BF16 INT2 scale-only kernels at block_size=64.

Summary of Changes

INT2 fpA_intB kernels

Adds INT2 numeric conversion, SM80 interleaved layout support, preprocessing, dispatch, and kernel instantiations for FP16 and BF16 activations.
Supports 2-bit fpA_intB with N % 128 == 0, block_size 64 or 128, and standard or SM80-prepacked weights. Native SM90-prepacked INT2 weights remain unsupported.
Fixes the two-deep B-fragment pipeline for the odd single-load INT2 case by rotating register buffers with compile-time indices. This avoids the 64-byte per-thread local-memory frame produced by a runtime cursor; generated 4-bit and 8-bit PTX remains byte-identical.

Compact build

Includes FP16 and BF16 INT2 scale-only GEMM/GEMV kernels at block_size=64.
Retains FP16 and BF16 INT4/INT8 scale-only kernels at block_size=32.
Keeps block_size=128, zero-point, bias, and native SM90-layout variants behind onnxruntime_USE_FPA_INTB_GEMM_FULL=ON.
Requires SM75 or later for FP16 and SM80 or later for BF16.

Coverage and documentation

Adds exact A = I layout round-trip tests for INT2 and INT4 weights.
Adds compact FP16/BF16 INT2 eligibility, architecture, GEMM/GEMV, and end-to-end MatMulNBits coverage.
Fixes the kernel test helper to account for per-column byte padding in packed INT2 zero points.
Documents 2-bit routing, compact/full build constraints, the B-register pipeline fix, and rejected optimization alternatives.

Performance

Measured on H200 with FP16 activations and block_size=128, the fused fpA_intB GEMV is fastest for every tested production shape at M=1 (5.41-17.46 us versus 5.67-22.23 us for the hand-written 2-bit GEMV and 16.19-141.11 us for dequantize+cuBLAS). The register-buffer rotation improves INT2 GEMM by 4.2-7.1% for M in {128, 512, 2048}.

Testing

Built onnxruntime_providers_cuda, onnxruntime_providers_cuda_ut, onnxruntime_provider_test, and onnxruntime_test_all with CUDA 13.0, compact fpA_intB enabled, SM80/SM90 targets, and CUDA EP internal tests enabled.
Passed the CUDA EP internal suite on A100: 288 passed, 3 skipped, 1 disabled.
Passed MatMul2BitsCuda.BFloat16_FpAIntB on A100 across four compact BF16/INT2 shapes.
Passed clang-format --dry-run --Werror for all changed C++/CUDA files and git diff --check.
Verified the compact build graph includes bf16_int2_gemm_scaleonly.cu and dispatcher_bf16_int2.cu.

Checklist

Tests added/updated
No breaking changes
Documentation updated

Copilot

Copilot review overview

🟡 Changes recommended

The fpA_intB layout tests can exercise fallback paths instead, and one production-shape test introduces excessive reference-computation cost.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 3

Open (3)

Avoid prohibitively expensive scalar reference for large-M test · New
Use fpA_intB-eligible block sizes in layout tests · New
Disable CPU fallback in CUDA-specific regression test · New

What changed in this PR

Adds CUDA 2-bit MatMulNBits support across fused GEMV, CUTLASS fpA_intB GEMM, preprocessing, fallback dequantization, tests, and documentation.

Changes:

Implements FP16/BF16 INT2 CUTLASS GEMM/GEMV dispatch and layouts.
Adds generic 2-bit CUDA kernels and dequantization fallback.
Expands validation, tests, compact-build configuration, and documentation.

File	Description
`onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc`	Adds exact weight-layout tests.
`onnxruntime/test/contrib_ops/matmul_2bits_test.cc`	Adds CUDA INT2 end-to-end coverage.
`onnxruntime/test/contrib_ops/cuda_kernels/fpA_intB_gemm_kernel_test.cc`	Extends kernel tests to INT2.
`onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h`	Adds INT2 validation and dispatch.
`onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cuh`	Declares INT2 fast-path dispatch.
`onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc`	Integrates preprocessing, profiling, and fallback.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits.cu`	Dispatches INT2 fused kernels.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1.cuh`	Declares M=1 kernel entry point.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_m1_impl.cuh`	Implements M=1 GEMV.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_half.cu`	Instantiates FP16 M=1 support.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_float.cu`	Instantiates FP32 M=1 support.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_common.cuh`	Adds shared INT2 conversion primitives.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_bfloat16.cu`	Instantiates BF16 M=1 support.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched.cuh`	Declares small-M dispatch.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_impl.cuh`	Implements batched INT2 kernels.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_half.cu`	Instantiates FP16 batched support.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_float.cu`	Instantiates FP32 batched support.
`onnxruntime/contrib_ops/cuda/quantization/matmul_2bits_batched_bfloat16.cu`	Instantiates BF16 batched support.
`onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise.cuh`	Adds generic INT2 dispatch.
`onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_2bits.cu`	Implements INT2 fallback dequantization.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.h`	Adds INT2 kernel types.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/fpA_intB_gemv.cu`	Adds INT2 launch and support checks.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher.h`	Updates reduction and INT2 dispatch geometry.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher_fp16_int2.cu`	Instantiates FP16 INT2 GEMV.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/dispatcher_bf16_int2.cu`	Instantiates BF16 INT2 GEMV.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h`	Defines INT2 type and conversion traits.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm_template.h`	Routes INT2 through SM80-compatible GEMM.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fp16_int2_gemm_scaleonly.cu`	Instantiates FP16 scale-only GEMM.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fp16_int2_gemm_scale_zeros.cu`	Instantiates FP16 affine GEMM.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/bf16_int2_gemm_scaleonly.cu`	Instantiates BF16 scale-only GEMM.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/bf16_int2_gemm_scale_zeros.cu`	Instantiates BF16 affine GEMM.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h`	Defines INT2 profiling width.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc`	Adds INT2 profiler sizing and execution.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors.h`	Adds the W2_A16 quantization type.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.h`	Defines INT2 layouts and permutations.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_preprocessors_impl.cu`	Implements INT2 transpose and interleave.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.h`	Generalizes packed zero-point handling.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_adaptor.cu`	Implements INT2 unpacking and zero points.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/interleaved_numeric_conversion.h`	Adds UINT2 FP16/BF16 converters.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_percol.h`	Fixes odd B-fragment pipelining.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_pipelined_finegrained.h`	Fixes fine-grained pipelining.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_percol.h`	Fixes multistage per-column pipelining.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/dq_mma_multistage_finegrained.h`	Fixes multistage fine-grained pipelining.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma.h`	Adds FP16/UINT2 MMA specializations.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_mma_bf16.h`	Adds BF16/UINT2 MMA specializations.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_pipelined.h`	Permits UINT2 pipelined MMA.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/threadblock/default_dq_mma_multistage.h`	Permits UINT2 multistage MMA.
`onnxruntime/contrib_ops/cuda/llm/cutlass_extensions/gemm/kernel/mixed_gemm_B_layout.h`	Defines the INT2 interleaved layout.
`onnxruntime/contrib_ops/cpu/quantization/matmul_nbits_helper.h`	Validates INT2 block-size limits.
`docs/contrib_ops/cuda/matmul_nbits.md`	Documents INT2 routing and layouts.
`cmake/onnxruntime_cuda_source_filters.cmake`	Includes INT2 kernels in compact builds.

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

Extends the CUTLASS mixed-input weight-only path (fpA_intB) to uint2b_t weights so MatMulNBits with bits=2 can use the interleaved fast kernels instead of dequantize+cuBLAS. - interleaved_numeric_conversion: Int2NumericArrayConverter for half_t (4 masks + 1 shift, fma.rn.f16x2) and bfloat16_t (shift per output, since bf16's 7 mantissa bits cannot hold a field above shift 4). - mixed_gemm_B_layout / default_mma{,_bf16} / default_dq_mma_*: uint2b_t layout details (ColumnsInterleaved = 8) and DefaultMma specializations. - dq_mma_{multistage,pipelined}_{finegrained,percol}: the warp_frag_B[2] double buffer was indexed by the within-tile load index, which is only a valid ping-pong when kWarpGemmIterationsForB is even. For 2-bit it is 1, so every load wrote slot 1 and every MMA read slot 0, silently computing B[n][k % 64]. Replaced with a running read cursor that flips after the last k-iteration of each B load. Bit-identical for 4/8-bit. - Prepack: QuantType::W2_A16, 64-entry permutation map, generalized sub-byte transpose over BITS_PER_ELT, int2 bias+interleave kernel, and unpack_uint2_transposed_to_int8_direct_cuda. Zero-point packing is now templated on ZeroPointBits (0/4/2). - GEMV: Int2 details/traits, generalized warp_reduce_sum for interleave 8, CtaN tuned for kStepK == 64. - MatMulNBits: eligibility requires N % 128 == 0 and block_size >= 64 for 2-bit (one B fragment spans a full 64-element K tile), and rejects the Hopper prepacked layout. Adds an exact (A = I) layout round-trip regression test for 2- and 4-bit, 2-bit fpA_intB shape/production-shape tests, and kernel-level Int2 tests.

…ursor The two-deep warp B register pipeline is indexed by the within-tile load index, which only ping-pongs when kWarpGemmIterationsForB is even. 2-bit weights make it 1, so the previous commit added a running read cursor. That cursor is correct but flips an odd number of times per mainloop iteration, so it is not loop-invariant: warp_frag_B becomes dynamically indexed and ptxas gives the CtaShape128x128x64 kernel a 64-byte per-thread stack frame (note the tell-tale "0 bytes spill stores, 0 bytes spill loads" -- indexed local memory, not spilling). Rotating the two buffers instead keeps both indices compile-time constants. Measured on H200 (fp16, block_size=128, bonsai2-27B shapes), 2-bit GEMM: -4.2% to -7.1% at M in {128, 512, 2048}, e.g. N=17408 K=5120 M=2048 goes 1811 -> 1703 us. 4- and 8-bit generated PTX is byte-identical to the pre-2-bit baseline (verified by diffing ptx for both widths). Also records in the docs that converting the B fragment one MMA k-step at a time -- the previously suggested fix for the large-M gap -- was implemented and measured as a 4-8% regression: it saves 6 registers but occupancy is 1 CTA/SM either way, so the gather is pure added ALU.

…l set 2-bit was reachable only in a full build: the compact source filter dropped every int2 translation unit, select_gs instantiated GroupSize 32 only, and the eligibility gate required block_size == 32 -- which 2-bit can never satisfy, since one B fragment spans a whole 64-element K tile. Adds the FP16 scale-only int2 variant at block_size 64, the smallest group 2-bit supports: - source filter keeps fp16_int2_gemm_scaleonly.cu and dispatcher_fp16_int2.cu - select_gs instantiates GroupSize 64, gated on Details::kStepK >= 64 so the int4/int8 units are unchanged - CheckFpAIntBEligibility and InitGemmProfiler accept bits=2 at block_size 64 with N % 128 == 0 Cost in a compact build: 4 -> 6 translation units, +1.5 MB of objects (dispatcher_fp16_int2 276 KB, fp16_int2_gemm_scaleonly 1259 KB, versus 283 KB / 1272 KB for the int4 pair already present). BF16, block_size 128 and zero-points still require onnxruntime_USE_FPA_INTB_GEMM_FULL=ON. Note that 2-bit models quantised with block_size 128 and zero-points -- the common case -- therefore still need the full build; this only covers the scale-only block_size 64 contract.

Records the analysis behind kPerm_W2_A16: it compensates for ldmatrix moving 16-bit elements when the weights are narrower, so it composes with the fixed ldmatrix pattern and converter shuffle to the identity and has no free parameters. Two checks back this up -- ncu measures 0 shared-memory bank conflicts in the 2-bit CtaShape128x128x64 kernel (4-bit is also 0, and the swizzled ColumnMajorTensorOpMultiplicandCrosswise smem layout is what handles conflicts), and perturbing the map fails the exact A=I round-trip test. Also corrects the sliced-conversion regression figure: 9-15% against the shipped rotation kernel, not 4-8% against the superseded cursor version.

## 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>

Copilot AI balanced review requested due to automatic review settings September 19, 2026 21:22

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) September 19, 2026 21:22

View session

Tianlei Wu (tianleiwu) mentioned this pull request Sep 19, 2026

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

Closed

3 tasks done

Copilot AI reviewed Sep 19, 2026

View reviewed changes

Comment thread onnxruntime/test/contrib_ops/matmul_2bits_test.cc Resolved

Comment thread onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc Resolved

Comment thread onnxruntime/test/contrib_ops/matmul_nbits_fpa_intb_layout_test.cc Resolved

Tianlei Wu (tianleiwu) mentioned this pull request Sep 20, 2026

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

Merged

Tianlei Wu (tianleiwu) force-pushed the tlwu/matmul-nbits-2bit-fpa-intb branch from 1b6be66 to 5a3d5e3 Compare September 21, 2026 22:37

Tianlei Wu (tianleiwu) requested review from Akshay Sonawane (apsonawane), David Fan (jiafatom) and kunal-vaishnavi September 21, 2026 23:12

Tianlei Wu (tianleiwu) added 6 commits September 22, 2026 09:03

Address MatMulNBits test feedback

12b3a3c

feat(cuda): enable compact BF16 INT2 kernels

7e698e3

Tianlei Wu (tianleiwu) force-pushed the tlwu/matmul-nbits-2bit-fpa-intb branch from 5a3d5e3 to 7e698e3 Compare September 22, 2026 11:35

David Fan (jiafatom) approved these changes Sep 22, 2026

View reviewed changes

Tianlei Wu (tianleiwu) merged commit 307768a into microsoft:main Sep 22, 2026
90 checks passed

Thiago Pereira Rocha (thpereir) mentioned this pull request Sep 23, 2026

[CUDA] QMoE: GeGLU-int2 activation + fractional zero_point_offset #32198

Open

3 tasks done

Navigation Menu

Repository navigation

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

[CUDA] Add 2-bit fpA_intB GEMM and GEMV support#32699
Tianlei Wu (tianleiwu) merged 6 commits into
microsoft:mainfrom
tianleiwu:tlwu/matmul-nbits-2bit-fpa-intb

Tianlei Wu (tianleiwu) commented Sep 19, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Description

Summary of Changes

INT2 fpA_intB kernels

Compact build

Coverage and documentation

Performance

Testing

Checklist

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!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Tianlei Wu (tianleiwu) commented Sep 19, 2026 •

edited

Loading