2, 3, 5, 6 bit quantization exploration

[David Fan (jiafatom)]

Description

Documents the design space for 2-, 3-, 5-, and 6-bit weight-only quantization and defines the end-to-end delivery plan for mixed-width INT2 QMoE. The prioritized QMoE recipe uses INT2 expert gate/up weights and INT4 expert down weights while keeping the portable model format separate from provider-specific runtime prepacking.

QMoE Progress (October 6, 2026)

• Mixed-width QMoE contract and shared validation merged in Add mixed-width weight contract for QMoE #32697; bounded CUDA correctness fallback merged in Add CUDA correctness fallback for INT2 and mixed-width QMoE #32743.
• Packed CUDA INT2 and mixed INT2/INT4 decode merged in Add packed INT2 and mixed-width CUDA QMoE decode #32761 on September 28.
• PR4a [CUDA] Add SM80 FP16/BF16 INT2 grouped GEMM foundation #32963 merged on October 1 (a4d55e2b6f): SM80 FP16/BF16 INT2 grouped GEMM foundation, including the corrected byte-address alignment contract.
• PR4a2 [CUDA] Follow up INT2 grouped GEMM test coverage and plugin build (PR4a2) #33045 merged on October 1 (78c834918d): SM8x test coverage, plugin test build footprint, diagnostic options and requested file headers. The QUICK_BUILD baseline tactic issue was addressed before merge; it is no longer an outstanding delivery task. The redundant quantization fix is already in main via Replace old use of routing_record with routing_snapshot_ in onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc. #32932.
• Historical PR4a2 local A100 validation in Docker jiafa-dev: actual CUDA 12 legacy provider-test target built successfully; 23 grouped tests passed (3 benchmarks disabled), QMoE regression 53 passed / 11 skipped. Both passed memcheck with zero errors/leaks. Actual plugin internal-test target and fixture compiled; a standalone harness linked from these generated objects passed 16 tests under memcheck. Standalone CUDA 12/13 plugin and CUDA 12 legacy harnesses also passed memcheck/synccheck/initcheck and three focused racecheck cases. These results do not establish SM86/SM89 coverage. That local full CUDA 12 plugin provider build was blocked by third-party-header warning-as-error diagnostics; the historical standalone results are not a full plugin integration pass.
• PR4b [CUDA] Enable SM80+ FP16/BF16 mixed INT2/INT4 QMoE prefill by default #33005 merged on October 2 (2cdff3ca8e): default-enabled SM80+ FP16/BF16 packed INT2 FC1 / INT4 FC2 prefill, symmetric quantization, block size 64 and interleaved fused SwiGLU. Set ORT_ENABLE_QMOE_INT2_PREFILL=0 to opt out. Packed decode retains priority; ORT_DISABLE_MOE_GEMV=1 disables eligible INT2 and existing INT4/INT8 GEMV dispatch. Set environment controls before process startup.
• Eligible packed prefill now tiles within a 256 MiB estimated temporary-scratch budget, including packed workspace, tile-local routing metadata and uncached scale transposes. Explicit row tiling and long prompts stay packed rather than triggering full dense weight dequantization. Full/final tile sizes are accounted for; a one-row overflow reports a packed-scratch error. This is not a whole-process GPU memory cap. Unsupported configurations retain the bounded fallback. Regression coverage includes 8193-row FP16/BF16 cases with a one-byte dense scratch guard, top-k, bias, runtime scales and partial final tiles.
• SM80+ eligibility includes later architectures but retains the SM80 CUTLASS template/layout; suitable target code or compatible virtual PTX is required. Merge and metadata tests do not establish hardware execution coverage on every architecture.
• PR4c Extend CUDA QMoE packed INT2 prefill block sizes and weight combinations #33092 merged on October 5 (0c67e77bb7cb): extends packed prefill to block sizes 32, 64 and 128 and (FC1, FC2) widths (2,4), (2,2) and (4,2), dispatching each projection to the INT2 or INT4 grouped GEMM. Existing (4,4), INT8, zero-point and non-fused paths and width inheritance defaults are unchanged. Block size 32 has no GEMV path. Final-head local validation at db8a7625b8 on A100 with non-plugin Release CUDA 12.8, BF16 and internal tests enabled passed 15 outer tests and 26 internal tests, with no failures/skips (three benchmarks disabled). Coverage includes the 216-invocation FP16/BF16 integration matrix, direct grouped-GEMM block/tile cases, cached-scale decode-to-prefill transition, packed-decode fallback and block32 dense fallback with prefill disabled. Validation details: Extend CUDA QMoE packed INT2 prefill block sizes and weight combinations #33092 (comment) . This is local A100 correctness evidence, not plugin integration, every-GPU coverage or a performance claim.
• Historical PR4b validation: six focused tests passed in a clean non-plugin wheel build at 463fa42cb5. Full mixed GPT-OSS-20B validation demonstrated GEMM prefill -> GEMV decode in one session, with finite logits/KV and explicit branch counts (24 prefill and 48 decode calls for a 309-token prompt). These results predate the final default-dispatch and packed-tiling changes and must not be treated as validation of those revisions.
• Initial direct-ORT A100 benchmark compared existing mixed INT2/INT4 and INT4 GPT-OSS-20B exports: batch 1, 128/512/2048/4096 input tokens, 128 generated tokens, three warmups and ten measured requests per configuration. Mixed prefill TPS was 22.9%-30.8% higher; TTFT P50 was 18.5%-22.8% lower; decode TPS was 2.3%-4.1% lower. Sampled process peak GPU memory was 0.20-1.39 GiB higher, not lower. Full tables and measurement boundaries are in [CUDA] Enable SM80+ FP16/BF16 mixed INT2/INT4 QMoE prefill by default #33005. These historical timings also predate the final packed-tiling changes.
• Mixed model artifacts total 7.462 GiB versus 10.985 GiB for INT4: 3.523 GiB / 32.07% smaller. This is ONNX graph plus referenced external weights, not GPU residency.
• Performance results are an initial synthetic baseline, not tau2-bench or accuracy qualification. Existing exports differ in block size (mixed 64, INT4 32) and dispatch; sequential runs on a shared machine are not an isolated bit-width ablation. Repeated prompts and differing greedy continuations can affect expert routing.
• Remaining work: controlled alternating and matched-block-size benchmarks for the merged implementation; exact runtime-memory attribution; broader fused export integration, CPU/CUDA model parity, deployment-platform qualification, real-workload and target-model quality evaluation; repeated transition/concurrency/capture validation. Default packed prefill enablement is merged; generalized speedup, lower runtime memory and end-to-end product qualification are not claimed.

