[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 introduced for MatMulNBits. This lets eligible FP16/BF16 bits=2 workloads use fused weight-only kernels instead of dequantize+cuBLAS, while preserving the existing 4-bit and 8-bit code generation.

This PR depends on #32693. Until that PR merges, the GitHub diff against main also includes its commits; the four fpA_intB commits can be reviewed independently in the feature-only comparison.

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 x INT2 scale-only kernels at block_size=64 in compact builds.
Keeps BF16, block_size=128, and zero-point variants behind onnxruntime_USE_FPA_INTB_GEMM_FULL=ON.

Coverage and documentation

Adds exact A = I layout round-trip tests for INT2 and INT4 weights.
Adds fpA_intB kernel and end-to-end MatMulNBits tests, including the Bonsai 2 production shape and ternary zero-point cases.
Documents 2-bit routing, layout 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

Incrementally rebuilt onnxruntime_test_all with CUDA 13.0, onnxruntime_USE_FPA_INTB_GEMM_FULL=ON, and onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON.
Passed 7 focused GPU tests from MatMul2BitsCuda and MatMulNBitsFpAIntBLayout, including the production-shape test.
Passed lintrunner -m origin/pr/32693.
Verified 4-bit and 8-bit generated PTX is byte-identical to the pre-INT2 baseline.

Checklist

Tests added/updated
No breaking changes
Documentation updated

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.

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.

+  for (int64_t m : {int64_t{1}, int64_t{8}, int64_t{512}}) {
+    TestOptions2Bits opts{};
+    opts.M = m;
+    opts.N = 10240;
+    opts.K = 5120;
+    opts.block_size = 128;
+    opts.has_zero_point = true;
+    opts.use_cuda = true;
+    opts.output_abs_error = 0.2f;
+    opts.output_rel_error = 0.02f;
+    RunTest2Bits<MLFloat16>(opts);


+// K spans four 64-element threadblock K tiles, which is what caught the B pipeline stalling on
+// the first tile. N is a multiple of 128 so 2-bit weights are fpA_intB-eligible.
+constexpr int64_t kN = 128, kK = 256, kBlockSize = 128;


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


Tianlei Wu (tianleiwu) added 7 commits September 18, 2026 19:00

[CUDA] Add 2-bit MatMulNBits support

ce56dcc

[CUDA] Address 2-bit MatMulNBits review feedback

76377a4

[CUDA] Add typed zero-point test coverage

c719cd4

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

Copilot AI reviewed Sep 19, 2026

View reviewed changes

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

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

Open

Navigation Menu

[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) wants to merge 7 commits into
microsoft:mainfrom
tianleiwu:tlwu/matmul-nbits-2bit-fpa-intb

Tianlei Wu (tianleiwu) commented Sep 19, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Tianlei Wu (tianleiwu) commented Sep 19, 2026

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants