2, 3, 5, 6 bit quantization exploration

[David Fan (jiafatom)]

Description

Motivation and Context

[Copilot]

🟡 Changes recommended

The support matrix incorrectly reports that WebGPU MatMulNBits lacks 2-bit support.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Explores INT2/INT6 weight-only quantization and recommends a scoped CUDA INT2 implementation.

Changes:

• Documents current support, external approaches, packing options, evaluation criteria, and delivery estimates.
• Proposes phased CUDA INT2 development with INT6 as follow-up.

File summaries

File	Description
docs/design/2bit-6bit-weight-only-quantization-exploration.md	Adds the quantization exploration and roadmap.

Review details

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

---

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

[Ti-Tai Wang (titaiwangms)]

Thanks for putting this exploration together. From the Olive/Mobius side, we want to clarify the end-to-end model-production dependency for the proposed INT2 CUDA work.

Even after ORT implements the CUDA MatMulNBits(bits=2) kernels, Olive and Mobius may need corresponding work before users can produce deployable INT2 models. There are two possible workflows:

1. ONNX-native: FP ONNX → Olive OnnxBlockWiseRtnQuantization → MatMulNBits(bits=2).
2. PyTorch/Mobius: HF/PyTorch → Olive RTN/GPTQ/SMP INT2 or mixed-precision checkpoint → Mobius → MatMulNBits(bits=2/4/8).

The current codebase has different gaps for these paths. Olive's PyTorch RTN/GPTQ/KQuant and SelectiveMixedPrecision already accept 2-bit configurations, but the INT2 checkpoint-to-ONNX path still needs end-to-end qualification. Olive's ONNX RTN and built-in quantized-linear export paths also need additional work for complete INT2 support. Mobius can construct MatMulNBits(bits=2), but Olive-format uniform INT2 and mixed INT2/INT4/INT8 checkpoint export need end-to-end graph, weight-binding, and numerical-parity coverage.

We expect both workflows to be supported eventually, but could the document clarify which one is required for the initial CUDA vertical slice? Given the GPTQ, SMP, and mixed-precision quality goals, the PyTorch/Mobius path may be the more relevant first milestone, with ONNX-native INT2 RTN following for ONNX-only workflows.

Could the scope also explicitly confirm whether the initial target is limited to dense MatMulNBits projections? If fused MoE/QMoE expert weights are included, that introduces a separate Olive/Mobius/runtime contract and should be tracked as a distinct workstream.

[David Fan (jiafatom)]

Thanks for clarifying the model-production dependency. Updated in f20da36 with two explicit scope decisions:

1. The initial vertical slice uses the HF/PyTorch -> Olive RTN/GPTQ/SelectiveMixedPrecision checkpoint -> Mobius -> dense mixed-bit MatMulNBits ONNX workflow, including graph conversion, initializer binding, tensor-wise bit-width, and numerical-parity qualification. ONNX-native Olive INT2 RTN/export follows as a separate milestone.
2. The initial CUDA target is limited to individual dense MatMulNBits projections. Expert gate/up INT2 is included only when exported as independent dense nodes; fused MoE/QMoE export, packing, and runtime contracts are explicitly out of scope and tracked separately.

The committed scope and dated schedule now reflect the export qualification work.

[Copilot]

🔵 Needs a closer look

The support matrix inaccurately describes existing 2-bit QDQ packing and runtime fusion support.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

docs/design/2bit-6bit-weight-only-quantization-exploration.md:37

• This note misidentifies the remaining QDQ gap. A native 2-bit QDQ packer is already exported as quantize_qdq_matmul_2bits (onnxruntime/python/onnxruntime_pybind_quant.cc:210-211), and runtime DQ→MatMulNBits(bits=2) fusion is implemented and tested (onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc:24-26,296-300; onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc:763-768). What remains 4-bit-only is the high-level Python quantizer's QDQ branch (matmul_nbits_quantizer.py:1032-1037). Please describe that distinction so the roadmap does not plan already-completed native/runtime work.

• Files reviewed: 1/1 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[David Fan (jiafatom)]

Added the INT2 QMoE investigation in e751729. The document now:

• keeps CUDA MatMulNBits(bits=2) as the first implementation priority;
• identifies the single expert_weight_bits attribute as the blocker for the target FC1 INT2 / FC2 INT4 recipe;
• records the current CPU, CUDA, WebGPU, and export gaps;
• treats full expert dequantization only as a correctness oracle, not a production milestone; and
• defines QMoE contract/export and fused decode as the next track after the MatMulNBits packing, correctness, and M=1 primitives stabilize.

The November dense MatMulNBits commitment remains unchanged; fused QMoE is still explicitly follow-up scope.

[David Fan (jiafatom)]

Added a standalone QMoE-first delivery plan in 19a723d: docs/design/int2-qmoe-end-to-end-delivery-plan.md.

It defines the end-to-end path from Olive calibration and Mobius export through CPU reference, CUDA correctness, packed decode, bounded/native prefill, and Qwen quality/performance qualification. The initial target is mixed FC1 INT2 / FC2 INT4, with backward-compatible schema semantics and explicit separation from IQ2_XS/IQ4_NL.

The proposed eight-week milestone is contract + reproducible export + CPU reference + bounded CUDA correctness + packed decode on one architecture. Production grouped-GEMM prefill and broad architecture coverage require follow-up time or additional CUDA staffing. The original exploration document now points to this plan and identifies QMoE as the product priority.

[David Fan (jiafatom)]

Extended the MatMulNBits exploration in 929ee53 to include INT3 and INT5 alongside INT6. The update adds:

• current support-matrix gaps for INT3/INT5/INT6;
• one canonical LSB-first contiguous bitstream contract;
• candidate runtime prepacking as 2+1 planes for INT3, 4+1 for INT5, and 4+2 for INT6;
• a common 2/3/4/5/6/8-bit quality and effective-size gate; and
• a cost model that selects at most one new width for initial CUDA decode/prefill prototypes instead of committing to three kernel families.

INT2 QMoE remains the product priority. INT3/INT5/INT6 are explicitly a parallel measured research track, with INT3 the first optimization candidate only if the data supports it.