Branch and Documentation Update

Rebased onto main at 2cdff3ca8e (including PR4a2 and PR4b), without conflicts. Both design documents record the merged work and distinguish historical estimates from remaining qualification tasks. The October 6 update records merged PR4c (#33092, 0c67e77bb7cb) and its final-head local A100 validation in both design documents, including the expanded support boundary and remaining qualification limits. This roadmap update is documentation-only; no new build or runtime tests were run for this update.

Motivation and Context

Expert weights dominate storage and memory traffic in the target Qwen MoE models. This plan prioritizes a deployable mixed-width QMoE path while reusing dense MatMulNBits packing, validation, and CUDA primitives where useful. It does not treat the bounded dense fallback as production performance support, and it does not claim compatibility with GGUF IQ2_XS or IQ4_NL formats.

[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.

[Ti-Tai Wang (titaiwangms)]

Thanks for making mixed-width QMoE the explicit product priority. We aligned the
existing Olive and Mobius tracking with the updated plan:

• microsoft/Olive#2638 owns fused-expert SMP/recipe work and qualification of
  an Olive checkpoint with FC1 gate/up INT2 and FC2 down INT4.
• onnxruntime/mobius#735 owns checkpoint ingestion, independent FC1/FC2 packed
  layouts, gate/up interleaving, down-projection orientation, graph emission,
  initializer/external-data binding, and ORT CPU parity.

The existing microsoft/Olive#2671 and onnxruntime/mobius#740 qualify the
dense checkpoint-to-MatMulNBits(bits=2/4/8) path only; they do not claim fused
mixed-width QMoE support.

From the downstream side, we propose treating the ORT mixed-width QMoE contract
as the graph-emission gate. Could the contract work settle:

1. FC-specific attributes versus a versioned descriptor.
2. FC1/FC2 scale and zero-point shapes and dtypes.
3. Packed tensor ordering and whether the initial configuration shares one
   block size.
4. Inheritance and backward compatibility with expert_weight_bits.
5. The schema commit that Olive/Mobius should pin for implementation and shared
   fixtures.

Before that contract is fixed, we will keep work schema-neutral: checkpoint
qualification, expert target classification, packing oracles, quality/effective
size measurements, and a tiny deterministic source fixture. Once pinned, the
Mobius export work can proceed in parallel with the ORT CPU reference and use
the same fixture for CPU/CUDA parity.

Please also confirm whether native Mobius QMoE construction is sufficient for
the first exporter milestone, with dense-graph-to-QMoE fusion following
separately.

[David Fan (jiafatom)]

Thanks for aligning the downstream tracking. The mixed-width contract is now fixed by #32697, merged as 090f6ba1cfd6f4a0cb2481bac4b92156828a0e33. Answers to the five contract questions:

1. We chose FC-specific optional attributes on the existing QMoE v1 schema: fc1_expert_weight_bits, fc2_expert_weight_bits, and fc3_expert_weight_bits; there is no versioned descriptor in the initial contract.
2. For integer QMoE, weights and packed zero points are uint8; scales may be float, float16, or bfloat16. FC1 scales are [E, fusion_size * inter_size] row-wise or [E, fusion_size * inter_size, hidden_size / block_size] block-wise. FC2 scales are [E, hidden_size] or [E, hidden_size, inter_size / block_size]. Row-wise FC1/FC2 zero points are [E, ceil(fusion_size * inter_size * fc1_bits / 8)] and [E, ceil(hidden_size * fc2_bits / 8)]; block-wise forms are [E, fusion_size * inter_size, ceil((hidden_size / block_size) * fc1_bits / 8)] and [E, hidden_size, ceil((inter_size / block_size) * fc2_bits / 8)]. Zero-point rows pad to whole bytes.
3. The portable raw weight form (weights_prepacked=0) is [E, N, K * effective_bits / 8], with K-contiguous values packed low bits first within each byte. Provider prepacking is runtime-internal and is not an exporter contract. The initial schema has one shared block_size attribute for FC1/FC2/FC3.
4. Each omitted FC override inherits expert_weight_bits (default 4), preserving existing uniform models. Effective widths are restricted to 2/4/8. For fused SwiGLU, effective FC3 must equal effective FC1 because FC3 is stored in FC1.
5. Olive/Mobius should pin schema and graph-emission fixtures to merge commit 090f6ba1cfd6f4a0cb2481bac4b92156828a0e33. For executable ORT parity, Add CUDA correctness fallback for INT2 and mixed-width QMoE #32743 is merged as ee5f6e7cfa1f9e4e4253154e7c719a846280ca83; Add packed INT2 and mixed-width CUDA QMoE decode #32761 adds packed CUDA decode and remains a runtime optimization, not an exporter contract dependency.

Native Mobius QMoE construction is sufficient for the first exporter milestone. Dense-graph-to-QMoE fusion should follow separately and should not block the deterministic fused-QMoE fixture/export path. I recorded that milestone boundary and the Olive#2638/Mobius#735 ownership in 35f76b8ce5.

[G. Ramalingam (gramalingam)]

Please see onnx/onnx#8576 ... whether it makes sense

[David Fan (jiafatom)]

G. Ramalingam (@gramalingam) Thank you for the feedback in #32657 (comment) and for pointing me to onnx/onnx#8576 in #32657 (comment).

I reviewed the current Pack/Unpack proposal and updated the INT3 QMoE delivery plan in commit 1340936:

• The proposed portable INT3 weight format already matches its UINT8, last-axis, row-local LSB-first packing, including the [0,1,2,3,4,5,6,7] -> [0x88,0xc6,0xfa] example. Explicit zero points use the same packing with count=B (number of scale blocks), while weights use count=K.
• P0 now requires interoperability checks against a recorded proposal revision and an executable standard-operator unpack/dequantization reference, rather than deferring that reference until the committee submission.
• P2 will compare the independent scalar oracle and CUDA results with that reference: unpack unsigned codes, cast before zero-point subtraction, and apply scales/zero points using logical K-block indexing.
• P5 now includes numerical equivalence between the native fused QMoE export and a small standard-operator MoE reference using the same packed initializers. Automatic graph recognition/fusion remains follow-up work.
• Section 7.3 clarifies that generic Unpack ignores high padding bits, whereas the proposed INT3 QMoE profile requires canonical zero padding; that stricter check belongs at the profile/export boundary.

This keeps the standard interchange semantics separate from backend execution. CUDA may still fuse decoding with computation or convert the portable bytes into a provider-private prepack; the reference graph does not require production execution to materialize all experts as floating-point weights.

Since #8576 is still an open proposal, the plan records the reviewed revision and does not treat its draft operator/opset as a released dependency. Until it is approved and supported, the reference unpack stage can use an equivalent decomposition of existing standard operators. Native INT3/UINT3 datatype standardization remains outside this delivery sequence.

These are plan and acceptance-criteria updates, not a claim that INT3 QMoE kernels or the interoperability tests are already implemented.