Repository navigation
Conversation
- Register MatMulTernary with T8 v1 validation and shape inference - Implement NEON masked-add GEMM and a scalar reference path - Add opt-in tiled and decoded-mask prepacking with two-row reuse - Support direct Bonsai checkpoint-to-T8 export and Q2 graph conversion - Add codec, exporter, kernel, and prepacking regression tests - Add Q2/T8 and runtime-layout benchmarks with documented results Validation: - 14 native tests and 12 Python tests passed - Experimental layouts produced bit-identical outputs on 24 matrix cases and both complete Bonsai models - Masks with two-row reuse achieved 1.47x/1.49x speedups over current T8 on 2B/27B fixed-window inference - ARM64 wheel rebuilt and verified in an isolated smoke test Keep the existing kernel as default and preserve the on-disk T8 format.
MatMulTernary had a NEON kernel on ARM64. Every other target used a scalar fallback that rebuilt each T8 byte bit by bit and decoded it with divide and modulo by 3. On x86 that made a 1x2048x2048 GEMV take 22 ms. x86 kernels: - Add an MLAS_PLATFORM::T8GemmKernel pointer. platform.cpp selects the AVX2 kernel on AVX2+FMA hosts and the AVX-512F kernel on AVX-512F hosts. MlasT8Gemm and MlasT8GemmPacked keep the threading (16-unit tiles, row tiles of 1 or 2) and call the kernel when one is set. - Both kernels use the NEON kernel's bitsliced gate network, shared as a template in t8gemm_kernel.h. AVX-512F decodes four groups per vector with one vpternlogd per gate; AVX2 decodes two groups per vector. - AVX-512F uses each 16-bit half of a decoded word as an opmask: a masked xor applies the sign and a masked add accumulates. AVX2 broadcasts each word and moves every lane's bit into the sign position with vpsllvd. A is read in natural order, so the workspace is not used on x86. - The masks layout skips the decode. For a given kernel, the planes, tiled and masks layouts with row tile 1 or 2 give bit-identical results, as on NEON. - FP16 scales are converted with vcvtph2ps inside each kernel, not with the inline MLFloat16::ToFloat, so no ISA-specific copy of a shared inline function can reach other translation units. Scalar fallback: - It uses the same gate network on 32-bit words and a branchless masked add. Inputs are still summed one at a time in input order, so results match the old code (except that an exactly-zero sum is +0 instead of -0). - The packed path decodes a tiled group once per row tile, and the scalar masks packing uses the same decode. Single-thread time on a Ryzen AI MAX+ 395 (Zen 5): shape (MxNxK), layout old scalar new scalar AVX2 AVX-512F 1x2048x2048, planes 21.95 ms 2.49 ms 0.269 ms 0.135 ms 1x2048x2048, masks 16.58 ms 2.40 ms 0.244 ms 0.118 ms 1x2048x6144, planes 65.76 ms 7.47 ms 0.812 ms 0.412 ms 16x2048x2048, masks/2 244.0 ms - 2.61 ms 1.32 ms masks packing, 2048^2 21.6 ms 0.057 ms - - Tests: - New MLAS test T8Gemm.ShortExecute checks MlasT8Gemm and MlasT8GemmPacked (tiled and masks, row tiles 1 and 2) against a double-precision reference for M 1-5, N with padded units and K of 1-9 groups. - On x86 it also runs the entry points with the platform kernel cleared, to cover the scalar fallback. It calls the AVX2 and AVX-512F kernels directly, so AVX2 is tested on AVX-512 hosts, and checks that results are bit-identical across layouts and row tiles. - Passed: onnxruntime_mlas_test (31853), MatMulTernary provider tests (12), and the bonsai Python tests (12). - The Windows x64 cmake entries are not built or tested.
mirounga
marked this pull request as draft
November 2, 2026 18:37
mirounga
requested review from
Justin Chu (justinchuby) and
Xavier Dupré (xadupre)
November 2, 2026 18:38
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved converter bugs bypass zero-point validation for cached matrices and misreference retained external tensors, risking silently incorrect results or model-load failures.
7 open issues
What changed in this PR
Adds T8 ternary-weight support to ONNX Runtime through the com.microsoft.MatMulTernary CPU operator, with optimized kernels and Bonsai model conversion tools.
Changes:
- Introduces the operator schema, CPU registration, validation, and configurable weight prepacking.
- Adds NEON, AVX2, AVX-512F, and scalar kernels with runtime dispatch.
- Adds Bonsai export, conversion, generation, benchmarks, documentation, and regression tests.
| File | Description |
|---|---|
onnxruntime/test/mlas/unittest/test_t8gemm.cpp |
Tests kernel accuracy and layout equivalence. |
onnxruntime/test/contrib_ops/matmul_ternary_test.cc |
Tests operator execution, packing, and invalid inputs. |
onnxruntime/python/tools/transformers/models/bonsai/test_t8.py |
Tests codec, conversion, and benchmark checks. |
onnxruntime/python/tools/transformers/models/bonsai/test_bonsai_export.py |
Tests direct export and attention behavior. |
onnxruntime/python/tools/transformers/models/bonsai/t8_codec.py |
Encodes and decodes T8 weights. |
onnxruntime/python/tools/transformers/models/bonsai/run_onnx.py |
Runs fixed-window greedy generation. |
onnxruntime/python/tools/transformers/models/bonsai/README.md |
Documents conversion, layouts, and benchmark results. |
onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py |
Converts existing Q2 graphs to T8. |
onnxruntime/python/tools/transformers/models/bonsai/convert_to_onnx.py |
Exports Bonsai checkpoints as Q2 or T8. |
onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8.py |
Compares Q2 and T8 correctness and latency. |
onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8_layouts.py |
Compares runtime packing and row-reuse variants. |
onnxruntime/core/mlas/lib/t8gemm.cpp |
Implements validation, packing, NEON, and scalar execution. |
onnxruntime/core/mlas/lib/t8gemm_kernel.h |
Shares decoding logic and layout helpers. |
onnxruntime/core/mlas/lib/t8gemm_kernel_avx512f.cpp |
Implements the AVX-512F kernel. |
onnxruntime/core/mlas/lib/t8gemm_kernel_avx2.cpp |
Implements the AVX2 kernel. |
onnxruntime/core/mlas/lib/platform.cpp |
Selects x86 kernels at startup. |
onnxruntime/core/mlas/lib/mlasi.h |
Declares internal kernel interfaces and dispatch state. |
onnxruntime/core/mlas/inc/mlas_t8.h |
Declares T8 MLAS entry points. |
onnxruntime/core/graph/contrib_ops/contrib_defs.cc |
Defines the operator schema and shape inference. |
onnxruntime/contrib_ops/cpu/quantization/matmul_ternary.cc |
Implements CPU execution and prepacking. |
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc |
Registers the CPU operator. |
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h |
Adds packing and row-tile configuration keys. |
cmake/onnxruntime_mlas.cmake |
Includes T8 sources with architecture-specific compiler flags. |
Review effort: Balanced
Comment on lines
+61
to
+78
| weight, scales, zeros = (resolve(name) for name in node.input[1:4]) | ||
| key = (weight.name, scales.name, int(attrs["N"]), int(attrs["K"])) | ||
| if key not in matrices: | ||
| packed, s, zp = array(weight), array(scales), array(zeros) | ||
| if s.dtype != np.float32: | ||
| raise ValueError(f"{node.name}: MatMulTernary currently supports float32 activations only") | ||
| n, k = key[2:] | ||
| if packed.shape != (n, k // 128, 32) or k <= 0 or k % 128: | ||
| raise ValueError(f"{node.name}: Q2 shape disagrees with K/N") | ||
| if zp.dtype != np.uint8: | ||
| raise ValueError(f"{node.name}: expected packed uint8 zero points") | ||
| expected_zp_bytes = (k // 128 + 3) // 4 | ||
| if zp.shape not in ((n, expected_zp_bytes), (n * expected_zp_bytes,)): | ||
| raise ValueError(f"{node.name}: invalid zero-point shape") | ||
| zp_codes = ((zp.reshape(n, -1, 1) >> np.array([0, 2, 4, 6], dtype=np.uint8)) & 3).reshape(n, -1) | ||
| if not np.all(zp_codes[:, : k // 128] == 1): | ||
| raise ValueError(f"{node.name}: every zero point must equal 1") | ||
| planes, t8_scales = encode_q2(packed, s, verify=True) |
Comment on lines
+116
to
+118
| for tensor in tensors: | ||
| if not tensor.external_data or external_data_helper.ExternalDataInfo(tensor).location != data_path.name: | ||
| store(tensor) |
Comment on lines
+39
to
+49
| Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr allocator, | ||
| bool& is_packed, PrePackedWeights*) override { | ||
| is_packed = false; | ||
| if (input_idx != 1 || !use_packed_) return Status::OK(); | ||
| ORT_RETURN_IF_ERROR(ValidatePlanes(tensor)); | ||
| packed_ = IAllocator::MakeUniquePtr<uint32_t>(allocator, PackedWords(), true); | ||
| MlasT8Pack(n_, k_, tensor.Data<uint32_t>(), packed_.get(), layout_); | ||
| constant_planes_ = true; | ||
| is_packed = true; | ||
| return Status::OK(); | ||
| } |
| if (m == 0) return Status::OK(); | ||
| AllocatorPtr allocator; | ||
| ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); | ||
| auto workspace = IAllocator::MakeUniquePtr<float>(allocator, SafeInt<size_t>(m) * k_, true); |
Comment on lines
+79
to
+81
| names = (weight.name + ".t8", scales.name + ".t8") | ||
| if any(name in initializers for name in names): | ||
| raise ValueError(f"{node.name}: T8 initializer name collision") |
| f"{sequence_length - len(token_ids)} generated tokens" | ||
| ) | ||
|
|
||
| eos_token_ids = set(tokenizer.all_special_ids) |
Comment on lines
+3858
to
+3860
| ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulTernary) | ||
| .SetDomain(kMSDomain) | ||
| .SinceVersion(1) |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| from pathlib import Path | ||
|
|
||
| import numpy as np | ||
| import onnx |
| from pathlib import Path | ||
|
|
||
| import numpy as np | ||
| import onnx |
| from pathlib import Path | ||
|
|
||
| import numpy as np | ||
| import onnx |
| from unittest.mock import Mock | ||
|
|
||
| import numpy as np | ||
| import onnx |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ternary models (weights in {-1, 0, +1} times a group scale) are now shipped as real checkpoints, for example PrismML's Bonsai 2 (2B and 27B). Today ONNX Runtime can only run them as 2-bit
MatMulNBits. One of the four 2-bit code values is never used, so about 21% of the weight bits carry no information. Every weight also costs a full multiply-add, even though a ternary weight only ever means add, subtract or skip.This PR adds T8, a ternary weight format, and the
com.microsoft.MatMulTernaryCPU operator that runs it:MatMulNBits. Linear-layer storage drops about 24%.The T8 format
T8 v1 (
layout = "unit5x128-planes-v1") packs five ternary weights into one byte, because 3^5 = 243 ≤ 256:uint32bit planes of shape[ceil(N/5), 8, K/32]. One plane word covers 32 consecutive inputs.[ceil(N/5), K/128, 5]. Padded rows have scale zero.Storage: 8 bits / 5 weights = 1.6 bits, plus a 16-bit scale per 128 weights = 0.125 bits, for 1.725 bits per weight. Q2
MatMulNBitsneeds 2.266 bits per weight (2-bit codes, FP32 scales, packed zero points). The format is close to the information-theoretic minimum of log2(3) ≈ 1.585 bits per weight while staying byte-addressable.Why bit planes: a bitsliced gate network turns the eight plane words of a group straight into per-row "zero" and "positive" masks, for 32 inputs (or 128–512 inputs in a SIMD register) at once. Nothing is expanded to a dense weight matrix and nothing is divided by 3. The same layout is used by an existing AVX-512 reference and by the MLX implementation, so one file serves all of them.
What's in the PR
Operator and graph support:
com.microsoft.MatMulTernaryschema with shape inference, plus the CPU kernel. It takes float32 activations[..., K],uint32planes and FP16 scales, and validates the planes once for constant weights.PrePack. The on-disk format does not change:session.ternary.prepack = current | tiled | masks(defaultcurrent).session.ternary.row_tile = 1 | 2: two activation rows share one decode (prefill).MLAS kernels:
vpternlogdper gatevpsllvdMLAS_PLATFORM::T8GemmKernelpointer.Tooling (
onnxruntime/python/tools/transformers/models/bonsai/):convert_to_onnx.py --weight-format t8exports Bonsai 2 straight from the MLX checkpoint to T8, with no intermediate Q2 graph and no dequantize or requantize step.convert_to_t8.pyre-encodes an existing Q2 graph losslessly.benchmark_t8.pyandbenchmark_t8_layouts.pyreproduce the numbers below.Results
Whole model, Windows ARM64
These figures come from a Snapdragon X X1E80100 running ONNX Runtime built from this branch, with the CPU EP, 8 threads and float32 activations. Each step runs a fixed 24-token window with a 21-token prompt. Values are medians of 7 alternating Q2/T8 iterations, excluding session load.
The opt-in
maskslayout with two-row reuse adds another 1.47x (2B) and 1.49x (27B) over the default T8 kernel in the same setup. Its output is bit-identical on 24 actual-weight matrix shapes and both full models.Kernel, x86
These are single-thread figures on a Ryzen AI MAX+ 395 (Zen 5), using random valid T8 weights:
Testing
T8Gemmtest:onnxruntime_mlas_test(31,853 tests) on Linux x86_64. Native and Python tests on Windows ARM64 and Linux x86_64.Scope and limitations
MatMulTernary, so models exported with--weight-format t8need a build that includes this PR. Leaving out--weight-formatkeeps the existing Q2 export.