From 50522cc3aaee3d1537eb5a9fe5137c31c467da16 Mon Sep 17 00:00:00 2001 From: David Fan Date: Wed, 16 Sep 2026 16:51:31 +0000 Subject: [PATCH 1/8] Document 2-bit and 6-bit quantization exploration --- ...it-weight-only-quantization-exploration.md | 289 ++++++++++++++++++ 1 file changed, 289 insertions(+) create mode 100644 docs/design/2bit-6bit-weight-only-quantization-exploration.md diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md new file mode 100644 index 0000000000000..5cc35a80407b4 --- /dev/null +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -0,0 +1,289 @@ +# 2-Bit and 6-Bit Weight-Only Quantization Exploration + +## Executive Summary + +ONNX Runtime already has a substantial 2-bit foundation: the `MatMulNBits` model format, Python quantization tooling, CPU kernels, and correctness tests support 2-bit weights. The largest practical gap is CUDA `MatMulNBits`, whose execution and prepacking paths currently support only 4-bit and 8-bit weights. + +The recommended first implementation target is therefore 2-bit CUDA `MatMulNBits`. A correctness-first dequantization plus cuBLAS fallback would establish end-to-end coverage before investing in fused GEMV and small-M kernels. + +Six-bit support should begin as a format and value-validation spike. It is not a small extension of the existing implementation: current packing code assumes that the bit width divides eight, while 6-bit values cross byte boundaries. There is also no native INT6 tensor type or NVIDIA Tensor Core instruction. A portable 6-bit representation and execution-provider prepacking contract should be agreed upon before optimized kernels are implemented. + +## Scope + +This document focuses on weight-only block quantization for `MatMulNBits` and related embedding gather operations. It covers: + +- Existing ONNX Runtime support and gaps. +- Relevant implementations in other inference and quantization projects. +- Model-format and packing choices. +- A staged implementation and evaluation plan. + +Activation quantization, KV-cache quantization, and floating-point formats such as FP4 and FP6 are outside the primary scope. + +## Current ONNX Runtime Support + +### Support Matrix + +| Component | 2-bit | 6-bit | Notes | +| --- | --- | --- | --- | +| `MatMulNBits` contrib schema | Yes | No | The schema attribute lists 2, 4, and 8 as supported values. | +| Python `MatMulNBits` quantizer | Yes | No | Dedicated 2-bit, 4-bit, and 8-bit native packers are exported. | +| CPU `MatMulNBits` | Yes | No | Includes fallback, AVX512/VNNI, and Arm64 work. | +| CUDA `MatMulNBits` | No | No | Kernel dispatch, prepacking, and memory estimation accept only 4 and 8 bits. | +| WebGPU `MatMulNBits` | No | No | Accepts only 4 and 8 bits. | +| CPU `GatherBlockQuantized` | Yes | No | Supports packed 2-bit, 4-bit, and 8-bit integer data. | +| CUDA `GatherBlockQuantized` | Yes | No | The generic packed-`uint8` path supports 2-bit extraction. | +| QDQ quantization helper | Limited | No | The dedicated native QDQ packing helper is currently 4-bit-oriented. | + +### Existing 2-Bit Assets + +The following pieces significantly reduce the cost and risk of a 2-bit CUDA implementation: + +- `onnxruntime/core/graph/contrib_ops/contrib_defs.cc` defines the portable `MatMulNBits` inputs and block layout. +- `onnxruntime/python/onnxruntime_pybind_quant.cc` exports `quantize_matmul_2bits`. +- `onnxruntime/python/tools/quantization/matmul_nbits_quantizer.py` creates 2-bit `MatMulNBits` models. +- `onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc` provides the CPU execution path. +- `onnxruntime/core/mlas/` contains optimized and fallback 2-bit implementations for supported CPU architectures. +- `onnxruntime/test/python/quantization/test_op_matmul_2bits.py` exercises model conversion and inference correctness. +- `onnxruntime/test/python/quantization/test_quantizeblockwise_2bits.py` verifies blockwise packing. +- `onnxruntime/test/mlas/unittest/test_sqnbitgemm_2bit_gemm.cpp` covers optimized MLAS behavior. + +The CPU history shows that 2-bit support has moved beyond a prototype. It includes AVX512/VNNI kernels, Arm64 kernels, optimized dequantization, float zero-point handling, and fallback paths. + +### CUDA Gap + +CUDA `MatMulNBits` currently assumes either 4-bit or 8-bit weights in several places: + +- Operator construction and dispatch. +- Blockwise dequantization specialization. +- Fused M=1 and small-M kernels. +- CUTLASS `fpA_intB` eligibility and weight conversion. +- Offline and runtime prepacked layouts. +- Workspace and persistent-memory estimates. + +This means enabling 2-bit requires more than relaxing one validation check. However, the generic dequantize-to-floating-point plus cuBLAS path provides a narrow route to initial correctness without extending CUTLASS prepacked formats. + +## External Landscape + +### llama.cpp + +llama.cpp has broad execution support for Q2_K, IQ2 variants, and Q6_K across CPU, CUDA, Metal, and additional backends. Its formats are not equivalent to the ONNX Runtime `MatMulNBits` format. + +Q6_K stores each 6-bit code as two planes: + +- Lower four bits: `QK_K / 2` bytes. +- Upper two bits: `QK_K / 4` bytes. +- Quantized sub-block scales plus a super-block scale. + +This layout is friendly to SIMD and GPU extraction, but its hierarchical scaling and metadata produce an effective size greater than exactly six bits per weight. Current llama.cpp documentation reports approximately 6.56 bits per weight for Q6_K in representative models. + +The main lesson for ONNX Runtime is that an optimized 6-bit implementation benefits from a deliberately designed physical layout. Q6_K should not be treated as a drop-in representation for uniform `MatMulNBits` quantization. + +### AutoGPTQ + +AutoGPTQ accepts 2-bit, 3-bit, 4-bit, and 8-bit GPTQ formats and contains legacy CUDA 2-bit matrix-vector kernels. Its most optimized modern paths, such as Marlin, remain 4-bit-specific. Mainline GPTQ packing does not provide a comparable production 6-bit path. + +AutoGPTQ is useful as a reference for 2-bit bit extraction and small-batch execution, but its shape restrictions, layouts, and project status make direct reuse less attractive than extending the existing ONNX Runtime 4-bit/8-bit CUDA structure. + +### TorchAO + +TorchAO contains generic x-bit packing and CPU/MPS low-bit operators. Its MPS implementation explicitly packs four 6-bit weights into three bytes. It also demonstrates an important architectural pattern: portable or generic quantized values can be converted into opaque, backend-specific packed formats selected by the runtime. + +Backend coverage and maturity vary by bit width. The most broadly optimized CUDA weight-only paths still center on 4-bit and 8-bit execution. + +### TensorRT-LLM + +TensorRT-LLM weight-only CUDA kernels and tests focus on INT4, INT8, and NVFP4. Its kernel type registry does not expose equivalent INT2 or INT6 weight-only paths. This is consistent with current NVIDIA hardware acceleration, which does not provide a native INT6 Tensor Core operation. + +### ONNX Data Types + +ONNX standardizes `INT2` and `UINT2`. Four consecutive 2-bit values are packed into one byte from least significant bits to most significant bits. This is compatible with the logical ordering already used by ONNX Runtime's packed 2-bit weights. + +ONNX currently has floating-point 6-bit types, but no standard `INT6` or `UINT6` tensor type. `MatMulNBits` can still carry 6-bit codes in an opaque `uint8` blob, but the contrib operator must define the packing contract precisely. + +## Design Considerations + +### 2-Bit Packing + +Two-bit packing naturally fits the current representation: + +```text +packed = (x0 & 0x03) + | ((x1 & 0x03) << 2) + | ((x2 & 0x03) << 4) + | ((x3 & 0x03) << 6) +``` + +No value crosses a byte boundary, and block sizes supported by `MatMulNBits` are multiples of four. The primary work is therefore execution-provider implementation and performance tuning rather than format design. + +### 6-Bit Packing Options + +#### Option A: Canonical Contiguous Bitstream + +Pack four values into three bytes, with each value occupying six consecutive bits. This gives exactly six payload bits per weight and a portable representation. + +Advantages: + +- Minimal model size. +- A simple canonical definition independent of an execution provider. +- Similar to generic x-bit packing in other frameworks. + +Disadvantages: + +- Some values cross byte boundaries. +- Existing `8 / bits` pack and unpack logic cannot be reused. +- Direct kernel extraction requires additional shifts and merged loads. + +#### Option B: Lower-4/Upper-2 Split Planes + +Store low four bits and high two bits in separate regions, following the basic physical idea used by Q6_K. + +Advantages: + +- Efficient aligned loads and extraction. +- Easier reuse of 4-bit and 2-bit unpacking primitives. +- Better fit for SIMD, DP4A, and backend-specific prepacking. + +Disadvantages: + +- Requires a new explicitly documented layout. +- Is less naturally represented as one generic bitstream. +- Can complicate interoperability with generic quantization tools. + +#### Recommended Contract + +Use a canonical contiguous bitstream in the portable model and allow each execution provider to prepack it into an opaque optimized representation, such as lower-4/upper-2 planes. Do not change the interpretation of existing models based only on `bits=6`. The exact canonical ordering, padding, zero-point packing, and versioning behavior must be specified first. + +## Proposed Roadmap + +### Phase 0: Define Targets and Baselines + +Before implementation, agree on: + +- Target execution providers and GPU architectures. +- Target models and model sizes. +- Decode, small-batch, and prefill workloads. +- Required symmetric and asymmetric quantization support. +- Required block sizes, bias, `g_idx`, and data types. +- Accuracy and performance acceptance criteria. + +### Phase 1: 2-Bit CUDA Correctness + +1. Add a CUDA 2-bit blockwise dequantization kernel. +2. Route 2-bit `MatMulNBits` through dequantization plus cuBLAS. +3. Keep CUTLASS and offline-prepacked paths disabled for 2-bit initially. +4. Update workspace and memory estimation for the fallback path. +5. Add CUDA tests for symmetric and asymmetric zero points, bias, tails, and supported block sizes. +6. Add model-level quantization and inference tests using the existing Python packer. + +This phase provides complete functionality and a stable reference for optimized kernels. + +### Phase 2: 2-Bit CUDA Performance + +1. Implement a fused M=1 GEMV path for token decode. +2. Implement or adapt a small-M batched kernel. +3. Evaluate an INT8 activation plus INT2 weight dot-product strategy. +4. Add runtime prepacking if profiling shows that the canonical layout limits load efficiency. +5. Tune dispatch thresholds against the dequantization plus cuBLAS fallback. + +Pure 2-bit round-to-nearest quantization may not meet model-quality targets. GPTQ, HQQ, K-quant-inspired optimization, importance-aware quantization, and mixed 2-bit/4-bit layer assignment should be evaluated alongside kernel work. + +### Phase 3: 6-Bit Format and Value Spike + +1. Write a precise portable packing specification. +2. Implement Python pack and unpack reference functions. +3. Add CPU reference dequantization and correctness tests. +4. Quantize representative models with 4-bit, 6-bit, and 8-bit configurations. +5. Measure quality, model size, load time, and dequantization overhead. +6. Prototype contiguous and split-plane CUDA extraction with a microbenchmark. + +The output of this phase should be a go/no-go decision for optimized MLAS and CUDA work. + +### Phase 4: 6-Bit Optimized Execution + +Proceed only if Phase 3 demonstrates a meaningful quality, memory, or latency niche that is not covered by mixed 4-bit/8-bit quantization. + +Potential work includes: + +- CPU SIMD dequantization and GEMM integration. +- CUDA fused decode GEMV. +- Small-M CUDA execution. +- Execution-provider-specific prepacking. +- Model conversion and compatibility tests. + +## Evaluation Plan + +### Correctness + +- Compare operator output against explicit dequantization plus floating-point MatMul. +- Cover FP32, FP16, and BF16 where supported. +- Cover symmetric, asymmetric, omitted, packed, and floating-point zero points as applicable. +- Cover K and N tails, empty dimensions, multiple leading dimensions, and block sizes 16 through 256. +- Verify model serialization and execution-provider fallback behavior. + +### Accuracy + +At minimum, compare: + +- 2-bit RTN, GPTQ, HQQ, and K-quant-inspired quantization. +- Uniform 2-bit versus mixed 2-bit/4-bit. +- Uniform 4-bit, 6-bit, and 8-bit. +- Perplexity plus representative downstream tasks used by the ONNX Runtime model-validation pipeline. + +### Performance + +Measure operator latency and end-to-end model performance for: + +- M=1 decode. +- Small M values such as 2, 4, 8, 16, and 32. +- Prefill M values of 128 and above. +- Representative transformer K and N dimensions. +- Each target GPU architecture. + +Report: + +- Latency and throughput. +- Effective memory bandwidth. +- Peak and persistent memory. +- Runtime prepacking cost. +- Model load time. +- End-to-end tokens per second. + +## Risks + +- Two-bit quality may require mixed precision or more expensive calibration algorithms. +- A 2-bit CUDA kernel can become unpacking-bound and fail to outperform a mature 4-bit kernel. +- Six-bit saves only 25 percent of weight payload relative to 8-bit while requiring non-native unpacking. +- A poorly specified 6-bit format could create incompatible models across execution providers. +- Extending CUTLASS prepacked formats too early would increase scope before value is established. +- Metadata overhead can materially change the effective bits per weight, especially for small blocks. + +## Recommended Initial Deliverables + +### 2-Bit Track + +- CUDA correctness fallback for `MatMulNBits`. +- Focused CUDA and Python tests. +- Decode and prefill benchmark baseline against 4-bit and CPU 2-bit. +- Accuracy report for uniform and mixed 2-bit models. + +### 6-Bit Track + +- Packing-format proposal. +- Python reference packer and CPU reference implementation. +- 4-bit/6-bit/8-bit quality and size comparison. +- CUDA unpack microbenchmark for contiguous and split-plane layouts. +- Go/no-go recommendation for optimized kernels. + +## Recommendation to Management + +Start implementation with 2-bit CUDA `MatMulNBits` because ONNX Runtime already has the portable format, quantizer, CPU implementation, and test foundation. Treat 6-bit as a separate format-and-value spike because it requires a new cross-provider packing contract and has no native Tensor Core path. + +## References + +- ONNX 2-bit integer types: https://onnx.ai/onnx/technical/int2.html +- ONNX Runtime `MatMulNBits` CUDA documentation: ../contrib_ops/cuda/matmul_nbits.md +- llama.cpp quantization documentation: https://github.com/ggml-org/llama.cpp/blob/master/tools/quantize/README.md +- llama.cpp Q6_K block definition: https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-common.h +- AutoGPTQ: https://github.com/AutoGPTQ/AutoGPTQ +- TorchAO: https://github.com/pytorch/ao +- TensorRT-LLM: https://github.com/NVIDIA/TensorRT-LLM \ No newline at end of file From c8a8520eb6412145aafd7444e8ec2dac39250441 Mon Sep 17 00:00:00 2001 From: David Fan Date: Wed, 16 Sep 2026 20:40:52 +0000 Subject: [PATCH 2/8] Prioritize 2-bit CUDA quantization exploration --- ...it-weight-only-quantization-exploration.md | 175 +++++++++++++----- 1 file changed, 128 insertions(+), 47 deletions(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index 5cc35a80407b4..07172d67bc518 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -2,11 +2,11 @@ ## Executive Summary -ONNX Runtime already has a substantial 2-bit foundation: the `MatMulNBits` model format, Python quantization tooling, CPU kernels, and correctness tests support 2-bit weights. The largest practical gap is CUDA `MatMulNBits`, whose execution and prepacking paths currently support only 4-bit and 8-bit weights. +ONNX Runtime already has a substantial 2-bit foundation: the `MatMulNBits` model format, Python quantization tooling, CPU kernels, and correctness tests support 2-bit weights. This makes INT2 the shortest path to extending CUDA weight-only execution without introducing a new portable format. -The recommended first implementation target is therefore 2-bit CUDA `MatMulNBits`. A correctness-first dequantization plus cuBLAS fallback would establish end-to-end coverage before investing in fused GEMV and small-M kernels. +The recommended first implementation target is therefore 2-bit CUDA `MatMulNBits`. The initial CUDA work should compare direct packed-INT2 execution with a GPU-native LUT approach, then integrate the best path for both M=1 decode and representative large-M prefill. This is required for a GitHub Copilot-style workload: M=1 decode affects token-generation speed, while long-context prefill affects time to first token (TTFT). -Six-bit support should begin as a format and value-validation spike. It is not a small extension of the existing implementation: current packing code assumes that the bit width divides eight, while 6-bit values cross byte boundaries. There is also no native INT6 tensor type or NVIDIA Tensor Core instruction. A portable 6-bit representation and execution-provider prepacking contract should be agreed upon before optimized kernels are implemented. +This choice prioritizes implementation readiness and maximum memory-bandwidth reduction while accepting material model-quality risk. The first week must establish uniform and mixed-precision INT2 quality baselines against INT4, INT8, and BF16. Mixed INT2/INT4/INT8 quantization may be required for sensitive layers. INT6 remains a follow-up option if INT2 cannot meet coding and tool-calling quality targets or if a less aggressive quality/size tradeoff is needed. ## Scope @@ -60,7 +60,20 @@ CUDA `MatMulNBits` currently assumes either 4-bit or 8-bit weights in several pl - Offline and runtime prepacked layouts. - Workspace and persistent-memory estimates. -This means enabling 2-bit requires more than relaxing one validation check. However, the generic dequantize-to-floating-point plus cuBLAS path provides a narrow route to initial correctness without extending CUTLASS prepacked formats. +This means enabling 2-bit requires more than relaxing one validation check. The implementation should preserve the portable model layout while allowing a CUDA-specific runtime prepack. It should execute packed weights directly rather than materializing a floating-point weight matrix. The CPU implementation remains the semantic reference, but its AVX/VNNI and LUT layouts should not be copied literally because CUDA requires different memory coalescing, register, shared-memory, and occupancy tradeoffs. + +### Existing M-Dependent CUDA Dispatch + +The existing INT4 and INT8 `MatMulNBits` CUDA implementations already select different execution paths according to the runtime row count M, although their fused-kernel coverage is not identical: + +| Weight type | M=1 | Small M | Larger M | +| --- | --- | --- | --- | +| INT4 | Dedicated fused GEMV | Fused batched/small-M kernels for M=2-16, subject to data type, shape, block-size, and shared-memory constraints | The FP16/BF16 `fpA_intB` path can select a CUTLASS weight-only GEMM; otherwise execution falls back to dequantization followed by cuBLAS | +| INT8 | Dedicated fused GEMV | Fused batched kernels for M=2-5, subject to shape and block-size constraints | The FP16/BF16 `fpA_intB` path can select a CUTLASS weight-only GEMM; otherwise execution falls back to dequantization followed by cuBLAS | + +For eligible prepacked FP16/BF16 configurations, the `fpA_intB` profiler selects between its CUDA GEMV and CUTLASS GEMM tactics using the actual M bucket rather than a single fixed threshold. Consequently, the table describes the current specialized direct-kernel coverage, not a universal three-way dispatcher that applies to every data type and shape. + +The proposed INT6 implementation should preserve this M-sensitive architecture: a fused GEMV for M=1 decode, a tiled fused GEMM for representative large-M prefill, and eventually a separate small-M path. It cannot simply reuse the existing INT4/INT8 CUTLASS kernels because NVIDIA hardware and the current CUTLASS integration do not expose a native INT6 weight-only operation. INT6 therefore also requires a new lower-4/upper-2 prepack, extraction logic, and fused compute kernels. Unsupported INT6 configurations should retain a correctness fallback, but that fallback is not a performance milestone. ## External Landscape @@ -166,49 +179,116 @@ Before implementation, agree on: - Required block sizes, bias, `g_idx`, and data types. - Accuracy and performance acceptance criteria. -### Phase 1: 2-Bit CUDA Correctness +### Phase 1: 2-Bit Quality and Kernel Gate + +1. Freeze representative Qwen shapes, workloads, quality metrics, and performance baselines. +2. Validate the existing portable INT2 packing and CPU implementation as the CUDA semantic reference. +3. Measure uniform and mixed INT2/INT4/INT8 coding, tool-calling, KL-divergence, and effective model size. +4. Prototype direct packed-INT2 and GPU-native LUT extraction for M=1 decode. +5. Prototype a tiled direct or LUT-based path for representative large-M prefill. +6. Select the CUDA execution and runtime-prepacking strategy using measured quality and performance data. + +The first week is a quality and workload gate, not a stop condition for all INT2 engineering. If uniform INT2 misses model-quality thresholds, the implementation should proceed with a mixed-precision recipe that preserves sensitive layers at INT4, INT8, or BF16. The CUDA kernel decision must be based on end-to-end value rather than unpack throughput alone. + +#### Required Deliverables + +- Uniform and mixed-precision INT2 quality and effective-size results on the agreed Qwen coding-model workload. +- Direct-unpack versus GPU-native LUT microbenchmarks for representative M=1 decode and large-M prefill shapes. +- A selected CUDA runtime-prepacking and kernel strategy, including memory overhead and architecture constraints. +- A written assessment of whether INT2 delivers useful end-to-end decode, TTFT, and memory improvements over INT4. + +#### Go/No-Go Criteria + +- **Quality:** A uniform or mixed-precision INT2 recipe meets agreed coding and tool-calling thresholds. +- **Size:** Effective model size, including higher-precision layers and metadata, remains materially below INT4. +- **CUDA value:** Fused M=1 and large-M prototypes show credible decode and prefill gains over INT4 without becoming dominated by unpacking, LUT, or occupancy costs. +- **Complexity:** Runtime prepacking, kernel coverage, and maintenance cost are justified by end-to-end model improvements. + +### Phase 2: Scoped 2-Bit CUDA Delivery + +Integrate the selected approach into `MatMulNBits(bits=2)` with: -1. Add a CUDA 2-bit blockwise dequantization kernel. -2. Route 2-bit `MatMulNBits` through dequantization plus cuBLAS. -3. Keep CUTLASS and offline-prepacked paths disabled for 2-bit initially. -4. Update workspace and memory estimation for the fallback path. -5. Add CUDA tests for symmetric and asymmetric zero points, bias, tails, and supported block sizes. -6. Add model-level quantization and inference tests using the existing Python packer. +- Native fused M=1 decode GEMV and representative large-M prefill GEMM for FP16 activations, symmetric weights, and one selected block size. +- Runtime prepacking where it provides a measured benefit while preserving the existing portable INT2 model layout. +- Correctness checks against the existing CPU implementation and explicit dequantization. +- Focused Qwen-shape performance tests and end-to-end decode throughput and TTFT measurements. -This phase provides complete functionality and a stable reference for optimized kernels. +### Phase 3: 2-Bit Production Expansion -### Phase 2: 2-Bit CUDA Performance +After the scoped delivery, expand to BF16, additional block sizes, asymmetric zero points, bias, tails, intermediate/small-M execution, offline prepacking, and broader GPU tuning as justified by measured demand. -1. Implement a fused M=1 GEMV path for token decode. -2. Implement or adapt a small-M batched kernel. -3. Evaluate an INT8 activation plus INT2 weight dot-product strategy. -4. Add runtime prepacking if profiling shows that the canonical layout limits load efficiency. -5. Tune dispatch thresholds against the dequantization plus cuBLAS fallback. +### Phase 4: 6-Bit Follow-Up Gate -Pure 2-bit round-to-nearest quantization may not meet model-quality targets. GPTQ, HQQ, K-quant-inspired optimization, importance-aware quantization, and mixed 2-bit/4-bit layer assignment should be evaluated alongside kernel work. +Evaluate INT6 if INT2 cannot achieve the required quality/size tradeoff or product requirements call for a less aggressive quantization option. Reuse the contiguous portable format and lower-4/upper-2 prepacking analysis in this document, but require a separate format review and measured advantage over mixed INT4/INT8 before implementation. -### Phase 3: 6-Bit Format and Value Spike +## Schedule Estimate -1. Write a precise portable packing specification. -2. Implement Python pack and unpack reference functions. -3. Add CPU reference dequantization and correctness tests. -4. Quantize representative models with 4-bit, 6-bit, and 8-bit configurations. -5. Measure quality, model size, load time, and dequantization overhead. -6. Prototype contiguous and split-plane CUDA extraction with a microbenchmark. +These estimates assume one engineer working full time with Copilot assistance, timely access to representative Ampere, Ada, and Hopper GPUs, and reusable ONNX exports for the target Qwen models. They include implementation, profiling, tests, documentation, and normal review fixes, but not unpredictable CI queue or external model-conversion blockers. -The output of this phase should be a go/no-go decision for optimized MLAS and CUDA work. +### 2-Bit CUDA `MatMulNBits` -### Phase 4: 6-Bit Optimized Execution +| Work item | Estimate | +| --- | ---: | +| Qwen shape inventory, quality baselines, benchmark harness, and CPU/reference validation | 1 week | +| Direct and GPU-native LUT M=1 INT2 GEMV prototypes and profiling | 1-2 weeks | +| Direct and LUT-based large-M INT2 GEMM prototypes and profiling | 2 weeks | +| ORT integration for the primary symmetric FP16 configuration and selected block size | 1 week | +| Focused correctness tests, end-to-end measurements, tuning, and PR cleanup | 2-3 weeks | -Proceed only if Phase 3 demonstrates a meaningful quality, memory, or latency niche that is not covered by mixed 4-bit/8-bit quantization. +Allow **7-9 engineering weeks** for the scoped M=1 decode and large-M prefill vertical slice under the primary configuration. Broader data types, block sizes, asymmetric quantization, small-M kernels, and multi-architecture tuning would extend the work beyond this initial delivery. Copilot reduces coding and test-authoring time, but it does not remove hardware profiling, kernel tuning, model-quality evaluation, or code-review time. -Potential work includes: +### 6-Bit `MatMulNBits` -- CPU SIMD dequantization and GEMM integration. -- CUDA fused decode GEMV. -- Small-M CUDA execution. -- Execution-provider-specific prepacking. -- Model conversion and compatibility tests. +| Work item | Estimate | +| --- | ---: | +| Portable format specification and schema/tooling design | 1-2 weeks | +| Python pack/unpack, CPU reference, tests, and 4/6/8-bit quality study | 2-3 weeks | +| CUDA contiguous-versus-split-plane extraction and fused M=1/large-M prototypes | 3-4 weeks | +| Go/no-go analysis and design review | 1 week | + +Allow **6-8 engineering weeks** for the complete 6-bit format-and-value study and an evidence-based final go/no-go decision when both decode and prefill prototypes are required. The format, quality experiments, reference implementation, and scoped CUDA work can overlap to target an earlier vertical slice. A production-ready CUDA 6-bit track with broad operator coverage remains approximately **11-17 engineering weeks total**. Adding optimized CPU/MLAS support or another execution provider would require separate estimates. + +## November 15 Delivery Plan + +There are approximately 8.5 calendar weeks from September 16 to November 15, 2026. With one engineer, the committed delivery should be a scoped INT2 CUDA vertical slice built on the existing portable format, Python tooling, and CPU reference implementation. The target is reviewable native fused CUDA execution for M=1 decode and representative large-M prefill under one primary symmetric FP16 configuration. This is narrower than complete production INT2 support. Upstream merge by November 15 cannot be guaranteed because review and CI timing are outside the implementation owner's control. + +### Committed INT2 Scope + +- Uniform and mixed INT2/INT4/INT8 quality and effective-size results for representative Qwen3.8 coding-model workloads. +- Direct packed-INT2 versus GPU-native LUT prototype results and a selected CUDA strategy. +- CUDA runtime prepacking if justified by profiling, plus fused M=1 GEMV and fused large-M GEMM for FP16 activations, symmetric weights, and one selected block size. +- Correctness and performance results for M=1 decode and representative prefill M values, such as 128, 512, and 2048, on Qwen3.8-27B and Qwen3.8-Flash-Next matrix shapes. +- End-to-end decode throughput and TTFT measurements for a Copilot-style long-context workload. +- ORT integration, focused tests, documentation, and a reviewable pull request or draft pull request, depending on review readiness. + +### Schedule + +| Dates | Milestone | +| --- | --- | +| September 16-20 | Freeze quality thresholds, Qwen workloads, candidate block sizes, and INT4/INT8/BF16 baselines. | +| September 21-27 | Run uniform and mixed-precision INT2 quality/size experiments; validate the existing portable format and CPU reference. | +| September 28-October 11 | Prototype and compare direct packed-INT2 and GPU-native LUT paths for M=1 decode and large-M prefill. | +| October 12-25 | Select the kernel/prepack strategy; implement fused M=1 GEMV and large-M GEMM for the primary FP16 configuration. | +| October 26-November 1 | Integrate both paths into CUDA `MatMulNBits(bits=2)` and validate representative Qwen decode and prefill shapes. | +| November 2-8 | Add focused correctness tests, serialization, memory estimates, build integration, and performance measurements. | +| November 9-15 | Regression testing, documentation, final quality/performance report, PR cleanup, and review buffer. | + +### Stretch Scope + +The following items should not put the November 15 commitment at risk: + +- Asymmetric zero points. +- BF16 activation support. +- Additional block sizes beyond the selected primary configuration. +- Bias and tail combinations not already covered by the selected kernel path. +- Optimized intermediate-M execution for M values between the committed decode and prefill ranges. +- Offline CUDA-specific prepacking. +- Broad multi-GPU tuning. +- INT6 format, tooling, or kernel implementation. + +### INT6 Scheduling Impact + +With the same engineer, INT6 implementation should not run concurrently if it threatens the INT2 deadline. Before November 15, INT6 work should be limited to preserving the format analysis and collecting quality data that directly informs the INT2 comparison. A second engineer could run the INT6 quality and format gate independently. ## Evaluation Plan @@ -259,24 +339,25 @@ Report: ## Recommended Initial Deliverables -### 2-Bit Track +### 2-Bit Priority Track -- CUDA correctness fallback for `MatMulNBits`. -- Focused CUDA and Python tests. -- Decode and prefill benchmark baseline against 4-bit and CPU 2-bit. -- Accuracy report for uniform and mixed 2-bit models. +- Accuracy and effective-size report for uniform and mixed 2-bit/4-bit/8-bit/BF16 coding-model recipes. +- Direct-unpack and GPU-native LUT comparison on Qwen3.8 decode and prefill shapes. +- CUDA M=1 GEMV and representative large-M GEMM vertical slice for the primary symmetric FP16 configuration. +- Qwen-shape decode, prefill, and TTFT correctness/performance report. +- Go/no-go recommendation for broader production investment. -### 6-Bit Track +### 6-Bit Follow-Up Track -- Packing-format proposal. -- Python reference packer and CPU reference implementation. -- 4-bit/6-bit/8-bit quality and size comparison. -- CUDA unpack microbenchmark for contiguous and split-plane layouts. -- Go/no-go recommendation for optimized kernels. +- Preserve the portable contiguous and lower-4/upper-2 prepacking design analysis. +- Run a 4-bit/6-bit/8-bit quality and effective-size comparison if INT2 quality is insufficient. +- Require a separate format and value gate before native CUDA implementation. ## Recommendation to Management -Start implementation with 2-bit CUDA `MatMulNBits` because ONNX Runtime already has the portable format, quantizer, CPU implementation, and test foundation. Treat 6-bit as a separate format-and-value spike because it requires a new cross-provider packing contract and has no native Tensor Core path. +Prioritize INT2 because ONNX Runtime already has a portable model format, quantization tooling, CPU kernels, and correctness coverage, and because INT2 offers the largest potential weight-memory and bandwidth reduction. Use the first week to freeze quality thresholds and identify a viable uniform or mixed-precision recipe, then target a scoped CUDA `MatMulNBits(bits=2)` vertical slice by November 15: direct-versus-LUT evidence, fused M=1 decode, and fused representative large-M prefill for the primary symmetric FP16 configuration. Both execution paths are required for a GitHub Copilot-style workload because decode determines generation speed and prefill determines TTFT for long repository context. + +This November scope is not complete production INT2 support. Broader data types, block sizes, asymmetric quantization, small-M execution, offline prepacking, and multi-architecture tuning remain follow-up work. Keep INT6 as the next quality-oriented option if INT2 cannot meet the agreed coding and tool-calling targets at a meaningful effective-size advantage over INT4. ## References From 8915955b998f78a2f828bbefc4685c7137b1f7a3 Mon Sep 17 00:00:00 2001 From: David Fan Date: Wed, 16 Sep 2026 20:53:13 +0000 Subject: [PATCH 3/8] Fix WebGPU INT2 support documentation --- docs/design/2bit-6bit-weight-only-quantization-exploration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index 07172d67bc518..a5849ed0b609f 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -29,7 +29,7 @@ Activation quantization, KV-cache quantization, and floating-point formats such | Python `MatMulNBits` quantizer | Yes | No | Dedicated 2-bit, 4-bit, and 8-bit native packers are exported. | | CPU `MatMulNBits` | Yes | No | Includes fallback, AVX512/VNNI, and Arm64 work. | | CUDA `MatMulNBits` | No | No | Kernel dispatch, prepacking, and memory estimation accept only 4 and 8 bits. | -| WebGPU `MatMulNBits` | No | No | Accepts only 4 and 8 bits. | +| WebGPU `MatMulNBits` | Yes | No | Supports 2-bit unpacking with symmetric and asymmetric test coverage. | | CPU `GatherBlockQuantized` | Yes | No | Supports packed 2-bit, 4-bit, and 8-bit integer data. | | CUDA `GatherBlockQuantized` | Yes | No | The generic packed-`uint8` path supports 2-bit extraction. | | QDQ quantization helper | Limited | No | The dedicated native QDQ packing helper is currently 4-bit-oriented. | From 6208f754d158993c1e438fbdff2434a16bb19286 Mon Sep 17 00:00:00 2001 From: David Fan Date: Thu, 17 Sep 2026 15:54:11 +0000 Subject: [PATCH 4/8] Document expert INT2 quantization target --- ...it-weight-only-quantization-exploration.md | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index a5849ed0b609f..70e00c3d8ee4b 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -81,6 +81,23 @@ The proposed INT6 implementation should preserve this M-sensitive architecture: llama.cpp has broad execution support for Q2_K, IQ2 variants, and Q6_K across CPU, CUDA, Metal, and additional backends. Its formats are not equivalent to the ONNX Runtime `MatMulNBits` format. +#### Qwen3.8-Flash-Next Published 2-Bit Baseline + +The published Unsloth `UD-Q2_K_XL` GGUF for Qwen3.8-Flash-Next is approximately 78.9 GB, but it is neither uniformly Q2_K nor a stock llama.cpp quantization preset. Its metadata labels the model `MOSTLY_Q2_K` while assigning formats per tensor using an importance matrix with 926 entries derived from 45 calibration chunks. The inspected assignments include: + +- Expert gate/up tensors: `IQ2_XS` in 47 layers and `IQ3_XXS` in layer 2. +- Expert down tensors: `IQ4_NL`. +- Attention projections: primarily `Q5_K` and `Q6_K`. +- Token and output embeddings: `Q5_K` and `Q4_K`. +- N-gram embedding: `IQ4_NL`. +- Hyper-connection matrices: `Q8_0`. +- Sparse-attention indexer Q/K tensors: BF16. +- Norms and small control tensors: primarily F32. + +`IQ2_XS` is a nominal 2-bit importance-quantized format with approximately 2.31 effective bits per weight after its scales and indexing metadata. It is not the same numerical format as blockwise affine `MatMulNBits(bits=2)`. Consequently, "2-bit" describes the lowest and dominant expert tier of this GGUF, not a two-bit average across every parameter. + +The first ONNX Runtime model-level experiment should nevertheless target the same tensor placement: use `MatMulNBits(bits=2)` for the 47 expert gate/up tensor groups, retain layer 2 and the remaining sensitive tensor classes at supported higher precision, and then measure quality and effective model size. This tests whether expert gate/up INT2 captures most of the useful compression and bandwidth reduction without claiming bit-exact equivalence to `IQ2_XS`. Comparisons with `UD-Q2_K_XL` must report tensor-type distribution and effective bits per parameter rather than comparing quantization names alone. + Q6_K stores each 6-bit code as two planes: - Lower four bits: `QK_K / 2` bytes. @@ -183,7 +200,7 @@ Before implementation, agree on: 1. Freeze representative Qwen shapes, workloads, quality metrics, and performance baselines. 2. Validate the existing portable INT2 packing and CPU implementation as the CUDA semantic reference. -3. Measure uniform and mixed INT2/INT4/INT8 coding, tool-calling, KL-divergence, and effective model size. +3. Measure uniform and mixed INT2/INT4/INT8 coding, tool-calling, KL-divergence, and effective model size, starting with `MatMulNBits(bits=2)` on the 47 expert gate/up tensor groups and higher precision elsewhere. 4. Prototype direct packed-INT2 and GPU-native LUT extraction for M=1 decode. 5. Prototype a tiled direct or LUT-based path for representative large-M prefill. 6. Select the CUDA execution and runtime-prepacking strategy using measured quality and performance data. @@ -193,6 +210,7 @@ The first week is a quality and workload gate, not a stop condition for all INT2 #### Required Deliverables - Uniform and mixed-precision INT2 quality and effective-size results on the agreed Qwen coding-model workload. +- An expert gate/up INT2 model variant compared with the approximately 78.9 GB Unsloth `UD-Q2_K_XL` baseline, including tensor-type distribution and effective bits per parameter. - Direct-unpack versus GPU-native LUT microbenchmarks for representative M=1 decode and large-M prefill shapes. - A selected CUDA runtime-prepacking and kernel strategy, including memory overhead and architecture constraints. - A written assessment of whether INT2 delivers useful end-to-end decode, TTFT, and memory improvements over INT4. From f20da3606f9a66cd8fc4a15d1da8b529ef1c27b7 Mon Sep 17 00:00:00 2001 From: David Fan Date: Thu, 17 Sep 2026 15:59:20 +0000 Subject: [PATCH 5/8] Clarify INT2 model production scope --- ...it-weight-only-quantization-exploration.md | 42 +++++++++++++++---- 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index 70e00c3d8ee4b..2c333352412e3 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -17,7 +17,9 @@ This document focuses on weight-only block quantization for `MatMulNBits` and re - Model-format and packing choices. - A staged implementation and evaluation plan. -Activation quantization, KV-cache quantization, and floating-point formats such as FP4 and FP6 are outside the primary scope. +The initial CUDA vertical slice is limited to dense projections represented as individual `MatMulNBits` nodes. Expert gate/up tensors are included only when the model-production path exports them as independent dense nodes. Fused MoE/QMoE operators and their packed expert-weight contracts are a separate workstream. + +Activation quantization, KV-cache quantization, floating-point formats such as FP4 and FP6, and fused MoE/QMoE execution are outside the primary scope. ## Current ONNX Runtime Support @@ -96,7 +98,7 @@ The published Unsloth `UD-Q2_K_XL` GGUF for Qwen3.8-Flash-Next is approximately `IQ2_XS` is a nominal 2-bit importance-quantized format with approximately 2.31 effective bits per weight after its scales and indexing metadata. It is not the same numerical format as blockwise affine `MatMulNBits(bits=2)`. Consequently, "2-bit" describes the lowest and dominant expert tier of this GGUF, not a two-bit average across every parameter. -The first ONNX Runtime model-level experiment should nevertheless target the same tensor placement: use `MatMulNBits(bits=2)` for the 47 expert gate/up tensor groups, retain layer 2 and the remaining sensitive tensor classes at supported higher precision, and then measure quality and effective model size. This tests whether expert gate/up INT2 captures most of the useful compression and bandwidth reduction without claiming bit-exact equivalence to `IQ2_XS`. Comparisons with `UD-Q2_K_XL` must report tensor-type distribution and effective bits per parameter rather than comparing quantization names alone. +The first ONNX Runtime model-level experiment should nevertheless target the same tensor placement: use `MatMulNBits(bits=2)` for the 47 expert gate/up tensor groups when they are exported as independent dense projections, retain layer 2 and the remaining sensitive tensor classes at supported higher precision, and then measure quality and effective model size. This tests whether expert gate/up INT2 captures most of the useful compression and bandwidth reduction without claiming bit-exact equivalence to `IQ2_XS`. Comparisons with `UD-Q2_K_XL` must report tensor-type distribution and effective bits per parameter rather than comparing quantization names alone. Q6_K stores each 6-bit code as two planes: @@ -196,14 +198,33 @@ Before implementation, agree on: - Required block sizes, bias, `g_idx`, and data types. - Accuracy and performance acceptance criteria. +### Initial Model-Production Contract + +The required model-production workflow for the initial CUDA vertical slice is: + +```text +HF/PyTorch model + -> Olive RTN, GPTQ, or SelectiveMixedPrecision INT2/mixed checkpoint + -> Mobius export + -> dense MatMulNBits(bits=2/4/8) ONNX model + -> ONNX Runtime CUDA execution +``` + +This path best matches the GPTQ, selective mixed-precision, and Qwen quality goals. The vertical slice must qualify graph conversion, initializer binding, tensor-wise bit-width selection, and numerical parity from the Olive checkpoint through Mobius export and ONNX Runtime execution. + +The ONNX-native workflow, `FP ONNX -> Olive OnnxBlockWiseRtnQuantization -> MatMulNBits(bits=2)`, remains required for broader product support but follows the initial PyTorch/Mobius milestone. Complete Olive ONNX RTN and built-in quantized-linear INT2 export support must be tracked with the owning Olive work rather than assumed to exist because the ONNX Runtime kernel is available. + +The initial contract produces dense `MatMulNBits` nodes only. Exporting expert tensors into a fused MoE/QMoE operator would require a separate Olive/Mobius/runtime schema, packing, weight-binding, kernel, and parity contract and is not part of this CUDA delivery. + ### Phase 1: 2-Bit Quality and Kernel Gate 1. Freeze representative Qwen shapes, workloads, quality metrics, and performance baselines. 2. Validate the existing portable INT2 packing and CPU implementation as the CUDA semantic reference. -3. Measure uniform and mixed INT2/INT4/INT8 coding, tool-calling, KL-divergence, and effective model size, starting with `MatMulNBits(bits=2)` on the 47 expert gate/up tensor groups and higher precision elsewhere. -4. Prototype direct packed-INT2 and GPU-native LUT extraction for M=1 decode. -5. Prototype a tiled direct or LUT-based path for representative large-M prefill. -6. Select the CUDA execution and runtime-prepacking strategy using measured quality and performance data. +3. Measure uniform and mixed INT2/INT4/INT8 coding, tool-calling, KL-divergence, and effective model size, starting with `MatMulNBits(bits=2)` on independently exported expert gate/up projections and higher precision elsewhere. +4. Qualify the Olive checkpoint-to-Mobius-to-ONNX path for graph structure, initializer binding, tensor-wise bit widths, and numerical parity. +5. Prototype direct packed-INT2 and GPU-native LUT extraction for M=1 decode. +6. Prototype a tiled direct or LUT-based path for representative large-M prefill. +7. Select the CUDA execution and runtime-prepacking strategy using measured quality and performance data. The first week is a quality and workload gate, not a stop condition for all INT2 engineering. If uniform INT2 misses model-quality thresholds, the implementation should proceed with a mixed-precision recipe that preserves sensitive layers at INT4, INT8, or BF16. The CUDA kernel decision must be based on end-to-end value rather than unpack throughput alone. @@ -211,6 +232,7 @@ The first week is a quality and workload gate, not a stop condition for all INT2 - Uniform and mixed-precision INT2 quality and effective-size results on the agreed Qwen coding-model workload. - An expert gate/up INT2 model variant compared with the approximately 78.9 GB Unsloth `UD-Q2_K_XL` baseline, including tensor-type distribution and effective bits per parameter. +- An end-to-end mixed INT2/INT4/INT8 model exported through Olive and Mobius with graph, weight-binding, and numerical-parity coverage. - Direct-unpack versus GPU-native LUT microbenchmarks for representative M=1 decode and large-M prefill shapes. - A selected CUDA runtime-prepacking and kernel strategy, including memory overhead and architecture constraints. - A written assessment of whether INT2 delivers useful end-to-end decode, TTFT, and memory improvements over INT4. @@ -230,6 +252,7 @@ Integrate the selected approach into `MatMulNBits(bits=2)` with: - Runtime prepacking where it provides a measured benefit while preserving the existing portable INT2 model layout. - Correctness checks against the existing CPU implementation and explicit dequantization. - Focused Qwen-shape performance tests and end-to-end decode throughput and TTFT measurements. +- A deployable dense `MatMulNBits` model produced through the qualified PyTorch/Olive-to-Mobius workflow. ### Phase 3: 2-Bit Production Expansion @@ -273,6 +296,7 @@ There are approximately 8.5 calendar weeks from September 16 to November 15, 202 ### Committed INT2 Scope - Uniform and mixed INT2/INT4/INT8 quality and effective-size results for representative Qwen3.8 coding-model workloads. +- A qualified PyTorch/Olive-to-Mobius export path producing dense mixed-bit `MatMulNBits` nodes, with graph, initializer-binding, and numerical-parity tests. - Direct packed-INT2 versus GPU-native LUT prototype results and a selected CUDA strategy. - CUDA runtime prepacking if justified by profiling, plus fused M=1 GEMV and fused large-M GEMM for FP16 activations, symmetric weights, and one selected block size. - Correctness and performance results for M=1 decode and representative prefill M values, such as 128, 512, and 2048, on Qwen3.8-27B and Qwen3.8-Flash-Next matrix shapes. @@ -284,8 +308,8 @@ There are approximately 8.5 calendar weeks from September 16 to November 15, 202 | Dates | Milestone | | --- | --- | | September 16-20 | Freeze quality thresholds, Qwen workloads, candidate block sizes, and INT4/INT8/BF16 baselines. | -| September 21-27 | Run uniform and mixed-precision INT2 quality/size experiments; validate the existing portable format and CPU reference. | -| September 28-October 11 | Prototype and compare direct packed-INT2 and GPU-native LUT paths for M=1 decode and large-M prefill. | +| September 21-27 | Run uniform and mixed-precision INT2 quality/size experiments; validate the portable format and qualify Olive checkpoint-to-Mobius graph and weight export. | +| September 28-October 11 | Complete export numerical-parity coverage; prototype and compare direct packed-INT2 and GPU-native LUT paths for M=1 decode and large-M prefill. | | October 12-25 | Select the kernel/prepack strategy; implement fused M=1 GEMV and large-M GEMM for the primary FP16 configuration. | | October 26-November 1 | Integrate both paths into CUDA `MatMulNBits(bits=2)` and validate representative Qwen decode and prefill shapes. | | November 2-8 | Add focused correctness tests, serialization, memory estimates, build integration, and performance measurements. | @@ -303,6 +327,8 @@ The following items should not put the November 15 commitment at risk: - Offline CUDA-specific prepacking. - Broad multi-GPU tuning. - INT6 format, tooling, or kernel implementation. +- ONNX-native Olive INT2 RTN and built-in quantized-linear export completion. +- Fused MoE/QMoE export, packing contracts, and runtime kernels. ### INT6 Scheduling Impact From e7517297d85520f6c0992503a762cf71e056da7f Mon Sep 17 00:00:00 2001 From: David Fan Date: Thu, 17 Sep 2026 17:47:25 +0000 Subject: [PATCH 6/8] Add INT2 QMoE follow-up assessment --- ...it-weight-only-quantization-exploration.md | 34 ++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index 2c333352412e3..c489459268336 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -77,6 +77,24 @@ For eligible prepacked FP16/BF16 configurations, the `fpA_intB` profiler selects The proposed INT6 implementation should preserve this M-sensitive architecture: a fused GEMV for M=1 decode, a tiled fused GEMM for representative large-M prefill, and eventually a separate small-M path. It cannot simply reuse the existing INT4/INT8 CUTLASS kernels because NVIDIA hardware and the current CUTLASS integration do not expose a native INT6 weight-only operation. INT6 therefore also requires a new lower-4/upper-2 prepack, extraction logic, and fused compute kernels. Unsupported INT6 configurations should retain a correctness fallback, but that fallback is not a performance milestone. +### INT2 QMoE Follow-Up Assessment + +INT2 QMoE is a meaningful follow-up because expert weights dominate the storage and memory traffic of large MoE models, and a fused operator preserves top-k routing instead of expanding every expert into independently scheduled dense operations. CPU QMoE already accepts blockwise INT2 and includes an MLAS LUT GEMM path, so it can provide a semantic reference. However, the current cross-provider implementation is not production-ready: + +| Area | Current state | Required work | +| --- | --- | --- | +| QMoE schema | Declares `expert_weight_bits` values 2, 4, and 8 | Define a versioned mixed-width contract before targeting the published Qwen recipe | +| CPU QMoE | Accepts blockwise INT2 and has an optimized LUT path | Add mixed-width semantics and model-level conformance coverage | +| CUDA QMoE | Constructor, packing, runner selection, and weight preprocessing assume INT4 or INT8 | Add explicit INT2 fallback and packed kernels without mapping INT2 to an INT4 CUTLASS type | +| WebGPU QMoE | Rejects INT2 and uses a 4/8-bit-specific pack-size calculation | Use `8 / bits`, complete reachable INT2 shader support, and add QMoE tests | +| Model production | Dense mixed-bit export is the initial Olive/Mobius target | Define and qualify a distinct fused QMoE graph and weight-binding contract | + +The largest model-contract blocker is that QMoE currently exposes one `expert_weight_bits` attribute for FC1 and FC2. This cannot represent the relevant Qwen3.8-Flash-Next placement, where expert gate/up tensors use an approximately 2-bit tier while expert down tensors use an approximately 4-bit tier. A production-oriented extension therefore needs independent FC1 and FC2 bit-width semantics, for example `fc1_expert_weight_bits=2` and `fc2_expert_weight_bits=4`, together with corresponding shape, scale, zero-point, prepacking, and backward-compatibility rules. The exact schema design requires review; these names are illustrative rather than a committed interface. + +CUDA dequantization to persistent FP16/BF16 expert weights is useful only as a correctness oracle because it expands INT2 payloads by approximately 8x and removes the deployment memory benefit. The first performance-relevant QMoE target should be packed INT2 fused decode for small expanded-row counts. Long-context prefill ultimately requires a native or equivalently bounded W2A16 grouped GEMM; full expert dequantization is not a production milestone. + +MatMulNBits remains the first implementation priority because its portable INT2 contract, CPU reference, quantization tooling, and isolated matrix tests reduce risk. Once its CUDA packing interpretation, correctness path, and fused M=1 load/dequantization primitives are stable, QMoE contract and export work can proceed in parallel rather than waiting for complete MatMulNBits coverage across every data type and block size. + ## External Landscape ### llama.cpp @@ -258,7 +276,19 @@ Integrate the selected approach into `MatMulNBits(bits=2)` with: After the scoped delivery, expand to BF16, additional block sizes, asymmetric zero points, bias, tails, intermediate/small-M execution, offline prepacking, and broader GPU tuning as justified by measured demand. -### Phase 4: 6-Bit Follow-Up Gate +### Phase 4: Mixed-Width INT2 QMoE Follow-Up + +Start the QMoE work after the MatMulNBits INT2 packing contract, CUDA correctness behavior, and fused M=1 primitives are stable: + +1. Define independent FC1 and FC2 bit-width semantics and preserve compatibility with the existing single-width QMoE contract. +2. Add CPU correctness and model-level tests for FC1 INT2 with FC2 INT4. +3. Qualify Olive/Mobius fused QMoE export, graph binding, and numerical parity separately from dense `MatMulNBits` export. +4. Implement CUDA packed INT2 QMoE decode by reusing validated INT2 extraction and dequantization primitives. +5. Implement or evaluate a bounded native W2A16 grouped-GEMM path for prefill. + +WebGPU QMoE enablement is an independent follow-up and should not block the CUDA model-level milestone. A cross-provider correctness fallback may be useful for conformance, but persistent full expert dequantization does not satisfy the production acceptance criteria. + +### Phase 5: 6-Bit Follow-Up Gate Evaluate INT6 if INT2 cannot achieve the required quality/size tradeoff or product requirements call for a less aggressive quantization option. Reuse the contiguous portable format and lower-4/upper-2 prepacking analysis in this document, but require a separate format review and measured advantage over mixed INT4/INT8 before implementation. @@ -330,6 +360,8 @@ The following items should not put the November 15 commitment at risk: - ONNX-native Olive INT2 RTN and built-in quantized-linear export completion. - Fused MoE/QMoE export, packing contracts, and runtime kernels. +These QMoE items remain outside the November commitment, but the next implementation track should begin with the mixed FC1/FC2 contract and export qualification after the MatMulNBits INT2 foundation is stable. + ### INT6 Scheduling Impact With the same engineer, INT6 implementation should not run concurrently if it threatens the INT2 deadline. Before November 15, INT6 work should be limited to preserving the format analysis and collecting quality data that directly informs the INT2 comparison. A second engineer could run the INT6 quality and format gate independently. From 19a723d6dd26828b3c731168226dc5aa8ba1bcd1 Mon Sep 17 00:00:00 2001 From: David Fan Date: Thu, 17 Sep 2026 19:23:30 +0000 Subject: [PATCH 7/8] Add INT2 QMoE end-to-end delivery plan --- ...it-weight-only-quantization-exploration.md | 6 +- .../int2-qmoe-end-to-end-delivery-plan.md | 323 ++++++++++++++++++ 2 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 docs/design/int2-qmoe-end-to-end-delivery-plan.md diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index c489459268336..4fc7e2f08e06b 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -4,7 +4,7 @@ ONNX Runtime already has a substantial 2-bit foundation: the `MatMulNBits` model format, Python quantization tooling, CPU kernels, and correctness tests support 2-bit weights. This makes INT2 the shortest path to extending CUDA weight-only execution without introducing a new portable format. -The recommended first implementation target is therefore 2-bit CUDA `MatMulNBits`. The initial CUDA work should compare direct packed-INT2 execution with a GPU-native LUT approach, then integrate the best path for both M=1 decode and representative large-M prefill. This is required for a GitHub Copilot-style workload: M=1 decode affects token-generation speed, while long-context prefill affects time to first token (TTFT). +The current product priority is mixed-width INT2 QMoE for expert-heavy Qwen models: INT2 expert gate/up projections with INT4 expert down projections. The end-to-end contract, export, runtime, validation, and staffing plan is defined in [INT2 QMoE End-to-End Delivery Plan](int2-qmoe-end-to-end-delivery-plan.md). Dense CUDA `MatMulNBits(bits=2)` remains valuable as a source of reusable packing, validation, and kernel primitives, but complete dense GEMV/GEMM delivery should not block the QMoE path. This choice prioritizes implementation readiness and maximum memory-bandwidth reduction while accepting material model-quality risk. The first week must establish uniform and mixed-precision INT2 quality baselines against INT4, INT8, and BF16. Mixed INT2/INT4/INT8 quantization may be required for sensitive layers. INT6 remains a follow-up option if INT2 cannot meet coding and tool-calling quality targets or if a less aggressive quality/size tradeoff is needed. @@ -17,7 +17,7 @@ This document focuses on weight-only block quantization for `MatMulNBits` and re - Model-format and packing choices. - A staged implementation and evaluation plan. -The initial CUDA vertical slice is limited to dense projections represented as individual `MatMulNBits` nodes. Expert gate/up tensors are included only when the model-production path exports them as independent dense nodes. Fused MoE/QMoE operators and their packed expert-weight contracts are a separate workstream. +The dense CUDA vertical slice described in this exploration remains a supporting implementation option for independently exported projections. The prioritized fused QMoE work and its packed expert-weight contract are scoped separately in the end-to-end QMoE delivery plan. Activation quantization, KV-cache quantization, floating-point formats such as FP4 and FP6, and fused MoE/QMoE execution are outside the primary scope. @@ -93,7 +93,7 @@ The largest model-contract blocker is that QMoE currently exposes one `expert_we CUDA dequantization to persistent FP16/BF16 expert weights is useful only as a correctness oracle because it expands INT2 payloads by approximately 8x and removes the deployment memory benefit. The first performance-relevant QMoE target should be packed INT2 fused decode for small expanded-row counts. Long-context prefill ultimately requires a native or equivalently bounded W2A16 grouped GEMM; full expert dequantization is not a production milestone. -MatMulNBits remains the first implementation priority because its portable INT2 contract, CPU reference, quantization tooling, and isolated matrix tests reduce risk. Once its CUDA packing interpretation, correctness path, and fused M=1 load/dequantization primitives are stable, QMoE contract and export work can proceed in parallel rather than waiting for complete MatMulNBits coverage across every data type and block size. +Shared MatMulNBits INT2 packing, CPU-reference, validation, and CUDA load/dequantization primitives remain early enabling work because they reduce QMoE implementation risk. QMoE contract and export work should proceed in parallel and must not wait for complete dense MatMulNBits coverage across every data type, block size, or M regime. ## External Landscape diff --git a/docs/design/int2-qmoe-end-to-end-delivery-plan.md b/docs/design/int2-qmoe-end-to-end-delivery-plan.md new file mode 100644 index 0000000000000..b6dc09ad0df25 --- /dev/null +++ b/docs/design/int2-qmoe-end-to-end-delivery-plan.md @@ -0,0 +1,323 @@ +# INT2 QMoE End-to-End Delivery Plan + +## Executive Summary + +The priority is an end-to-end deployable INT2 QMoE path for Qwen3.8-Flash-Next-class models. The target recipe uses INT2 for expert gate/up projections and INT4 for expert down projections, while sensitive non-expert tensors remain at higher precision. This follows the tensor placement of the published Unsloth `UD-Q2_K_XL` model without claiming numerical compatibility with its `IQ2_XS` and `IQ4_NL` formats. + +The delivery is complete only when a supported source checkpoint can be quantized, exported to ONNX, loaded by ONNX Runtime, and executed by CUDA with demonstrated model quality, memory reduction, decode throughput, and prefill TTFT. Operator-only correctness is necessary but is not an end-to-end deliverable. + +Dense INT2 `MatMulNBits` optimization is no longer the primary product milestone. QMoE should reuse provider-neutral INT2 packing, dequantization, validation, and CUDA load primitives where practical, but completion of dense fused GEMV/GEMM coverage must not block the QMoE contract, export, or packed expert-kernel work. + +The critical path is: + +```text +Mixed-width QMoE contract + -> CPU reference and deterministic tests + -> Olive quantization and Mobius export + -> CUDA correctness path + -> packed CUDA decode kernel + -> bounded or native CUDA prefill path + -> Qwen model validation and release evidence +``` + +## Product Target + +### Initial Quantization Recipe + +| Tensor class | Initial ONNX target | Rationale | +| --- | --- | --- | +| Expert gate/up, QMoE FC1 | Blockwise affine INT2 | Dominant low-bit tier and primary storage/bandwidth target | +| Expert down, QMoE FC2 | Blockwise affine INT4 | Lower quality risk than INT2 for the down projection | +| Attention projections | INT4, INT8, or BF16 according to measured quality | Outside the first QMoE kernel contract | +| Embeddings, indexers, norms, and control tensors | Existing supported higher precision | Preserve quality and avoid unrelated kernel work | + +The first supported configuration is: + +- CUDA execution provider. +- FP16 activations with FP32 accumulation where required for accuracy. +- Symmetric blockwise integer quantization. +- FC1 INT2 and FC2 INT4. +- One selected block size, chosen from 32, 64, or 128 after quality and kernel profiling. +- Interleaved fused SwiGLU (`swiglu_fusion=1`). +- Top-k routing with multiple tokens and experts. +- Raw portable model weights plus execution-provider-specific runtime prepacking. + +BF16, asymmetric zero points, additional block sizes, WebGPU, and alternative SwiGLU layouts are follow-up coverage. + +### Format Boundary + +ORT blockwise INT2 uses four affine integer codes per byte: + +```text +packed = q0 | (q1 << 2) | (q2 << 4) | (q3 << 6) +dequantized = (q - zero_point) * scale +``` + +This is not `IQ2_XS`. The published Qwen GGUF is a comparison point for tensor placement, quality, and effective model size, not a bit-compatible input format. Conversion must quantize from the source floating-point checkpoint or another representation that preserves sufficient information; it must not relabel IQ2_XS bytes as `MatMulNBits` or QMoE INT2 bytes. + +## Required Operator Contract + +### Mixed FC1 and FC2 Widths + +QMoE currently has one `expert_weight_bits` attribute shared by FC1 and FC2. That cannot represent FC1 INT2 with FC2 INT4. The schema must support independent bit widths while preserving existing models. + +One candidate contract is: + +```text +expert_weight_bits # existing default for backward compatibility +fc1_expert_weight_bits # optional override +fc2_expert_weight_bits # optional override +fc3_expert_weight_bits # optional override when FC3 is present +``` + +The names are provisional and require schema review. Semantics must include: + +1. An omitted FC-specific value inherits `expert_weight_bits`. +2. Each effective width is restricted to a supported value. +3. Each weight input uses its own `pack_size = 8 / effective_bits`. +4. FC-specific scales, zero points, shape checks, strides, and prepacked layouts use the corresponding effective width. +5. Existing single-width INT4 and INT8 models remain byte-for-byte compatible. +6. Prepacked layouts encode their bit width and layout version; INT2 must never be interpreted as INT4. + +### Initial Shapes + +For $E$ experts, hidden size $H$, and intermediate size $I$: + +```text +FC1 logical: [E, 2 * I, H] +FC1 INT2: [E, 2 * I, H / 4] + +FC2 logical: [E, H, I] +FC2 INT4: [E, H, I / 2] +``` + +Scale shapes remain per output row and K-axis block. Validation must use independent FC1 and FC2 bit widths and pack sizes. + +## End-to-End Model Production + +### Required Workflow + +```text +HF/PyTorch Qwen checkpoint + -> Olive calibration and tensor-wise mixed quantization + -> mixed FC1 INT2 / FC2 INT4 checkpoint + -> Mobius fused QMoE export + -> ONNX model with portable raw expert weights + -> ONNX Runtime CUDA runtime prepack + -> QMoE inference +``` + +The Olive/Mobius path must provide: + +- Stable identification of expert gate/up and down tensors. +- Per-tensor quantization decisions and a persisted recipe manifest. +- Deterministic INT2 and INT4 packing. +- Correct FC1 interleaving for fused SwiGLU. +- Correct scales, optional zero points, attributes, and initializer bindings. +- External-data support for large models. +- Graph validation that no selected expert tensor silently remains unquantized. +- Numerical parity between the exported graph and the quantized PyTorch reference. + +The recipe manifest should report tensor names, logical shapes, selected widths, block sizes, scale types, packed bytes, and fallback precision. Effective bits per parameter must include scales, padding, metadata, and higher-precision tensors. + +## Implementation Workstreams + +### Workstream 1: Contract and Shared Validation + +Deliverables: + +- Approve the mixed-width QMoE schema design. +- Centralize bit-width, pack-size, default-zero-point, and packed-shape calculations. +- Update FC1, FC2, and FC3 validation to use independent effective widths. +- Define raw and provider-prepacked layout behavior. +- Add schema inference, invalid-shape, inheritance, and backward-compatibility tests. + +Exit gate: one raw mixed FC1 INT2 / FC2 INT4 model validates identically on all registered QMoE providers, even where execution returns a clear unsupported-status error. + +### Workstream 2: CPU Reference + +CPU is the semantic oracle for model production and CUDA validation. + +Deliverables: + +- Extend CPU QMoE from one shared width to independent FC widths. +- Reuse the existing INT2 MLAS LUT path for FC1 where eligible. +- Preserve the existing INT4 path for FC2. +- Cover routing, fused SwiGLU, bias, empty experts, multiple tokens, and top-k greater than one. +- Compare against a scalar explicit-dequantization reference using all INT2 codes 0, 1, 2, and 3. + +Exit gate: deterministic mixed-width CPU QMoE tests pass and the exported model matches the quantized PyTorch reference within agreed tolerances. + +### Workstream 3: Olive and Mobius + +Deliverables: + +- Add or qualify selective mixed-precision QMoE quantization. +- Export the approved FC-specific bit-width contract. +- Emit portable raw `[E, N, K / pack_size]` weights. +- Validate FC1 gate/up interleaving and FC2 orientation. +- Add graph, initializer-binding, external-data, and numerical-parity tests. +- Produce a small checked-in synthetic model and a reproducible Qwen conversion command. + +Exit gate: a clean environment can convert the selected checkpoint and run the exported model on the CPU reference path without manual graph edits. + +### Workstream 4: CUDA Correctness + +The correctness path must not reuse an INT4 type or layout for INT2. + +Deliverables: + +- Accept the approved FC1 INT2 / FC2 INT4 contract in CUDA QMoE. +- Explicitly exclude INT2 from existing INT4/INT8 CUTLASS preprocessing and tactic selection. +- Implement bit-exact INT2 dequantization with block scales and correct row-local packed zero-point addressing. +- Execute FC1 through a bounded dequantization plus existing dense MoE runner for correctness. +- Keep FC2 on its existing INT4 path where possible. +- Add CPU/CUDA parity tests. + +Persistent full-model INT2-to-FP16 dequantization expands the affected payload by approximately 8x. It is acceptable only for small tests and as an oracle. The test implementation must have an explicit memory bound and must not be presented as production support. + +Exit gate: CUDA executes synthetic and reduced Qwen mixed-width models correctly without interpreting INT2 as INT4 or allocating unbounded scratch. + +### Workstream 5: CUDA Packed Decode + +The first performance milestone is fused packed execution for decode and low expanded-row counts, where expanded rows are approximately `num_tokens * top_k`. + +Deliverables: + +- Define a versioned CUDA INT2 runtime-prepacked layout. +- Implement vectorized packed INT2 loads, extraction, scale application, and FP32 accumulation. +- Integrate packed FC1 INT2 with routing and interleaved SwiGLU. +- Reuse or preserve the optimized FC2 INT4 path. +- Support the selected block size and representative Qwen dimensions. +- Profile register pressure, occupancy, effective bandwidth, and prepack cost. +- Add fused-versus-reference correctness tests and M=1 end-to-end benchmarks. + +Exit gate: packed QMoE improves decode latency or throughput over the agreed INT4 baseline while preserving the accepted model quality. + +### Workstream 6: CUDA Prefill + +Long-context prefill must not dequantize every expert into persistent FP16 storage. + +Evaluate in this order: + +1. Bounded selected-expert or row-tiled dequantization for functional integration. +2. Chunked execution with a documented scratch cap. +3. Native W2A16 grouped GEMM for the production performance target. + +A native grouped path requires packed INT2 iterators, converters, block-scale loading, grouped expert pointers and strides, tactic selection, and architecture-specific tuning. + +Exit gate: representative prefill values such as M=128, 512, and 2048 complete within the memory budget and improve TTFT or provide an explicitly accepted intermediate baseline. + +### Workstream 7: Model Quality and Performance + +Quality comparisons must include: + +- BF16 baseline. +- Existing INT4 baseline. +- Uniform expert INT2, for diagnosis only. +- Mixed FC1 INT2 / FC2 INT4. +- The published approximately 78.9 GB `UD-Q2_K_XL` result when reproducible metrics are available. + +Measure coding, tool-calling, long-generation stability, perplexity or KL divergence, and task-specific acceptance metrics. Performance reporting must include model size, peak and persistent memory, load/prepack time, M=1 latency, tokens per second, prefill TTFT, and effective memory bandwidth. + +## Delivery Sequence + +### PR 1: Mixed-Width Contract + +- Schema and inheritance semantics. +- Shared shape and packing helpers. +- Backward-compatibility and validation tests. + +### PR 2: CPU Mixed QMoE + +- FC1 INT2 and FC2 INT4 execution. +- Scalar and MLAS parity tests. +- Routing and SwiGLU coverage. + +### PR 3: Olive/Mobius Export + +- Selective recipe and manifest. +- Fused QMoE graph export. +- Weight-binding and numerical-parity tests. + +PR 2 and PR 3 can proceed in parallel after the contract is approved. + +### PR 4: CUDA Correctness + +- INT2 validation and dequantization. +- Bounded fallback. +- CPU/CUDA parity tests. + +### PR 5: CUDA Packed Decode + +- Runtime prepack. +- Fused FC1 INT2 decode. +- FC2 INT4 integration. +- Correctness and decode benchmarks. + +### PR 6: CUDA Prefill + +- Bounded integration path. +- Native grouped GEMM when required by the performance gate. +- TTFT and memory benchmarks. + +### PR 7: End-to-End Qualification + +- Full conversion and inference automation. +- Quality and performance report. +- Documentation and support matrix updates. + +## Schedule and Staffing + +With one engineer, a production-quality mixed-width QMoE path spanning schema, tooling, CPU, CUDA decode, CUDA prefill, and model qualification is not a credible single eight-week task. A practical schedule uses parallel owners: + +| Weeks | Contract/CPU owner | Olive/Mobius owner | CUDA owner | Model-validation owner | +| --- | --- | --- | --- | --- | +| 1-2 | Freeze contract and add validation | Prototype recipe/export against draft contract | Prototype INT2 extraction and layouts | Freeze models, metrics, and baselines | +| 3-4 | Complete CPU mixed-width reference | Complete graph and parity tests | Complete bounded correctness path | Run first quality comparison | +| 5-6 | Support integration fixes | Produce full external-data model | Implement and tune packed decode | Validate decode quality and memory | +| 7-8 | Regression and compatibility tests | Reproducible conversion package | Integrate bounded prefill; prototype grouped GEMM | End-to-end decode and TTFT report | +| 9-12 | Follow-up coverage | Export hardening | Native grouped-GEMM tuning and architecture coverage | Release qualification | + +An eight-week milestone should commit to a reviewable mixed-width contract, reproducible export, CPU reference, CUDA correctness, and packed decode on one GPU architecture. Production prefill performance and broad architecture coverage are follow-up commitments unless additional CUDA staffing is assigned. + +## Acceptance Criteria + +### Functional Completion + +- A documented command converts the selected Qwen checkpoint without manual graph edits. +- The ONNX graph contains FC1 INT2 and FC2 INT4 QMoE weights with correct shapes and metadata. +- CPU and CUDA outputs match the quantized PyTorch reference within frozen tolerances. +- Routing, top-k, fused SwiGLU, multiple tokens, empty experts, and bias behave correctly. +- Existing single-width INT4 and INT8 QMoE models do not regress. +- Failures for unsupported configurations are explicit and occur before unsafe preprocessing. + +### Product Completion + +- The mixed recipe meets frozen coding and tool-calling quality thresholds. +- Effective model size is materially below the INT4 baseline after all metadata and higher-precision tensors are counted. +- CUDA executes packed INT2 FC1 without persistent full expert dequantization. +- Decode demonstrates an accepted improvement over INT4 on representative Qwen workloads. +- Prefill meets the agreed TTFT and peak-memory targets using a bounded or native path. +- Conversion, model loading, prepacking, inference, and evaluation are automated in CI or a reproducible qualification pipeline. + +## Risks and Stop Conditions + +| Risk | Mitigation or stop condition | +| --- | --- | +| Affine INT2 quality is materially below IQ2_XS | Improve calibration or mixed precision; stop kernel expansion if no viable recipe exists | +| Mixed-width schema creates excessive compatibility cost | Evaluate separate FC attributes versus a versioned quantization descriptor before implementation | +| INT2 unpacking removes decode benefit | Compare direct extraction and LUT/prepacked layouts before committing to one kernel family | +| Prefill activates most experts | Prioritize native grouped GEMM; do not rely on selected-expert dequantization as the final solution | +| Runtime prepack increases load time or memory | Report persistent bytes and prepack latency; support offline versioned packing only after the layout stabilizes | +| Export and runtime contracts diverge | Use one deterministic golden model across Olive, Mobius, CPU, and CUDA tests | + +Stop production kernel expansion if no mixed affine INT2 recipe meets the frozen quality and effective-size gates. A correctness implementation alone does not justify shipping large-model QMoE INT2 support. + +## Immediate Decisions Needed + +1. Select the schema approach for independent FC1 and FC2 widths. +2. Select the first block size and CUDA architecture. +3. Assign owners for ORT schema/CPU, Olive/Mobius, CUDA, and model validation. +4. Freeze quality, memory, decode, and TTFT thresholds before collecting candidate results. +5. Decide whether native grouped prefill is part of the first committed date or a follow-up performance milestone. \ No newline at end of file From 929ee534a4ec304b9801863f83a783b71d1dc1a6 Mon Sep 17 00:00:00 2001 From: David Fan Date: Thu, 17 Sep 2026 19:31:39 +0000 Subject: [PATCH 8/8] Extend low-bit MatMulNBits exploration --- ...it-weight-only-quantization-exploration.md | 116 ++++++++++-------- 1 file changed, 66 insertions(+), 50 deletions(-) diff --git a/docs/design/2bit-6bit-weight-only-quantization-exploration.md b/docs/design/2bit-6bit-weight-only-quantization-exploration.md index 4fc7e2f08e06b..62002c5642e5b 100644 --- a/docs/design/2bit-6bit-weight-only-quantization-exploration.md +++ b/docs/design/2bit-6bit-weight-only-quantization-exploration.md @@ -1,4 +1,4 @@ -# 2-Bit and 6-Bit Weight-Only Quantization Exploration +# 2-, 3-, 5-, and 6-Bit Weight-Only Quantization Exploration ## Executive Summary @@ -6,7 +6,7 @@ ONNX Runtime already has a substantial 2-bit foundation: the `MatMulNBits` model The current product priority is mixed-width INT2 QMoE for expert-heavy Qwen models: INT2 expert gate/up projections with INT4 expert down projections. The end-to-end contract, export, runtime, validation, and staffing plan is defined in [INT2 QMoE End-to-End Delivery Plan](int2-qmoe-end-to-end-delivery-plan.md). Dense CUDA `MatMulNBits(bits=2)` remains valuable as a source of reusable packing, validation, and kernel primitives, but complete dense GEMV/GEMM delivery should not block the QMoE path. -This choice prioritizes implementation readiness and maximum memory-bandwidth reduction while accepting material model-quality risk. The first week must establish uniform and mixed-precision INT2 quality baselines against INT4, INT8, and BF16. Mixed INT2/INT4/INT8 quantization may be required for sensitive layers. INT6 remains a follow-up option if INT2 cannot meet coding and tool-calling quality targets or if a less aggressive quality/size tradeoff is needed. +This choice prioritizes implementation readiness and maximum memory-bandwidth reduction while accepting material model-quality risk. The first week must establish uniform and mixed-precision INT2 quality baselines against INT4, INT8, and BF16. Mixed INT2/INT4/INT8 quantization may be required for sensitive layers. INT3, INT5, and INT6 remain follow-up quality/size options if INT2 cannot meet coding and tool-calling targets or if a less aggressive quantization point is needed. ## Scope @@ -25,16 +25,16 @@ Activation quantization, KV-cache quantization, floating-point formats such as F ### Support Matrix -| Component | 2-bit | 6-bit | Notes | -| --- | --- | --- | --- | -| `MatMulNBits` contrib schema | Yes | No | The schema attribute lists 2, 4, and 8 as supported values. | -| Python `MatMulNBits` quantizer | Yes | No | Dedicated 2-bit, 4-bit, and 8-bit native packers are exported. | -| CPU `MatMulNBits` | Yes | No | Includes fallback, AVX512/VNNI, and Arm64 work. | -| CUDA `MatMulNBits` | No | No | Kernel dispatch, prepacking, and memory estimation accept only 4 and 8 bits. | -| WebGPU `MatMulNBits` | Yes | No | Supports 2-bit unpacking with symmetric and asymmetric test coverage. | -| CPU `GatherBlockQuantized` | Yes | No | Supports packed 2-bit, 4-bit, and 8-bit integer data. | -| CUDA `GatherBlockQuantized` | Yes | No | The generic packed-`uint8` path supports 2-bit extraction. | -| QDQ quantization helper | Limited | No | The dedicated native QDQ packing helper is currently 4-bit-oriented. | +| Component | 2-bit | 3-bit | 5-bit | 6-bit | Notes | +| --- | --- | --- | --- | --- | --- | +| `MatMulNBits` contrib schema | Yes | No | No | No | The schema attribute lists 2, 4, and 8 as supported values. | +| Python `MatMulNBits` quantizer | Yes | No | No | No | Dedicated 2-bit, 4-bit, and 8-bit native packers are exported. | +| CPU `MatMulNBits` | Yes | No | No | No | Includes fallback, AVX512/VNNI, and Arm64 work for 2-bit. | +| CUDA `MatMulNBits` | No | No | No | No | Kernel dispatch, prepacking, and memory estimation accept only 4 and 8 bits. | +| WebGPU `MatMulNBits` | Yes | No | No | No | Supports 2-bit unpacking with symmetric and asymmetric test coverage. | +| CPU `GatherBlockQuantized` | Yes | No | No | No | Supports packed 2-bit, 4-bit, and 8-bit integer data. | +| CUDA `GatherBlockQuantized` | Yes | No | No | No | The generic packed-`uint8` path supports 2-bit extraction. | +| QDQ quantization helper | Limited | No | No | No | The dedicated native QDQ packing helper is currently 4-bit-oriented. | ### Existing 2-Bit Assets @@ -75,9 +75,9 @@ The existing INT4 and INT8 `MatMulNBits` CUDA implementations already select dif For eligible prepacked FP16/BF16 configurations, the `fpA_intB` profiler selects between its CUDA GEMV and CUTLASS GEMM tactics using the actual M bucket rather than a single fixed threshold. Consequently, the table describes the current specialized direct-kernel coverage, not a universal three-way dispatcher that applies to every data type and shape. -The proposed INT6 implementation should preserve this M-sensitive architecture: a fused GEMV for M=1 decode, a tiled fused GEMM for representative large-M prefill, and eventually a separate small-M path. It cannot simply reuse the existing INT4/INT8 CUTLASS kernels because NVIDIA hardware and the current CUTLASS integration do not expose a native INT6 weight-only operation. INT6 therefore also requires a new lower-4/upper-2 prepack, extraction logic, and fused compute kernels. Unsupported INT6 configurations should retain a correctness fallback, but that fallback is not a performance milestone. +Any selected INT3, INT5, or INT6 implementation should preserve this M-sensitive architecture: a fused GEMV for M=1 decode, a tiled fused GEMM for representative large-M prefill, and eventually a separate small-M path. These widths cannot simply reuse the existing INT4/INT8 CUTLASS kernels because NVIDIA hardware and the current CUTLASS integration do not expose corresponding native weight-only operations. They therefore require width-aware runtime prepacking, extraction logic, and fused compute kernels. Unsupported configurations should retain a correctness fallback, but that fallback is not a performance milestone. -### INT2 QMoE Follow-Up Assessment +### INT2 QMoE Product Assessment INT2 QMoE is a meaningful follow-up because expert weights dominate the storage and memory traffic of large MoE models, and a fused operator preserves top-k routing instead of expanding every expert into independently scheduled dense operations. CPU QMoE already accepts blockwise INT2 and includes an MLAS LUT GEMM path, so it can provide a semantic reference. However, the current cross-provider implementation is not production-ready: @@ -142,13 +142,13 @@ Backend coverage and maturity vary by bit width. The most broadly optimized CUDA ### TensorRT-LLM -TensorRT-LLM weight-only CUDA kernels and tests focus on INT4, INT8, and NVFP4. Its kernel type registry does not expose equivalent INT2 or INT6 weight-only paths. This is consistent with current NVIDIA hardware acceleration, which does not provide a native INT6 Tensor Core operation. +TensorRT-LLM weight-only CUDA kernels and tests focus on INT4, INT8, and NVFP4. Its kernel type registry does not expose equivalent INT2, INT3, INT5, or INT6 weight-only paths. This is consistent with current NVIDIA hardware acceleration, which does not provide native INT3, INT5, or INT6 Tensor Core operations. ### ONNX Data Types ONNX standardizes `INT2` and `UINT2`. Four consecutive 2-bit values are packed into one byte from least significant bits to most significant bits. This is compatible with the logical ordering already used by ONNX Runtime's packed 2-bit weights. -ONNX currently has floating-point 6-bit types, but no standard `INT6` or `UINT6` tensor type. `MatMulNBits` can still carry 6-bit codes in an opaque `uint8` blob, but the contrib operator must define the packing contract precisely. +ONNX does not currently provide standard `INT3`, `UINT3`, `INT5`, `UINT5`, `INT6`, or `UINT6` integer tensor types. `MatMulNBits` can still carry these codes in an opaque `uint8` blob, but the contrib operator must define the packing contract precisely. ## Design Considerations @@ -165,11 +165,21 @@ packed = (x0 & 0x03) No value crosses a byte boundary, and block sizes supported by `MatMulNBits` are multiples of four. The primary work is therefore execution-provider implementation and performance tuning rather than format design. -### 6-Bit Packing Options +### 3-, 5-, and 6-Bit Packing Options + +All three widths cross byte boundaries under a compact representation. A common packing quantum of eight logical values gives: + +| Width | Values per quantum | Payload bytes | Candidate runtime bit planes | +| ---: | ---: | ---: | --- | +| 3 | 8 | 3 | Lower 2 bits plus upper 1 bit | +| 5 | 8 | 5 | Lower 4 bits plus upper 1 bit | +| 6 | 8 | 6 | Lower 4 bits plus upper 2 bits | + +The portable format should not pad INT3 to four bits or INT5/INT6 to eight bits because that would erase much of the intended model-size benefit. #### Option A: Canonical Contiguous Bitstream -Pack four values into three bytes, with each value occupying six consecutive bits. This gives exactly six payload bits per weight and a portable representation. +Pack each value into a little-endian, least-significant-bit-first contiguous bitstream. Logical value `i` begins at bit offset `i * bits`; any unused tail bits in the final byte are zero. This gives exactly three, five, or six payload bits per weight before block metadata and padding. Advantages: @@ -183,14 +193,14 @@ Disadvantages: - Existing `8 / bits` pack and unpack logic cannot be reused. - Direct kernel extraction requires additional shifts and merged loads. -#### Option B: Lower-4/Upper-2 Split Planes +#### Option B: Backend-Specific Split Planes -Store low four bits and high two bits in separate regions, following the basic physical idea used by Q6_K. +Prepack the portable bitstream into aligned planes selected by the execution provider: lower-2/upper-1 for INT3, lower-4/upper-1 for INT5, and lower-4/upper-2 for INT6. The INT6 layout follows the basic physical idea used by Q6_K without adopting its hierarchical quantization format. Advantages: - Efficient aligned loads and extraction. -- Easier reuse of 4-bit and 2-bit unpacking primitives. +- Easier reuse of 1-bit, 2-bit, and 4-bit unpacking primitives. - Better fit for SIMD, DP4A, and backend-specific prepacking. Disadvantages: @@ -201,10 +211,12 @@ Disadvantages: #### Recommended Contract -Use a canonical contiguous bitstream in the portable model and allow each execution provider to prepack it into an opaque optimized representation, such as lower-4/upper-2 planes. Do not change the interpretation of existing models based only on `bits=6`. The exact canonical ordering, padding, zero-point packing, and versioning behavior must be specified first. +Use one canonical contiguous-bitstream rule for `bits=3`, `bits=5`, and `bits=6`, and allow each execution provider to prepack it into an opaque optimized representation. Do not change the interpretation of existing models based only on a new `bits` value. The exact ordering, row and block boundaries, tail padding, zero-point packing, and versioning behavior must be specified before enabling any new width. ## Proposed Roadmap +The dense `MatMulNBits` phases below describe supporting technology and independent projection coverage. The active QMoE product roadmap, milestones, and acceptance criteria are maintained in the separate end-to-end QMoE delivery plan. + ### Phase 0: Define Targets and Baselines Before implementation, agree on: @@ -276,9 +288,9 @@ Integrate the selected approach into `MatMulNBits(bits=2)` with: After the scoped delivery, expand to BF16, additional block sizes, asymmetric zero points, bias, tails, intermediate/small-M execution, offline prepacking, and broader GPU tuning as justified by measured demand. -### Phase 4: Mixed-Width INT2 QMoE Follow-Up +### Phase 4: Mixed-Width INT2 QMoE Product Track -Start the QMoE work after the MatMulNBits INT2 packing contract, CUDA correctness behavior, and fused M=1 primitives are stable: +Execute the prioritized QMoE work according to the standalone end-to-end delivery plan. Reuse MatMulNBits INT2 packing, validation, and CUDA extraction primitives as they stabilize, but do not gate the QMoE schema or Olive/Mobius export work on complete dense kernel coverage: 1. Define independent FC1 and FC2 bit-width semantics and preserve compatibility with the existing single-width QMoE contract. 2. Add CPU correctness and model-level tests for FC1 INT2 with FC2 INT4. @@ -288,9 +300,11 @@ Start the QMoE work after the MatMulNBits INT2 packing contract, CUDA correctnes WebGPU QMoE enablement is an independent follow-up and should not block the CUDA model-level milestone. A cross-provider correctness fallback may be useful for conformance, but persistent full expert dequantization does not satisfy the production acceptance criteria. -### Phase 5: 6-Bit Follow-Up Gate +### Phase 5: 3-, 5-, and 6-Bit MatMulNBits Gate + +Evaluate INT3, INT5, and INT6 if INT2 cannot achieve the required quality/size tradeoff or product requirements call for less aggressive quantization. Use one generic contiguous pack/unpack reference and compare all three widths against INT2, INT4, and INT8 on quality, effective model size, and extraction cost. INT3 is the first candidate for optimized implementation because it directly fills the gap between INT2 and INT4, but measured results must select the width. -Evaluate INT6 if INT2 cannot achieve the required quality/size tradeoff or product requirements call for a less aggressive quantization option. Reuse the contiguous portable format and lower-4/upper-2 prepacking analysis in this document, but require a separate format review and measured advantage over mixed INT4/INT8 before implementation. +Do not build three independent production kernel families in parallel. After the format and quality study, select at most one width for CUDA M=1 and large-M prototypes. A production commitment requires that the selected width materially outperform neighboring supported formats or mixed-precision recipes. ## Schedule Estimate @@ -308,22 +322,23 @@ These estimates assume one engineer working full time with Copilot assistance, t Allow **7-9 engineering weeks** for the scoped M=1 decode and large-M prefill vertical slice under the primary configuration. Broader data types, block sizes, asymmetric quantization, small-M kernels, and multi-architecture tuning would extend the work beyond this initial delivery. Copilot reduces coding and test-authoring time, but it does not remove hardware profiling, kernel tuning, model-quality evaluation, or code-review time. -### 6-Bit `MatMulNBits` +### 3-, 5-, and 6-Bit `MatMulNBits` Study | Work item | Estimate | | --- | ---: | -| Portable format specification and schema/tooling design | 1-2 weeks | -| Python pack/unpack, CPU reference, tests, and 4/6/8-bit quality study | 2-3 weeks | -| CUDA contiguous-versus-split-plane extraction and fused M=1/large-M prototypes | 3-4 weeks | +| Generic contiguous format specification and schema/tooling design | 1-2 weeks | +| Python pack/unpack, CPU reference, tests, and 2/3/4/5/6/8-bit quality study | 2-3 weeks | +| CUDA contiguous-versus-bit-plane extraction for all candidates | 1-2 weeks | +| Fused M=1 and large-M prototypes for one selected width | 3-4 weeks | | Go/no-go analysis and design review | 1 week | -Allow **6-8 engineering weeks** for the complete 6-bit format-and-value study and an evidence-based final go/no-go decision when both decode and prefill prototypes are required. The format, quality experiments, reference implementation, and scoped CUDA work can overlap to target an earlier vertical slice. A production-ready CUDA 6-bit track with broad operator coverage remains approximately **11-17 engineering weeks total**. Adding optimized CPU/MLAS support or another execution provider would require separate estimates. +Allow **7-10 engineering weeks** for the generic format-and-value study plus decode and prefill prototypes for one selected width. This estimate does not include production kernels for all three widths. A production-ready CUDA track for the selected width with broad operator coverage remains approximately **11-17 engineering weeks total** after selection. Each additional optimized width requires a separate estimate. Optimized CPU/MLAS support or another execution provider also requires separate estimates. -## November 15 Delivery Plan +## Supporting Dense MatMulNBits Delivery Estimate -There are approximately 8.5 calendar weeks from September 16 to November 15, 2026. With one engineer, the committed delivery should be a scoped INT2 CUDA vertical slice built on the existing portable format, Python tooling, and CPU reference implementation. The target is reviewable native fused CUDA execution for M=1 decode and representative large-M prefill under one primary symmetric FP16 configuration. This is narrower than complete production INT2 support. Upstream merge by November 15 cannot be guaranteed because review and CI timing are outside the implementation owner's control. +If staffed as a separate supporting track, approximately 8.5 calendar weeks would allow one engineer to target a scoped dense INT2 CUDA vertical slice built on the existing portable format, Python tooling, and CPU reference implementation. The candidate target is reviewable native fused CUDA execution for M=1 decode and representative large-M prefill under one primary symmetric FP16 configuration. This estimate is not the current QMoE product commitment, and upstream merge timing cannot be guaranteed because review and CI are outside the implementation owner's control. -### Committed INT2 Scope +### Candidate Dense INT2 Scope - Uniform and mixed INT2/INT4/INT8 quality and effective-size results for representative Qwen3.8 coding-model workloads. - A qualified PyTorch/Olive-to-Mobius export path producing dense mixed-bit `MatMulNBits` nodes, with graph, initializer-binding, and numerical-parity tests. @@ -347,7 +362,7 @@ There are approximately 8.5 calendar weeks from September 16 to November 15, 202 ### Stretch Scope -The following items should not put the November 15 commitment at risk: +The following items are excluded from the supporting dense estimate: - Asymmetric zero points. - BF16 activation support. @@ -356,15 +371,13 @@ The following items should not put the November 15 commitment at risk: - Optimized intermediate-M execution for M values between the committed decode and prefill ranges. - Offline CUDA-specific prepacking. - Broad multi-GPU tuning. -- INT6 format, tooling, or kernel implementation. +- INT3, INT5, or INT6 format, tooling, or kernel implementation. - ONNX-native Olive INT2 RTN and built-in quantized-linear export completion. -- Fused MoE/QMoE export, packing contracts, and runtime kernels. - -These QMoE items remain outside the November commitment, but the next implementation track should begin with the mixed FC1/FC2 contract and export qualification after the MatMulNBits INT2 foundation is stable. +- Fused MoE/QMoE export, packing contracts, and runtime kernels, which are tracked as the separate product-priority plan. -### INT6 Scheduling Impact +### INT3/INT5/INT6 Scheduling Impact -With the same engineer, INT6 implementation should not run concurrently if it threatens the INT2 deadline. Before November 15, INT6 work should be limited to preserving the format analysis and collecting quality data that directly informs the INT2 comparison. A second engineer could run the INT6 quality and format gate independently. +With the same engineer, INT3/INT5/INT6 implementation should not run concurrently if it threatens the prioritized INT2 QMoE delivery. Before November 15, work should be limited to preserving the generic format analysis and collecting quality data that informs width selection. A second engineer could run the generic reference and quality gate independently. ## Evaluation Plan @@ -382,7 +395,7 @@ At minimum, compare: - 2-bit RTN, GPTQ, HQQ, and K-quant-inspired quantization. - Uniform 2-bit versus mixed 2-bit/4-bit. -- Uniform 4-bit, 6-bit, and 8-bit. +- Uniform 2-bit, 3-bit, 4-bit, 5-bit, 6-bit, and 8-bit, plus selected mixed-precision recipes. - Perplexity plus representative downstream tasks used by the ONNX Runtime model-validation pipeline. ### Performance @@ -408,8 +421,10 @@ Report: - Two-bit quality may require mixed precision or more expensive calibration algorithms. - A 2-bit CUDA kernel can become unpacking-bound and fail to outperform a mature 4-bit kernel. -- Six-bit saves only 25 percent of weight payload relative to 8-bit while requiring non-native unpacking. -- A poorly specified 6-bit format could create incompatible models across execution providers. +- INT3, INT5, and INT6 require non-native unpacking; extraction overhead may erase their bandwidth advantage over neighboring formats. +- INT5 and INT6 may not save enough effective model size over higher-quality neighboring formats after scales, padding, and metadata. +- A poorly specified non-byte-aligned format could create incompatible models across execution providers. +- Supporting every candidate width would multiply kernel, prepacking, test, and maintenance cost without guaranteeing model-level value. - Extending CUTLASS prepacked formats too early would increase scope before value is established. - Metadata overhead can materially change the effective bits per weight, especially for small blocks. @@ -423,17 +438,18 @@ Report: - Qwen-shape decode, prefill, and TTFT correctness/performance report. - Go/no-go recommendation for broader production investment. -### 6-Bit Follow-Up Track +### 3-/5-/6-Bit Follow-Up Track -- Preserve the portable contiguous and lower-4/upper-2 prepacking design analysis. -- Run a 4-bit/6-bit/8-bit quality and effective-size comparison if INT2 quality is insufficient. -- Require a separate format and value gate before native CUDA implementation. +- Freeze one generic contiguous packing rule and width-specific 2+1, 4+1, and 4+2 runtime-prepack candidates. +- Run a 2/3/4/5/6/8-bit quality and effective-size comparison if INT2 quality is insufficient. +- Select at most one candidate width for initial CUDA decode and prefill prototypes. +- Require a separate value gate before production CUDA implementation. ## Recommendation to Management -Prioritize INT2 because ONNX Runtime already has a portable model format, quantization tooling, CPU kernels, and correctness coverage, and because INT2 offers the largest potential weight-memory and bandwidth reduction. Use the first week to freeze quality thresholds and identify a viable uniform or mixed-precision recipe, then target a scoped CUDA `MatMulNBits(bits=2)` vertical slice by November 15: direct-versus-LUT evidence, fused M=1 decode, and fused representative large-M prefill for the primary symmetric FP16 configuration. Both execution paths are required for a GitHub Copilot-style workload because decode determines generation speed and prefill determines TTFT for long repository context. +Prioritize the mixed-width INT2 QMoE delivery defined in the standalone plan because expert weights dominate the target model's storage and memory traffic. Develop shared dense MatMulNBits packing, validation, and CUDA extraction primitives where they directly reduce QMoE risk, but do not require complete dense GEMV/GEMM coverage before QMoE schema, export, CPU reference, and packed decode work proceeds. -This November scope is not complete production INT2 support. Broader data types, block sizes, asymmetric quantization, small-M execution, offline prepacking, and multi-architecture tuning remain follow-up work. Keep INT6 as the next quality-oriented option if INT2 cannot meet the agreed coding and tool-calling targets at a meaningful effective-size advantage over INT4. +Keep INT3, INT5, and INT6 MatMulNBits as a parallel measured quality/size study rather than a near-term production commitment. INT3 should receive the first optimization consideration because it fills the INT2-to-INT4 gap, but only after the common reference implementation and model-quality evidence select a width with a defensible advantage. Do not commit to production kernels for all candidate widths. ## References