From 0044ecdb660ead8385ec8c8a6c9c7b6e8274cf98 Mon Sep 17 00:00:00 2001 From: Roman Snytsar Date: Mon, 2 Nov 2026 08:41:01 -0800 Subject: [PATCH 1/3] Add ARM64 NEON T8 ternary inference and Bonsai export support - 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. --- cmake/onnxruntime_mlas.cmake | 2 + .../onnxruntime_session_options_config_keys.h | 6 + .../contrib_ops/cpu/cpu_contrib_kernels.cc | 2 + .../cpu/quantization/matmul_ternary.cc | 117 ++++ .../core/graph/contrib_ops/contrib_defs.cc | 49 ++ onnxruntime/core/mlas/inc/mlas_t8.h | 28 + onnxruntime/core/mlas/lib/t8gemm.cpp | 382 +++++++++++++ .../transformers/models/bonsai/README.md | 264 +++++++++ .../models/bonsai/benchmark_t8.py | 251 +++++++++ .../models/bonsai/benchmark_t8_layouts.py | 208 +++++++ .../models/bonsai/convert_to_onnx.py | 508 ++++++++++++++++++ .../models/bonsai/convert_to_t8.py | 151 ++++++ .../transformers/models/bonsai/run_onnx.py | 106 ++++ .../transformers/models/bonsai/t8_codec.py | 65 +++ .../models/bonsai/test_bonsai_export.py | 155 ++++++ .../transformers/models/bonsai/test_t8.py | 174 ++++++ .../test/contrib_ops/matmul_ternary_test.cc | 202 +++++++ 17 files changed, 2670 insertions(+) create mode 100644 onnxruntime/contrib_ops/cpu/quantization/matmul_ternary.cc create mode 100644 onnxruntime/core/mlas/inc/mlas_t8.h create mode 100644 onnxruntime/core/mlas/lib/t8gemm.cpp create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/README.md create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8_layouts.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/convert_to_onnx.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/run_onnx.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/t8_codec.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/test_bonsai_export.py create mode 100644 onnxruntime/python/tools/transformers/models/bonsai/test_t8.py create mode 100644 onnxruntime/test/contrib_ops/matmul_ternary_test.cc diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 02168b4d579a6..746995d8cacea 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -61,6 +61,7 @@ onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/qdwconv_kernelsize.cpp ${MLAS_SRC_DIR}/qnbitgemm.h ${MLAS_SRC_DIR}/qnbitgemm.cpp + ${MLAS_SRC_DIR}/t8gemm.cpp ${MLAS_SRC_DIR}/qlutgemm.h ${MLAS_SRC_DIR}/qlutgemm.cpp ${MLAS_SRC_DIR}/sqnbitgemm_q8_block.h @@ -83,6 +84,7 @@ target_sources(onnxruntime_mlas PRIVATE ${MLAS_INC_DIR}/mlas_gemm_postprocessor.h ${MLAS_INC_DIR}/mlas_q4.h ${MLAS_INC_DIR}/mlas_qnbit.h + ${MLAS_INC_DIR}/mlas_t8.h ${MLAS_INC_DIR}/mlas_qkv_quant.h ${MLAS_INC_DIR}/mlas.h ) diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 02537bb643fe7..768cad3cf1f97 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -20,6 +20,12 @@ // If the config value is set to "1" then the prepacking is disabled, otherwise prepacking is enabled (default value) static const char* const kOrtSessionOptionsConfigDisablePrepacking = "session.disable_prepacking"; +// Experimental MatMulTernary CPU packing: "current" (default), "tiled", or "masks". +// Disk tensors remain T8 v1. Constant weights are packed at session initialization. +static const char* const kOrtSessionOptionsConfigTernaryPrepack = "session.ternary.prepack"; +// Activation rows sharing decoded masks: "1" (default) or "2"; "current" requires "1". +static const char* const kOrtSessionOptionsConfigTernaryRowTile = "session.ternary.row_tile"; + // Log MoE expert-counter updates at INFO severity. This enables counter updates even when expert counting is disabled. // At most 1024 updates are logged per Run; further updates emit one WARNING truncation marker but still count. // "0": disable (default); "1": enable. diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index 194c4dafbe518..5766877b691fa 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc @@ -75,6 +75,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, FusedMatMul); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MatMulNBits); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MLFloat16, MatMulNBits); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulTernary); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, int32_t, GatherBlockQuantized); class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, uint8_t, int64_t, GatherBlockQuantized); @@ -422,6 +423,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_ternary.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_ternary.cc new file mode 100644 index 0000000000000..5c6db7d0de681 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_ternary.cc @@ -0,0 +1,117 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/narrow.h" +#include "core/common/safeint.h" +#include "core/framework/op_kernel.h" +#include "core/mlas/inc/mlas_t8.h" +#include "core/providers/common.h" +#include "core/session/onnxruntime_session_options_config_keys.h" + +namespace onnxruntime::contrib { + +class MatMulTernary final : public OpKernel { + public: + explicit MatMulTernary(const OpKernelInfo& info) : OpKernel(info) { + const int64_t k = info.GetAttr("K"), n = info.GetAttr("N"); + ORT_ENFORCE(k > 0 && k % 128 == 0 && n > 0, "MatMulTernary requires K > 0 divisible by 128 and N > 0"); + ORT_ENFORCE(info.GetAttr("format_version") == 1 && + info.GetAttr("layout") == "unit5x128-planes-v1", + "Unsupported MatMulTernary T8 format version or layout"); + k_ = narrow(k); + n_ = narrow(n); + const auto packing = info.GetConfigOptions().GetConfigOrDefault(kOrtSessionOptionsConfigTernaryPrepack, "current"); + ORT_ENFORCE(packing == "current" || packing == "tiled" || packing == "masks", + "session.ternary.prepack must be current, tiled, or masks"); + use_packed_ = packing != "current"; + layout_ = packing == "masks" ? MLAS_T8_PACKED_LAYOUT::Masks : MLAS_T8_PACKED_LAYOUT::Tiled; + const auto row_tile = info.GetConfigOptions().GetConfigOrDefault(kOrtSessionOptionsConfigTernaryRowTile, "1"); + ORT_ENFORCE(row_tile == "1" || row_tile == "2", "session.ternary.row_tile must be 1 or 2"); + row_tile_ = row_tile == "2" ? 2 : 1; + ORT_ENFORCE(use_packed_ || row_tile_ == 1, "current T8 kernel requires row_tile 1"); + const Tensor* planes = nullptr; + if (info.TryGetConstantInput(1, &planes) && planes != nullptr) { + ORT_THROW_IF_ERROR(ValidatePlanes(*planes)); + constant_planes_ = true; + } + } + + 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(allocator, PackedWords(), true); + MlasT8Pack(n_, k_, tensor.Data(), packed_.get(), layout_); + constant_planes_ = true; + is_packed = true; + return Status::OK(); + } + + Status Compute(OpKernelContext* context) const override { + const Tensor* a = context->Input(0); + const Tensor* planes = packed_ ? nullptr : context->Input(1); + const Tensor* scales = context->Input(2); + const size_t units = n_ / 5 + (n_ % 5 != 0); + ORT_RETURN_IF_NOT(a->Shape().NumDimensions() >= 1 && a->Shape().GetDims().back() == static_cast(k_), + "MatMulTernary A must have rank >= 1 and last dimension K"); + ORT_RETURN_IF_NOT(scales->Shape() == TensorShape({narrow(units), narrow(k_ / 128), 5}), + "MatMulTernary scales must have shape [ceil(N/5), K/128, 5]"); + if (!constant_planes_) ORT_RETURN_IF_ERROR(ValidatePlanes(*planes)); + TensorShapeVector shape(a->Shape().GetDims().begin(), a->Shape().GetDims().end()); + shape.back() = narrow(n_); + Tensor* output = context->Output(0, TensorShape(shape)); + const size_t m = narrow(a->Shape().Size() / narrow(k_)); + if (m == 0) return Status::OK(); + AllocatorPtr allocator; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); + auto workspace = IAllocator::MakeUniquePtr(allocator, SafeInt(m) * k_, true); + if (use_packed_) { + IAllocatorUniquePtr dynamic_packed; + const uint32_t* packed = packed_.get(); + if (packed == nullptr) { + dynamic_packed = IAllocator::MakeUniquePtr(allocator, PackedWords(), true); + MlasT8Pack(n_, k_, planes->Data(), dynamic_packed.get(), layout_); + packed = dynamic_packed.get(); + } + MlasT8GemmPacked(m, n_, k_, a->Data(), packed, scales->Data(), + output->MutableData(), workspace.get(), context->GetOperatorThreadPool(), layout_, row_tile_); + } else { + MlasT8Gemm(m, n_, k_, a->Data(), planes->Data(), scales->Data(), + output->MutableData(), workspace.get(), context->GetOperatorThreadPool()); + } + return Status::OK(); + } + + private: + size_t PackedWords() const { + return SafeInt(n_ / 5 + (n_ % 5 != 0)) * (k_ / 128) * + (layout_ == MLAS_T8_PACKED_LAYOUT::Tiled ? 32 : 40); + } + + Status ValidatePlanes(const Tensor& planes) const { + const size_t units = n_ / 5 + (n_ % 5 != 0); + ORT_RETURN_IF_NOT(planes.Shape() == TensorShape({narrow(units), 8, narrow(k_ / 32)}), + "MatMulTernary B must have shape [ceil(N/5), 8, K/32]"); + ORT_RETURN_IF_NOT(MlasT8Validate(n_, k_, planes.Data()), + "MatMulTernary B contains an invalid T8 byte (243..255)"); + return Status::OK(); + } + + size_t k_{}, n_{}; + bool constant_planes_{false}; + bool use_packed_{false}; + size_t row_tile_{1}; + MLAS_T8_PACKED_LAYOUT layout_{MLAS_T8_PACKED_LAYOUT::Tiled}; + IAllocatorUniquePtr packed_{}; +}; + +ONNX_OPERATOR_KERNEL_EX( + MatMulTernary, kMSDomain, 1, kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("TPlanes", DataTypeImpl::GetTensorType()) + .TypeConstraint("TScale", DataTypeImpl::GetTensorType()), + MatMulTernary); + +} // namespace onnxruntime::contrib diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index eed58ac863515..3591c553da140 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -3855,6 +3855,55 @@ This op functions in much the same was as Dropout-11 and Dropout-13 do, except t } }); + ONNX_CONTRIB_OPERATOR_SCHEMA(MatMulTernary) + .SetDomain(kMSDomain) + .SinceVersion(1) + .SetDoc( + "Matmul with T8 v1 ternary weights: five trits per byte, stored as uint32 bitplanes. " + "Trit 0 is zero, 1 is +scale, 2 is -scale. Rows are padded to a multiple of five. " + "No bias or zero point is stored. Activations and output are float32; scales are float16.") + .Attr("K", "Positive input width, divisible by 128.", AttributeProto::INT) + .Attr("N", "Positive output width, excluding padded rows.", AttributeProto::INT) + .Attr("format_version", "T8 format version; only 1 is supported.", AttributeProto::INT, int64_t{1}) + .Attr("layout", "T8 tensor layout.", AttributeProto::STRING, std::string("unit5x128-planes-v1")) + .Input(0, "A", "Activations of shape [..., K], rank at least one.", "T") + .Input(1, "B", "Bitplanes [ceil(N/5), 8, K/32]; bit l holds lane l of a plane-set.", "TPlanes") + .Input(2, "scales", "Per-group scales [ceil(N/5), K/128, 5], padded rows have scale zero.", "TScale") + .Output(0, "Y", "Output of shape [..., N].", "T") + .TypeConstraint("T", {"tensor(float)"}, "Float32 activations and output.") + .TypeConstraint("TPlanes", {"tensor(uint32)"}, "Little-endian uint32 bitplanes.") + .TypeConstraint("TScale", {"tensor(float16)"}, "Float16 scales.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const int64_t k = getAttribute(ctx, "K", -1), n = getAttribute(ctx, "N", -1); + if (k <= 0 || k % 128 != 0 || n <= 0) { + fail_shape_inference("MatMulTernary requires K > 0 divisible by 128 and N > 0"); + } + const int64_t units = n / 5 + (n % 5 != 0); + const int64_t expected[2][3] = {{units, 8, k / 32}, {units, k / 128, 5}}; + for (size_t input = 1; input <= 2; ++input) { + if (!hasInputShape(ctx, input)) continue; + const auto& shape = getInputShape(ctx, input); + if (shape.dim_size() != 3) fail_shape_inference("MatMulTernary B and scales must have rank 3"); + for (int d = 0; d < 3; ++d) { + if (shape.dim(d).has_dim_value() && shape.dim(d).dim_value() != expected[input - 1][d]) { + fail_shape_inference("MatMulTernary invalid packed tensor dimension"); + } + } + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + const int rank = shape.dim_size(); + if (rank < 1) fail_shape_inference("MatMulTernary A must have rank >= 1"); + if (shape.dim(rank - 1).has_dim_value() && shape.dim(rank - 1).dim_value() != k) { + fail_shape_inference("MatMulTernary A last dimension must equal K"); + } + shape.mutable_dim(rank - 1)->clear_dim_param(); + shape.mutable_dim(rank - 1)->set_dim_value(n); + updateOutputShape(ctx, 0, shape); + } + }); + static const char* MatMulNBits_ver1_doc = R"DOC( MatMulNBits performs a matrix multiplication where the right-hand-side matrix (weights) is quantized to N bits. diff --git a/onnxruntime/core/mlas/inc/mlas_t8.h b/onnxruntime/core/mlas/inc/mlas_t8.h new file mode 100644 index 0000000000000..b99ca45090c97 --- /dev/null +++ b/onnxruntime/core/mlas/inc/mlas_t8.h @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "mlas.h" + +// T8 v1: uint32 planes [ceil(N/5), 8, K/32], fp16 scales [ceil(N/5), K/128, 5]. +// K is positive and divisible by 128. A and C are contiguous fp32 matrices. +// Workspace contains M*K floats; no expanded weight matrix is allocated. +void MLASCALL +MlasT8Gemm(size_t M, size_t N, size_t K, const float* A, const uint32_t* Planes, const MLAS_FP16* Scales, float* C, float* Workspace, MLAS_THREADPOOL* ThreadPool); + +bool MLASCALL +MlasT8Validate(size_t N, size_t K, const uint32_t* Planes); + +enum class MLAS_T8_PACKED_LAYOUT { + Tiled, + Masks, +}; + +// Each 5x128 group contains 32 tiled T8 words or 40 nonzero/sign mask words. +void MLASCALL +MlasT8Pack(size_t N, size_t K, const uint32_t* Planes, uint32_t* Packed, MLAS_T8_PACKED_LAYOUT Layout); + +// RowTile is 1 or 2. Workspace has the same size as MlasT8Gemm. +void MLASCALL +MlasT8GemmPacked(size_t M, size_t N, size_t K, const float* A, const uint32_t* Packed, const MLAS_FP16* Scales, float* C, float* Workspace, MLAS_THREADPOOL* ThreadPool, MLAS_T8_PACKED_LAYOUT Layout, size_t RowTile); diff --git a/onnxruntime/core/mlas/lib/t8gemm.cpp b/onnxruntime/core/mlas/lib/t8gemm.cpp new file mode 100644 index 0000000000000..cc4f49a54ea5d --- /dev/null +++ b/onnxruntime/core/mlas/lib/t8gemm.cpp @@ -0,0 +1,382 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/mlas/inc/mlas_t8.h" +#include "mlasi.h" + +#if defined(MLAS_TARGET_ARM64) +#include +#endif + +namespace +{ + +#if defined(MLAS_TARGET_ARM64) +using Bits = uint32x4_t; + +MLAS_FORCEINLINE Bits +Nt(Bits a, Bits b, Bits c) +{ + return vandq_u32(veorq_u32(a, b), veorq_u32(b, c)); +} +MLAS_FORCEINLINE Bits +Pt(Bits a, Bits b, Bits c, Bits nt) +{ + return vbicq_u32(veorq_u32(veorq_u32(a, b), c), nt); +} +MLAS_FORCEINLINE Bits +Bq(Bits a, Bits b, Bits c) +{ + return vandq_u32(veorq_u32(a, b), vorrq_u32(a, c)); +} +MLAS_FORCEINLINE Bits +Neg(Bits a, Bits b, Bits c) +{ + return veorq_u32(b, vandq_u32(veorq_u32(a, b), c)); +} +MLAS_FORCEINLINE Bits +Pos(Bits a, Bits b, Bits c) +{ + return vbicq_u32(veorq_u32(a, c), b); +} +MLAS_FORCEINLINE Bits +Zero(Bits a, Bits b, Bits c) +{ + return vmvnq_u32(vorrq_u32(a, veorq_u32(b, c))); +} +MLAS_FORCEINLINE Bits +Carry(Bits a, Bits b, Bits c) +{ + return vorrq_u32(a, vandq_u32(b, c)); +} + +// The decode3 gate network used by the AVX512 reference and the MLX T8 implementation. +MLAS_FORCEINLINE void +Decode(Bits b0, Bits b1, Bits b2, Bits b3, Bits b4, Bits b5, Bits b6, Bits b7, Bits* z, Bits* p) +{ + Bits nt = Nt(b7, b6, b5), pt = Pt(b7, b6, b5, nt); + b5 = Bq(b7, b6, b5); + b6 = vandq_u32(b6, b7); + Bits n = Neg(nt, pt, b4), q = Pos(nt, pt, b4); + b4 = Carry(nt, pt, b4); + nt = Neg(n, q, b3); + pt = Pos(n, q, b3); + b3 = Carry(n, q, b3); + n = Neg(nt, pt, b2); + q = Pos(nt, pt, b2); + b2 = Carry(nt, pt, b2); + nt = Neg(n, q, b1); + pt = Pos(n, q, b1); + b1 = Carry(n, q, b1); + z[0] = Zero(nt, pt, b0); + p[0] = Pos(nt, pt, b0); + b0 = Carry(nt, pt, b0); + + nt = Nt(b6, b5, b4); + pt = Pt(b6, b5, b4, nt); + b4 = Bq(b6, b5, b4); + b5 = vandq_u32(b5, b6); + n = Neg(nt, pt, b3); + q = Pos(nt, pt, b3); + b3 = Carry(nt, pt, b3); + nt = Neg(n, q, b2); + pt = Pos(n, q, b2); + b2 = Carry(n, q, b2); + n = Neg(nt, pt, b1); + q = Pos(nt, pt, b1); + b1 = Carry(nt, pt, b1); + z[1] = Zero(n, q, b0); + p[1] = Pos(n, q, b0); + b0 = Carry(n, q, b0); + + nt = Nt(b4, b3, b2); + pt = Pt(b4, b3, b2, nt); + b2 = Bq(b4, b3, b2); + b3 = vandq_u32(b3, b4); + n = Neg(nt, pt, b1); + q = Pos(nt, pt, b1); + b1 = Carry(nt, pt, b1); + z[2] = Zero(n, q, b0); + p[2] = Pos(n, q, b0); + b0 = Carry(n, q, b0); + + nt = Nt(b3, b2, b1); + pt = Pt(b3, b2, b1, nt); + Bits n4 = Bq(b3, b2, b1); + z[3] = Zero(nt, pt, b0); + p[3] = Pos(nt, pt, b0); + Bits p4 = Carry(nt, pt, b0); + z[4] = vmvnq_u32(vorrq_u32(n4, p4)); + p[4] = p4; +} + +void +PrepareA(size_t M, size_t K, const float* A, float* Workspace) +{ + for (size_t m = 0; m < M; ++m) { + for (size_t g = 0; g < K; g += 128) { + for (size_t i = 0; i < 32; i += 4) { + float32x4_t a = vld1q_f32(A + m * K + g + i); + float32x4_t b = vld1q_f32(A + m * K + g + 32 + i); + float32x4_t c = vld1q_f32(A + m * K + g + 64 + i); + float32x4_t d = vld1q_f32(A + m * K + g + 96 + i); + float32x4_t ab0 = vzip1q_f32(a, b), ab1 = vzip2q_f32(a, b); + float32x4_t cd0 = vzip1q_f32(c, d), cd1 = vzip2q_f32(c, d); + float* dst = Workspace + m * K + g + i * 4; + vst1q_f32(dst, vcombine_f32(vget_low_f32(ab0), vget_low_f32(cd0))); + vst1q_f32(dst + 4, vcombine_f32(vget_high_f32(ab0), vget_high_f32(cd0))); + vst1q_f32(dst + 8, vcombine_f32(vget_low_f32(ab1), vget_low_f32(cd1))); + vst1q_f32(dst + 12, vcombine_f32(vget_high_f32(ab1), vget_high_f32(cd1))); + } + } + } +} + +void +ComputeUnit(size_t N, size_t K, size_t unit, const float* A, const uint32_t* Planes, const MLAS_FP16* Scales, float* C) +{ + const size_t sets = K / 32; + const uint32_t* planes = Planes + unit * 8 * sets; + const MLAS_FP16* scales = Scales + unit * (K / 128) * 5; + float32x4_t total[5]{}; + const Bits sign_bit = vdupq_n_u32(0x80000000u); + for (size_t s = 0; s < sets; s += 4) { + Bits z[5], p[5], n[5]; + Decode(vld1q_u32(planes + s), vld1q_u32(planes + sets + s), vld1q_u32(planes + 2 * sets + s), vld1q_u32(planes + 3 * sets + s), vld1q_u32(planes + 4 * sets + s), vld1q_u32(planes + 5 * sets + s), vld1q_u32(planes + 6 * sets + s), vld1q_u32(planes + 7 * sets + s), z, p); + for (size_t j = 0; j < 5; ++j) n[j] = vmvnq_u32(vorrq_u32(z[j], p[j])); + float32x4_t sum[5]{}; + for (int i = 0; i < 32; ++i) { + Bits x = vreinterpretq_u32_f32(vld1q_f32(A + s * 32 + i * 4)); + int32x4_t shift = vdupq_n_s32(31 - i); + for (size_t j = 0; j < 5; ++j) { + Bits nonzero = vmvnq_u32(vreinterpretq_u32_s32( + vshrq_n_s32(vreinterpretq_s32_u32(vshlq_u32(z[j], shift)), 31) + )); + Bits sign = vandq_u32(vshlq_u32(n[j], shift), sign_bit); + sum[j] = vaddq_f32(sum[j], vreinterpretq_f32_u32(vandq_u32(veorq_u32(x, sign), nonzero))); + } + } + for (size_t j = 0; j < 5; ++j) { + total[j] = vfmaq_n_f32(total[j], sum[j], scales[(s / 4) * 5 + j].ToFloat()); + } + } + for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) { + C[unit * 5 + j] = vaddvq_f32(total[j]); + } +} +#else +void +ComputeUnit(size_t N, size_t K, size_t unit, const float* A, const uint32_t* Planes, const MLAS_FP16* Scales, float* C) +{ + const size_t sets = K / 32; + const uint32_t* planes = Planes + unit * 8 * sets; + float total[5]{}; + for (size_t g = 0; g < K / 128; ++g) { + float sum[5]{}; + for (size_t i = 0; i < 128; ++i) { + const size_t input = g * 128 + i; + uint32_t byte = 0; + for (size_t b = 0; b < 8; ++b) { + byte |= ((planes[b * sets + input / 32] >> (input % 32)) & 1u) << b; + } + for (size_t j = 0; j < 5; ++j, byte /= 3) { + const uint32_t trit = byte % 3; + if (trit == 1) sum[j] += A[input]; + if (trit == 2) sum[j] -= A[input]; + } + } + for (size_t j = 0; j < 5; ++j) { + total[j] += sum[j] * Scales[(unit * (K / 128) + g) * 5 + j].ToFloat(); + } + } + for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) C[unit * 5 + j] = total[j]; +} +#endif + +template +void +ComputePackedUnit(size_t M, size_t N, size_t K, size_t unit, const float* A, const uint32_t* Packed, const MLAS_FP16* Scales, float* C, MLAS_T8_PACKED_LAYOUT Layout) +{ + const size_t groups = K / 128; + const size_t words = Layout == MLAS_T8_PACKED_LAYOUT::Tiled ? 32 : 40; + const uint32_t* packed = Packed + unit * groups * words; + const MLAS_FP16* scales = Scales + unit * groups * 5; +#if defined(MLAS_TARGET_ARM64) + float32x4_t total[RowTile][5]{}; + const Bits sign_bit = vdupq_n_u32(0x80000000u); + for (size_t g = 0; g < groups; ++g) { + const uint32_t* block = packed + g * words; + Bits nz[5], neg[5]; + if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { + Bits z[5], p[5]; + Decode(vld1q_u32(block), vld1q_u32(block + 4), vld1q_u32(block + 8), vld1q_u32(block + 12), vld1q_u32(block + 16), vld1q_u32(block + 20), vld1q_u32(block + 24), vld1q_u32(block + 28), z, p); + for (size_t j = 0; j < 5; ++j) { + nz[j] = vmvnq_u32(z[j]); + neg[j] = vmvnq_u32(vorrq_u32(z[j], p[j])); + } + } else { + for (size_t j = 0; j < 5; ++j) { + nz[j] = vld1q_u32(block + j * 8); + neg[j] = vld1q_u32(block + j * 8 + 4); + } + } + float32x4_t sum[RowTile][5]{}; + for (int i = 0; i < 32; ++i) { + Bits x[RowTile]; + for (size_t r = 0; r < RowTile; ++r) { + if (r < M) x[r] = vreinterpretq_u32_f32(vld1q_f32(A + r * K + g * 128 + i * 4)); + } + const int32x4_t shift = vdupq_n_s32(31 - i); + for (size_t j = 0; j < 5; ++j) { + const Bits mask = vreinterpretq_u32_s32( + vshrq_n_s32(vreinterpretq_s32_u32(vshlq_u32(nz[j], shift)), 31) + ); + const Bits sign = vandq_u32(vshlq_u32(neg[j], shift), sign_bit); + for (size_t r = 0; r < RowTile; ++r) { + if (r < M) sum[r][j] = vaddq_f32(sum[r][j], vreinterpretq_f32_u32(vandq_u32(veorq_u32(x[r], sign), mask))); + } + } + } + for (size_t r = 0; r < RowTile; ++r) { + if (r < M) { + for (size_t j = 0; j < 5; ++j) { + total[r][j] = vfmaq_n_f32(total[r][j], sum[r][j], scales[g * 5 + j].ToFloat()); + } + } + } + } + for (size_t r = 0; r < RowTile; ++r) { + if (r < M) { + for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) { + C[r * N + unit * 5 + j] = vaddvq_f32(total[r][j]); + } + } + } +#else + for (size_t r = 0; r < std::min(RowTile, M); ++r) { + float total[5]{}; + for (size_t g = 0; g < groups; ++g) { + float sum[5]{}; + const uint32_t* block = packed + g * words; + for (size_t i = 0; i < 128; ++i) { + uint32_t byte = 0; + if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { + for (size_t b = 0; b < 8; ++b) byte |= ((block[b * 4 + i / 32] >> (i % 32)) & 1u) << b; + } + for (size_t j = 0; j < 5; ++j, byte /= 3) { + bool nonzero, negative; + if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { + nonzero = byte % 3 != 0; + negative = byte % 3 == 2; + } else { + nonzero = ((block[j * 8 + i / 32] >> (i % 32)) & 1u) != 0; + negative = ((block[j * 8 + 4 + i / 32] >> (i % 32)) & 1u) != 0; + } + if (nonzero) sum[j] += negative ? -A[r * K + g * 128 + i] : A[r * K + g * 128 + i]; + } + } + for (size_t j = 0; j < 5; ++j) total[j] += sum[j] * scales[g * 5 + j].ToFloat(); + } + for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) C[r * N + unit * 5 + j] = total[j]; + } +#endif +} + +} // namespace + +bool MLASCALL +MlasT8Validate(size_t N, size_t K, const uint32_t* Planes) +{ + const size_t units = N / 5 + (N % 5 != 0), sets = K / 32; + for (size_t q = 0; q < units; ++q) { + const uint32_t* p = Planes + q * 8 * sets; + for (size_t s = 0; s < sets; ++s) { + // 243..255 have the high nibble set and a low nibble >= 3. + if ((p[7 * sets + s] & p[6 * sets + s] & p[5 * sets + s] & p[4 * sets + s] & + (p[3 * sets + s] | p[2 * sets + s] | (p[sets + s] & p[s]))) != 0) return false; + } + } + return true; +} + +void MLASCALL +MlasT8Gemm(size_t M, size_t N, size_t K, const float* A, const uint32_t* Planes, const MLAS_FP16* Scales, float* C, float* Workspace, MLAS_THREADPOOL* ThreadPool) +{ + if (M == 0) return; +#if defined(MLAS_TARGET_ARM64) + PrepareA(M, K, A, Workspace); + A = Workspace; +#else + MLAS_UNREFERENCED_PARAMETER(Workspace); +#endif + const size_t units = N / 5 + (N % 5 != 0); + const size_t tiles = units / 16 + (units % 16 != 0); + MlasTrySimpleParallel(ThreadPool, static_cast(M * tiles), [&](ptrdiff_t task) { + const size_t m = static_cast(task) / tiles; + const size_t begin = (static_cast(task) % tiles) * 16; + for (size_t q = begin; q < std::min(begin + 16, units); ++q) { + ComputeUnit(N, K, q, A + m * K, Planes, Scales, C + m * N); + } + }); +} + +void MLASCALL +MlasT8Pack(size_t N, size_t K, const uint32_t* Planes, uint32_t* Packed, MLAS_T8_PACKED_LAYOUT Layout) +{ + const size_t groups = K / 128, sets = K / 32; + const size_t words = Layout == MLAS_T8_PACKED_LAYOUT::Tiled ? 32 : 40; + for (size_t q = 0; q < N / 5 + (N % 5 != 0); ++q) { + for (size_t g = 0; g < groups; ++g) { + const uint32_t* src = Planes + q * 8 * sets + g * 4; + uint32_t* dst = Packed + (q * groups + g) * words; + if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { + for (size_t b = 0; b < 8; ++b) std::copy_n(src + b * sets, 4, dst + b * 4); + } else { +#if defined(MLAS_TARGET_ARM64) + Bits z[5], p[5]; + Decode(vld1q_u32(src), vld1q_u32(src + sets), vld1q_u32(src + 2 * sets), vld1q_u32(src + 3 * sets), vld1q_u32(src + 4 * sets), vld1q_u32(src + 5 * sets), vld1q_u32(src + 6 * sets), vld1q_u32(src + 7 * sets), z, p); + for (size_t j = 0; j < 5; ++j) { + vst1q_u32(dst + j * 8, vmvnq_u32(z[j])); + vst1q_u32(dst + j * 8 + 4, vmvnq_u32(vorrq_u32(z[j], p[j]))); + } +#else + std::fill_n(dst, 40, 0u); + for (size_t i = 0; i < 128; ++i) { + uint32_t byte = 0; + for (size_t b = 0; b < 8; ++b) byte |= ((src[b * sets + i / 32] >> (i % 32)) & 1u) << b; + for (size_t j = 0; j < 5; ++j, byte /= 3) { + if (byte % 3 != 0) dst[j * 8 + i / 32] |= 1u << (i % 32); + if (byte % 3 == 2) dst[j * 8 + 4 + i / 32] |= 1u << (i % 32); + } + } +#endif + } + } + } +} + +void MLASCALL +MlasT8GemmPacked(size_t M, size_t N, size_t K, const float* A, const uint32_t* Packed, const MLAS_FP16* Scales, float* C, float* Workspace, MLAS_THREADPOOL* ThreadPool, MLAS_T8_PACKED_LAYOUT Layout, size_t RowTile) +{ + if (M == 0) return; +#if defined(MLAS_TARGET_ARM64) + PrepareA(M, K, A, Workspace); + A = Workspace; +#else + MLAS_UNREFERENCED_PARAMETER(Workspace); +#endif + const size_t units = N / 5 + (N % 5 != 0); + const size_t tiles = units / 16 + (units % 16 != 0); + const size_t row_tiles = M / RowTile + (M % RowTile != 0); + MlasTrySimpleParallel(ThreadPool, static_cast(row_tiles * tiles), [&](ptrdiff_t task) { + const size_t m = (static_cast(task) / tiles) * RowTile; + const size_t begin = (static_cast(task) % tiles) * 16; + for (size_t q = begin; q < std::min(begin + 16, units); ++q) { + if (RowTile == 2) { + ComputePackedUnit<2>(M - m, N, K, q, A + m * K, Packed, Scales, C + m * N, Layout); + } else { + ComputePackedUnit<1>(M - m, N, K, q, A + m * K, Packed, Scales, C + m * N, Layout); + } + } + }); +} diff --git a/onnxruntime/python/tools/transformers/models/bonsai/README.md b/onnxruntime/python/tools/transformers/models/bonsai/README.md new file mode 100644 index 0000000000000..921c43a24b65d --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/README.md @@ -0,0 +1,264 @@ +# Bonsai 2 ONNX conversion + +These tools convert the language component of a PrismML Bonsai 2 MLX pack to ONNX and run greedy text generation with +ONNX Runtime. Both the 2B (tied output head) and 27B (untied output head, swish attention gate) packs are supported. + +The converter preserves the published 2-bit, group-128 affine weights: + +- `MatMulNBits` runs packed linear layers with zero point 1. +- `GatherBlockQuantized` runs the packed token embedding. +- ONNX tensor operations apply the signed, normalized Walsh-Hadamard transforms (block 512 for 2B, 1024 for 27B). +- `LinearAttention` runs the Qwen3.5 gated-delta recurrence. + +## Dependencies + +Install `numpy`, `onnx`, `onnxruntime`, `safetensors`, `torch`, and `transformers`. + +Download a complete Bonsai 2 MLX pack. For example: + +```powershell +hf download prism-ml/Ternary-Bonsai-2-2B-mlx-2bit --local-dir C:\models\bonsai-2-2b +``` + +## Convert + +```powershell +python convert_to_onnx.py ` + --model-dir C:\models\bonsai-2-2b ` + --output C:\models\bonsai-2-2b.onnx ` + --sequence-length 24 +``` + +The exported model has a fixed sequence length and returns logits only for the position selected by `logits_index`. +Generation recomputes the fixed window instead of using a KV or recurrent-state cache. +Export a fresh Q2 graph with this converter before comparing formats. Older graphs can have a different +recurrent-attention implementation; matching T8 and Q2 outputs alone does not establish source-model correctness. + +## Run + +```powershell +python run_onnx.py ` + --model C:\models\bonsai-2-2b.onnx ` + --model-dir C:\models\bonsai-2-2b ` + --prompt "Capital of France? Answer with one word." ` + --max-new-tokens 1 +``` + +The prompt plus generated tokens must fit within the sequence length selected during conversion. + +## T8 ternary bitplanes on ARM64 + +Export directly from the original Bonsai MLX checkpoint with `--weight-format t8`: + +```powershell +python convert_to_onnx.py ` + --model-dir C:\models\bonsai-2-2b ` + --output C:\models\bonsai-2-2b-t8\model.onnx ` + --sequence-length 24 --weight-format t8 + +python convert_to_onnx.py ` + --model-dir C:\models\bonsai-2-27b ` + --output C:\models\bonsai-2-27b-t8\model.onnx ` + --sequence-length 24 --weight-format t8 +``` + +This path reads the original `model.safetensors`, encodes each linear matrix directly into T8 buffers, +and exports `MatMulTernary` nodes. It does not create or read an intermediate Q2 ONNX graph, dequantize +the full weights, or requantize the model. The checkpoint's affine 2-bit codes are losslessly repacked and +round-trip checked before export. Embeddings remain Q2; a tied output head is encoded independently as T8. +Omitting `--weight-format` preserves the existing Q2 export behavior. + +Use a separate output directory for each export and retain the graph together with **all** external-data +files written beside it. Small graphs can embed their weights; large exports use external data. +Direct exports record `t8.export=direct_checkpoint` and the format/source/round-trip metadata in the graph. +The direct 2B and 27B exports were checked against the earlier T8 graphs: all 552 matrices' bitplanes +and scales are bit-identical, and both direct exports pass the eager PyTorch full-logit comparison described below. + +For an existing Q2 ONNX graph, `convert_to_t8.py` remains available to losslessly re-encode its +linear weights as T8 v1, `unit5x128-planes-v1`. +This is the same tensor layout as the AVX512 reference and MLX implementation: + +- Five ternary rows share each base-3 byte: trit 0 is zero, 1 is positive, 2 is negative. +- Eight little-endian uint32 bitplanes have shape `[ceil(N/5), 8, K/32]`. +- Float16 scales have shape `[ceil(N/5), K/128, 5]`; padded rows have scale zero. +- Every matrix is expanded back to the original Q2 codes and checked bit-for-bit during conversion. +- Embedding gathers retain their original Q2 layout. A tied output head gets a separate T8 copy. + +T8 stores 1.725 bits per weight including scales, before row padding. The existing ONNX Q2 format uses +2.266 bits per weight with float32 scales and packed zero points (the source MLX format uses 2.25). +Total model savings depend on the embedding and other unchanged tensors; the 2B tied head requires an additional +copy in T8. + +```powershell +python convert_to_t8.py ` + --input C:\models\bonsai-2-2b.onnx ` + --output C:\models\bonsai-2-2b-t8.onnx +``` + +The Q2-graph re-encoder's output consists of the graph, an `.onnx.data` weight file and a `.validation.json` conversion report. +Existing output files are never overwritten. + +Run the result with `run_onnx.py` as above, using an ARM64 build of this branch. The new +`com.microsoft.MatMulTernary` CPU operator uses the T8 gate decoder and NEON masked add/subtract with float32 +accumulation; it does not expand a dense weight matrix. Other CPU architectures have a scalar reference path. +Unmodified ONNX Runtime releases do not recognize this operator. This is not an `onnxruntime-genai` model, +and does not add KV/recurrent caching or vision support. + +## Comparing T8 and Q2 + +Use the same newly built runtime for both formats: + +```powershell +python benchmark_t8.py --mode model ` + --model-dir C:\models\bonsai-2-2b ` + --q2 C:\models\bonsai-2-2b.onnx ` + --t8 C:\models\bonsai-2-2b-t8.onnx ` + --threads 8 --iterations 5 --output C:\models\bonsai-2-2b-benchmark.json + +python benchmark_t8.py --mode matmul ` + --model-dir C:\models\bonsai-2-27b ` + --rows 1 8 24 --threads 8 --iterations 5 ` + --output C:\models\bonsai-2-27b-matmul-benchmark.json +``` + +The model benchmark rejects non-finite outputs and checks logit agreement (`rtol=3e-4`, `atol=3e-3`), identical greedy next tokens and +KL divergence at most `1e-5`. Timings exclude loading and warmup, alternate Q2/T8 order and report medians +and all samples. Both formats use float32 activations, CPU EP and `ORT_ENABLE_ALL`; Q2 uses the existing +`accuracy_level=0` behavior, not an int8 activation approximation. +Thread-pool spinning is disabled so the inactive session does not compete with the session being measured. + +The reported model steps/second are **fixed-window recomputation**, not cached decode throughput or prefill +tokens/second. The matrix benchmark uses actual checkpoint weights for the MLP, recurrent projection and output head, +excluding Hadamard transforms equally from both formats. + +### Measured Windows ARM64 results + +Snapdragon X X1E80100, native ARM64 ONNX Runtime 1.32.0 built from this branch, CPU EP, +8 threads, float32 activations, sequence length 24, 21 prompt tokens. Each format received +3 warmups and 7 alternating timed iterations; numbers below are medians, excluding session loading. +Sizes include the graph and all referenced external data and use decimal GB. + +| Model | Q2 size | T8 size | Q2 step | T8 step | Q2/T8 speedup | +|---|---:|---:|---:|---:|---:| +| 2B | 0.669 GB | 0.562 GB | 2.807 s | 0.815 s | 3.44x | +| 27B | 7.682 GB | 6.002 GB | 93.207 s | 60.367 s | 1.54x | + +These are fixed-window recomputation steps, **not cached tokens/second**. The 27B timed samples ranged +from 88.554 to 98.779 seconds for Q2 and 58.974 to 62.179 seconds for T8; earlier runs varied substantially, +so compare the alternating samples within a run rather than treating absolute timings as hardware guarantees. +Maximum Q2/T8 logit differences were `1.12e-5` (2B) and `4.33e-4` (27B), with matching greedy tokens. +All 151 matrices in 2B and 401 matrices in 27B passed packed-code round-trip validation. +T8 matrix storage fell about 23.8%; whole-model savings were 16.0% and 21.9%, respectively. + +The actual-weight matrix benchmark also tested MLP up/down, recurrent QKV, and output-head matrices +at `M=1,8,24` for both model sizes (24 cases). All numerical checks passed; the largest absolute output +difference was `1.20e-6`. These matrix timings exclude Hadamard transforms and should not be substituted +for end-to-end generation throughput. + +### Validation scope + +The ARM64 tests cover every valid T8 byte and lane, padding, dynamic weights, batched activations, invalid +encodings, and both 16-head and 48-head recurrent projections. Nonzero recurrent outputs are also compared with +the independent PyTorch gated-delta reference. Conversion verifies the packed codes of every matrix. + +Full-model validation additionally compared both ONNX formats with eager PyTorch Qwen3.5 using the checkpoint's +original weights dequantized in float32. This checks the adapted language-only converter, not native MLX inference, +multimodal behavior, a model-quality evaluation, or cached generation. On the France prompt above, the fresh 2B +graphs and eager reference return `Paris`; the 27B graphs and eager reference all return `afs`. The latter is not a +correct answer, and format agreement must not be interpreted as validation of the 27B model's answer quality. + +## Experimental runtime weight layouts + +The default `MatMulTernary` kernel remains unchanged. Two opt-in CPU prepacking paths consume the same +T8 v1 graphs; no model conversion or on-disk format change is required: + +- `session.ternary.prepack=current` (default): original plane-major layout and kernel. +- `session.ternary.prepack=tiled`: reorder each 5-row/128-input group into eight adjacent NEON vectors, + `[row_unit, K/128, 8, 4]`. Packed weights remain 1.725 bits/weight including FP16 scales, before row padding. +- `session.ternary.prepack=masks`: decode once into nonzero/sign bitplanes, + `[row_unit, K/128, 5, 2, 4]`. Packed weights use 2.125 bits/weight including scales, about 23.2% more than T8. + +For either experimental layout, `session.ternary.row_tile=2` shares decoding and mask extraction between +two activation rows; the default is `1`. The current kernel only supports row tile `1`. +Constant planes are packed through ONNX Runtime's initialization-time `PrePack` hook. Dynamic planes, or +sessions with prepacking disabled, are packed per invocation instead, and that work is included in inference. +The packed sizes above describe weight buffers, not measured total session memory. + +```python +settings = onnxruntime.SessionOptions() +settings.add_session_config_entry("session.ternary.prepack", "masks") +settings.add_session_config_entry("session.ternary.row_tile", "2") +``` + +Compare both layouts, with and without two-row reuse, against the current kernel: + +```powershell +python benchmark_t8_layouts.py --mode matmul ` + --model-dir C:\models\bonsai-2-2b --rows 1 8 24 --threads 8 ` + --warmup 3 --iterations 10 --output C:\models\layout-matmul-2b.json + +python benchmark_t8_layouts.py --mode model ` + --model-dir C:\models\bonsai-2-2b --model C:\models\bonsai-2-2b-t8\model.onnx ` + --threads 8 --warmup 1 --iterations 5 --output C:\models\layout-model-2b.json +``` + +The layout benchmark requires bit-identical FP32 outputs, rotates measurement order, disables thread-pool +spinning, and records every sample plus initialization times. Packing is excluded from timed inference for +constant weights. Matrix tests use the checkpoint's actual MLP up/down, recurrent QKV, and output-head weights; +model timings remain fixed-window recomputation rather than cached decoding. + +### Measured layout results + +On Snapdragon X X1E80100, Windows ARM64, eight CPU threads and FP32 activations, all experimental outputs +were bit-identical to the current kernel: 24 actual-weight matrix shapes and both complete direct-T8 models. +The 14 native regression tests and 12 Python tests also passed, including dynamic weights, disabled prepacking, +empty activations, padding and odd two-row tails. + +Whole-model medians below use a 24-token window containing the same 21-token prompt. The 2B run used one warmup +and five timed samples per variant; the more expensive 27B run used one warmup and three timed samples and +tested only the two-row variants against current. Speedups refer to the current kernel in the same run, +not the earlier Q2/T8 measurements. + +| Runtime variant | 2B median (ms) | 2B speedup | 27B median (ms) | 27B speedup | +|---|---:|---:|---:|---:| +| Current | 3026.087 | 1.000x | 49050.671 | 1.000x | +| Tiled, one row | 3108.241 | 0.974x | Not measured | Not measured | +| Masks, one row | 2711.846 | 1.116x | Not measured | Not measured | +| Tiled, two-row reuse | 2200.528 | 1.375x | 34350.285 | 1.428x | +| Masks, two-row reuse | 2056.497 | 1.471x | 32872.387 | 1.492x | + +For 27B, the timed ranges were 45.966-49.954 s (current), 34.014-35.171 s (tiled/two-row) and +30.487-34.064 s (masks/two-row). Matrix samples also showed substantial timing variation; small differences +around 1.00x should not be interpreted as a reliable gain. + +Matrix results are geometric means of the four actual-weight matrices' current/variant median ratios, +with three warmups and ten timed samples per variant: + +| Checkpoint | Activation rows M | Tiled | Masks | Tiled/two-row | Masks/two-row | +|---|---:|---:|---:|---:|---:| +| 2B | 1 | 1.036x | 1.066x | 0.952x | 1.068x | +| 2B | 8 | 1.039x | 1.112x | 1.481x | 1.585x | +| 2B | 24 | 0.995x | 1.122x | 1.417x | 1.524x | +| 27B | 1 | 1.021x | 1.129x | 0.985x | 1.075x | +| 27B | 8 | 1.016x | 1.125x | 1.427x | 1.513x | +| 27B | 24 | 1.006x | 1.153x | 1.473x | 1.494x | + +Two-row reuse provides the main improvement for prefill; it offers no reuse when M=1. Tiling alone did not +give a consistent speedup. Prefer tiled/two-row when memory is important, or masks/two-row for the best +measured complete-model latency. Neither is enabled by default. + +Tiled packing retains the encoded plane size. Masks require 25% more plane-buffer memory: an additional +94,041,088 bytes for the direct 2B model (about 89.7 MiB), or 1,280,012,288 bytes for direct 27B (about 1.19 GiB). +This is not a total resident-memory measurement. On-disk graphs and weights remain unchanged. +Initialization times are recorded, but their fixed session-creation order and filesystem cache state mean +they are not an isolated packing-cost benchmark. + +The raw reports are `layout-matmul-2b.json`, `layout-matmul-27b.json`, `layout-model-2b.json` and +`layout-model-27b.json` in the local model-output directory. To reproduce the reduced 27B model run: + +```powershell +python benchmark_t8_layouts.py --mode model ` + --model-dir C:\models\bonsai-2-27b --model C:\models\bonsai-2-27b-t8\model.onnx ` + --variants current tiled_reuse2 masks_reuse2 --threads 8 ` + --warmup 1 --iterations 3 --output C:\models\layout-model-27b.json +``` diff --git a/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8.py b/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8.py new file mode 100644 index 0000000000000..749c068920c59 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8.py @@ -0,0 +1,251 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import json +import platform +import statistics +import time +from pathlib import Path + +import numpy as np +import onnx +from onnx import helper, numpy_helper +from safetensors import safe_open +from t8_codec import encode_q2 +from transformers import AutoTokenizer + +import onnxruntime as ort + + +def options(threads: int) -> ort.SessionOptions: + result = ort.SessionOptions() + result.intra_op_num_threads = threads + result.add_session_config_entry("session.intra_op.allow_spinning", "0") + result.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL + result.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + return result + + +def compare(q2: ort.InferenceSession, t8: ort.InferenceSession, feeds: dict, warmup: int, iterations: int) -> dict: + sessions = {"q2": q2, "t8": t8} + outputs = {name: session.run(None, feeds)[0] for name, session in sessions.items()} + for name, output in outputs.items(): + if not np.isfinite(output).all(): + raise ValueError(f"{name}: non-finite benchmark output") + np.testing.assert_allclose(outputs["t8"], outputs["q2"], rtol=3e-4, atol=3e-3) + maximum_error = float(np.max(np.abs(outputs["t8"] - outputs["q2"]))) + samples = {name: [] for name in sessions} + for _ in range(warmup): + for session in sessions.values(): + session.run(None, feeds) + for iteration in range(iterations): + order = ("q2", "t8") if iteration % 2 == 0 else ("t8", "q2") + for name in order: + start = time.perf_counter() + sessions[name].run(None, feeds) + samples[name].append(time.perf_counter() - start) + medians = {name: statistics.median(times) for name, times in samples.items()} + return { + "q2_ms": medians["q2"] * 1000, + "t8_ms": medians["t8"] * 1000, + "speedup_q2_over_t8": medians["q2"] / medians["t8"], + "max_abs_error": maximum_error, + "samples_seconds": samples, + "warmup": warmup, + "iterations": iterations, + } + + +def model_benchmark(args) -> list[dict]: + if args.q2 is None or args.t8 is None: + raise ValueError("Model benchmarks require --q2 and --t8") + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=False, fix_mistral_regex=True) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": args.prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False + ) + ids = tokenizer(prompt, add_special_tokens=False).input_ids + results = [] + for threads in args.threads: + start = time.perf_counter() + q2 = ort.InferenceSession(str(args.q2.resolve()), options(threads), providers=["CPUExecutionProvider"]) + q2_load = time.perf_counter() - start + start = time.perf_counter() + t8 = ort.InferenceSession(str(args.t8.resolve()), options(threads), providers=["CPUExecutionProvider"]) + t8_load = time.perf_counter() - start + shape = q2.get_inputs()[0].shape + if shape != t8.get_inputs()[0].shape or len(shape) != 2 or not isinstance(shape[1], int): + raise ValueError(f"Expected identical fixed-window input shapes; got {shape}") + length = shape[1] + if not ids or len(ids) > length: + raise ValueError(f"Prompt has {len(ids)} tokens; fixed window is {length}") + input_ids = np.zeros((1, length), dtype=np.int64) + input_ids[0, : len(ids)] = ids + attention_mask = np.zeros_like(input_ids) + attention_mask[0, : len(ids)] = 1 + feeds = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": np.arange(length, dtype=np.int64)[None, :], + "logits_index": np.array([len(ids) - 1], dtype=np.int64), + } + logits_q2, logits_t8 = (session.run(["logits"], feeds)[0] for session in (q2, t8)) + token_q2, token_t8 = (int(np.argmax(logits[0, 0])) for logits in (logits_q2, logits_t8)) + if token_q2 != token_t8: + raise ValueError(f"Greedy token mismatch: Q2={token_q2}, T8={token_t8}") + distributions = [] + for logits in (logits_q2, logits_t8): + values = logits[0, 0].astype(np.float64) + shifted = values - np.max(values) + distributions.append(shifted - np.log(np.exp(shifted).sum())) + kl = float(np.sum(np.exp(distributions[0]) * (distributions[0] - distributions[1]))) + if not np.isfinite(kl) or kl > 1e-5: + raise ValueError(f"Logit KL divergence {kl} must be finite and <= 1e-5") + result = compare(q2, t8, feeds, args.warmup, args.iterations) + result.update( + { + "kind": "fixed_window_model", + "threads": threads, + "sequence_length": length, + "prompt_tokens": len(ids), + "greedy_token": token_q2, + "greedy_text": tokenizer.decode([token_q2]), + "kl_q2_to_t8": kl, + "q2_load_seconds": q2_load, + "t8_load_seconds": t8_load, + "q2_recomputed_steps_per_second": 1000 / result["q2_ms"], + "t8_recomputed_steps_per_second": 1000 / result["t8_ms"], + } + ) + results.append(result) + print(json.dumps(result), flush=True) + del q2, t8 + return results + + +def checkpoint_paths(model_dir: Path) -> list[str]: + config = json.loads((model_dir / "config.json").read_text(encoding="utf-8")) + head = "model.embed_tokens" if config["text_config"]["tie_word_embeddings"] else "lm_head" + return [ + "model.layers.0.mlp.up_proj", + "model.layers.0.mlp.down_proj", + "model.layers.0.linear_attn.in_proj_qkv", + head, + ] + + +def load_checkpoint_matrix(weights, path: str) -> tuple[np.ndarray, np.ndarray]: + prefix = "language_model." + path + words = weights.get_tensor(prefix + ".weight") + n, k = words.shape[0], words.shape[1] * 16 + packed = words.view(np.uint8).reshape(n, k // 128, 32) + scales = weights.get_tensor(prefix + ".scales").astype(np.float32) + biases = weights.get_tensor(prefix + ".biases").astype(np.float32) + if not np.array_equal(biases, -scales): + raise ValueError(f"{path}: biases must equal -scales") + return packed, scales + + +def matmul_benchmark(args) -> list[dict]: + results = [] + with safe_open(args.model_dir / "model.safetensors", framework="np") as weights: + for path in checkpoint_paths(args.model_dir): + packed, scales = load_checkpoint_matrix(weights, path) + n, k = packed.shape[0], packed.shape[1] * 128 + planes, t8_scales = encode_q2(packed, scales) + zeros = np.full((n, (k // 128 + 3) // 4), 0x55, dtype=np.uint8) + q2_node = helper.make_node( + "MatMulNBits", + ["A", "B", "S", "Z"], + ["Y"], + domain="com.microsoft", + K=k, + N=n, + bits=2, + block_size=128, + accuracy_level=0, + ) + t8_node = helper.make_node("MatMulTernary", ["A", "B", "S"], ["Y"], domain="com.microsoft", K=k, N=n) + for m in args.rows: + input_info = helper.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [m, k]) + output_info = helper.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [m, n]) + models = [] + for node, tensors in ( + (q2_node, ((packed, "B"), (scales, "S"), (zeros, "Z"))), + (t8_node, ((planes, "B"), (t8_scales, "S"))), + ): + graph = helper.make_graph( + [node], + "matmul", + [input_info], + [output_info], + [numpy_helper.from_array(value, name) for value, name in tensors], + ) + models.append( + helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ir_version=10, + ).SerializeToString() + ) + feeds = {"A": np.random.default_rng(7).normal(0, 0.125, (m, k)).astype(np.float32)} + for threads in args.threads: + q2, t8 = ( + ort.InferenceSession(model, options(threads), providers=["CPUExecutionProvider"]) + for model in models + ) + result = compare(q2, t8, feeds, args.warmup, args.iterations) + result.update( + { + "kind": "matmul", + "path": path, + "M": m, + "N": n, + "K": k, + "threads": threads, + "q2_accuracy_level": 0, + "q2_weight_bytes": packed.nbytes + scales.nbytes + zeros.nbytes, + "t8_weight_bytes": planes.nbytes + t8_scales.nbytes, + } + ) + results.append(result) + print(json.dumps(result), flush=True) + del q2, t8 + del models + return results + + +def main() -> None: + parser = argparse.ArgumentParser(description="Benchmark native CPU T8 against the existing Bonsai ONNX Q2 format.") + parser.add_argument("--mode", choices=("model", "matmul"), default="model") + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--q2", type=Path) + parser.add_argument("--t8", type=Path) + parser.add_argument("--prompt", default="Capital of France? Answer with one word.") + parser.add_argument("--threads", type=int, nargs="+", default=[8]) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 8, 24]) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.warmup < 0 or args.iterations < 1 or any(value < 1 for value in args.threads + args.rows): + raise ValueError("Iterations, thread counts and row counts must be positive; warmup must be nonnegative") + report = { + "machine": platform.machine(), + "platform": platform.platform(), + "onnxruntime": ort.__version__, + "runtime_path": ort.__file__, + "model_dir": str(args.model_dir.resolve()), + "provider": "CPUExecutionProvider", + "graph_optimization": "ORT_ENABLE_ALL", + "activation_dtype": "float32", + "intra_op_spinning": False, + "results": (model_benchmark(args) if args.mode == "model" else matmul_benchmark(args)), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8_layouts.py b/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8_layouts.py new file mode 100644 index 0000000000000..de445f0139087 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/benchmark_t8_layouts.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import json +import platform +import statistics +import time +from pathlib import Path + +import numpy as np +import onnx +from benchmark_t8 import checkpoint_paths, load_checkpoint_matrix, options +from onnx import helper, numpy_helper +from safetensors import safe_open +from t8_codec import encode_q2 +from transformers import AutoTokenizer + +import onnxruntime as ort + +VARIANTS = { + "current": ("current", "1"), + "tiled": ("tiled", "1"), + "masks": ("masks", "1"), + "tiled_reuse2": ("tiled", "2"), + "masks_reuse2": ("masks", "2"), +} + + +def create_sessions(model, threads, variants): + sessions, loads = {}, {} + for name in variants: + packing, rows = VARIANTS[name] + settings = options(threads) + settings.add_session_config_entry("session.ternary.prepack", packing) + settings.add_session_config_entry("session.ternary.row_tile", rows) + start = time.perf_counter() + sessions[name] = ort.InferenceSession(model, settings, providers=["CPUExecutionProvider"]) + loads[name] = time.perf_counter() - start + return sessions, loads + + +def compare_variants(sessions, feeds, warmup, iterations): + expected = sessions["current"].run(None, feeds)[0] + if not np.isfinite(expected).all(): + raise ValueError("Non-finite current-kernel output") + errors = {} + for name, session in sessions.items(): + actual = session.run(None, feeds)[0] + if not np.isfinite(actual).all(): + raise ValueError(f"{name}: non-finite output") + np.testing.assert_array_equal(actual.view(np.uint32), expected.view(np.uint32), err_msg=name) + errors[name] = float(np.max(np.abs(actual - expected))) + names = list(sessions) + samples = {name: [] for name in names} + for _ in range(warmup): + for session in sessions.values(): + session.run(None, feeds) + for iteration in range(iterations): + offset = iteration % len(names) + order = names[offset:] + names[:offset] + if (iteration // len(names)) % 2: + order.reverse() + for name in order: + start = time.perf_counter() + sessions[name].run(None, feeds) + samples[name].append(time.perf_counter() - start) + medians = {name: statistics.median(values) for name, values in samples.items()} + return { + "median_ms": {name: value * 1000 for name, value in medians.items()}, + "speedup_over_current": {name: medians["current"] / value for name, value in medians.items()}, + "max_abs_error": errors, + "bit_identical": True, + "samples_seconds": samples, + "warmup": warmup, + "iterations": iterations, + } + + +def matmul_benchmark(args): + results = [] + with safe_open(args.model_dir / "model.safetensors", framework="np") as weights: + for path in checkpoint_paths(args.model_dir): + packed, scales = load_checkpoint_matrix(weights, path) + n, k = packed.shape[0], packed.shape[1] * 128 + planes, t8_scales = encode_q2(packed, scales) + units, groups = (n + 4) // 5, k // 128 + for m in args.rows: + graph = helper.make_graph( + [helper.make_node("MatMulTernary", ["A", "B", "S"], ["Y"], domain="com.microsoft", K=k, N=n)], + "layout_benchmark", + [helper.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [m, k])], + [helper.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [m, n])], + [numpy_helper.from_array(planes, "B"), numpy_helper.from_array(t8_scales, "S")], + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ir_version=10, + ).SerializeToString() + feeds = {"A": np.random.default_rng(7).normal(0, 0.125, (m, k)).astype(np.float32)} + for threads in args.threads: + sessions, loads = create_sessions(model, threads, args.variants) + result = compare_variants(sessions, feeds, args.warmup, args.iterations) + result.update( + { + "kind": "matmul", + "path": path, + "M": m, + "N": n, + "K": k, + "threads": threads, + "session_load_seconds": loads, + "encoded_plane_bytes": planes.nbytes, + "scale_bytes": t8_scales.nbytes, + "packed_plane_bytes": { + name: units * groups * (40 if VARIANTS[name][0] == "masks" else 32) * 4 + for name in args.variants + }, + } + ) + results.append(result) + print(json.dumps(result), flush=True) + del sessions + return results + + +def model_benchmark(args): + if args.model is None: + raise ValueError("--mode model requires --model") + tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=False, fix_mistral_regex=True) + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": args.prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False + ) + ids = tokenizer(prompt, add_special_tokens=False).input_ids + results = [] + for threads in args.threads: + sessions, loads = create_sessions(str(args.model.resolve()), threads, args.variants) + shape = sessions["current"].get_inputs()[0].shape + if len(shape) != 2 or not isinstance(shape[1], int) or not 0 < len(ids) <= shape[1]: + raise ValueError(f"Prompt of {len(ids)} tokens does not fit fixed input shape {shape}") + length = shape[1] + input_ids = np.zeros((1, length), dtype=np.int64) + input_ids[0, : len(ids)] = ids + mask = np.zeros_like(input_ids) + mask[0, : len(ids)] = 1 + feeds = { + "input_ids": input_ids, + "attention_mask": mask, + "position_ids": np.arange(length, dtype=np.int64)[None, :], + "logits_index": np.array([len(ids) - 1], dtype=np.int64), + } + result = compare_variants(sessions, feeds, args.warmup, args.iterations) + result.update( + { + "kind": "fixed_window_model", + "sequence_length": length, + "prompt_tokens": len(ids), + "threads": threads, + "session_load_seconds": loads, + } + ) + results.append(result) + print(json.dumps(result), flush=True) + del sessions + return results + + +def main(): + parser = argparse.ArgumentParser(description="Compare experimental T8 runtime layouts with the unchanged kernel.") + parser.add_argument("--mode", choices=("matmul", "model"), default="matmul") + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--model", type=Path) + parser.add_argument("--prompt", default="Capital of France? Answer with one word.") + parser.add_argument("--threads", type=int, nargs="+", default=[8]) + parser.add_argument("--rows", type=int, nargs="+", default=[1, 8, 24]) + parser.add_argument("--variants", choices=tuple(VARIANTS), nargs="+", default=list(VARIANTS)) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.warmup < 0 or args.iterations < 1 or any(value < 1 for value in args.threads + args.rows): + raise ValueError("Positive iterations/threads/rows and nonnegative warmup are required") + if "current" not in args.variants or len(set(args.variants)) != len(args.variants): + raise ValueError("--variants must include current and must not contain duplicates") + report = { + "platform": platform.platform(), + "machine": platform.machine(), + "onnxruntime": ort.__version__, + "runtime_path": ort.__file__, + "model_dir": str(args.model_dir.resolve()), + "model": str(args.model.resolve()) if args.model else None, + "variants": { + name: {"packing": VARIANTS[name][0], "row_tile": int(VARIANTS[name][1])} for name in args.variants + }, + "provider": "CPUExecutionProvider", + "activation_dtype": "float32", + "graph_optimization": "ORT_ENABLE_ALL", + "intra_op_spinning": False, + "results": model_benchmark(args) if args.mode == "model" else matmul_benchmark(args), + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2), encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_onnx.py new file mode 100644 index 0000000000000..3838bfda73041 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_onnx.py @@ -0,0 +1,508 @@ +#!/usr/bin/env python +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import numpy as np +import onnx +import torch +from safetensors import safe_open +from t8_codec import LAYOUT, encode_q2 +from torch import nn +from transformers.models.qwen3_5 import modeling_qwen3_5 +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig +from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5Attention, + Qwen3_5ForCausalLM, + Qwen3_5RMSNorm, + Qwen3_5TextRotaryEmbedding, +) + +QUANT_BITS = 2 +QUANT_BLOCK_SIZE = 128 +ZERO_POINT = 1 + + +class _MatMulNBits(torch.autograd.Function): + @staticmethod + def forward(ctx, x, packed_weight, scales, zero_points, in_features, out_features): + del ctx, packed_weight, scales, zero_points, in_features + return x.new_zeros((*x.shape[:-1], out_features)) + + @staticmethod + def symbolic(g, x, packed_weight, scales, zero_points, in_features, out_features): + return g.op( + "com.microsoft::MatMulNBits", + x, + packed_weight, + scales, + zero_points, + K_i=in_features, + N_i=out_features, + bits_i=QUANT_BITS, + block_size_i=QUANT_BLOCK_SIZE, + ) + + +class _GatherBlockQuantized(torch.autograd.Function): + @staticmethod + def forward(ctx, packed_weight, indices, scales, zero_points, hidden_size): + del ctx, packed_weight, scales, zero_points + return torch.zeros((*indices.shape, hidden_size), dtype=torch.float32, device=indices.device) + + @staticmethod + def symbolic(g, packed_weight, indices, scales, zero_points, hidden_size): + return g.op( + "com.microsoft::GatherBlockQuantized", + packed_weight, + indices, + scales, + zero_points, + gather_axis_i=0, + quantize_axis_i=1, + bits_i=QUANT_BITS, + block_size_i=QUANT_BLOCK_SIZE, + ) + + +class _MatMulTernary(torch.autograd.Function): + @staticmethod + def forward(ctx, x, planes, scales, in_features, out_features): + del ctx, planes, scales, in_features + return x.new_zeros((*x.shape[:-1], out_features)) + + @staticmethod + def symbolic(g, x, planes, scales, in_features, out_features): + return g.op( + "com.microsoft::MatMulTernary", + x, + planes, + scales, + K_i=in_features, + N_i=out_features, + format_version_i=1, + layout_s=LAYOUT, + ) + + +class _LinearAttention(torch.autograd.Function): + @staticmethod + def forward(ctx, query, key, value, unused_state, decay, beta, heads, head_dim): + del ctx, query, key, unused_state, decay, beta, heads, head_dim + return torch.zeros_like(value) + + @staticmethod + def symbolic(g, query, key, value, unused_state, decay, beta, heads, head_dim): + output_type = value.type() + flat_shape = g.op("Constant", value_t=torch.tensor([0, 0, -1], dtype=torch.int64)) + output_shape = g.op("Constant", value_t=torch.tensor([0, 0, heads, head_dim], dtype=torch.int64)) + query = g.op("Reshape", query, flat_shape) + key = g.op("Reshape", key, flat_shape) + value = g.op("Reshape", value, flat_shape) + output, _ = g.op( + "com.microsoft::LinearAttention", + query, + key, + value, + unused_state, + decay, + beta, + q_num_heads_i=heads, + kv_num_heads_i=heads, + update_rule_s="gated_delta", + outputs=2, + ) + return g.op("Reshape", output, output_shape).setType(output_type) + + +def _fwht(x: torch.Tensor, block_size: int) -> torch.Tensor: + original_shape = x.shape + x = x.reshape(-1, block_size) + stride = 1 + while stride < block_size: + x = x.reshape(-1, block_size // (2 * stride), 2, stride) + left = x[:, :, 0, :] + right = x[:, :, 1, :] + x = torch.stack((left + right, left - right), dim=2).reshape(-1, block_size) + stride *= 2 + return (x * (1.0 / math.sqrt(block_size))).reshape(original_shape) + + +def _gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, +) -> tuple[torch.Tensor, None]: + del kwargs + if initial_state is not None or output_final_state: + raise ValueError("The fixed-window Bonsai export does not use recurrent state inputs or outputs") + if use_qk_l2norm_in_kernel: + query = modeling_qwen3_5.l2norm(query, dim=-1, eps=1e-6) + key = modeling_qwen3_5.l2norm(key, dim=-1, eps=1e-6) + unused_state = query.new_empty((0,)) + return _LinearAttention.apply( + query, key, value, unused_state, g, beta, int(value.shape[2]), int(value.shape[3]) + ), None + + +class BonsaiSwishAttention(Qwen3_5Attention): + """The 27B checkpoint uses a swish output gate rather than Qwen3.5's sigmoid gate.""" + + def forward(self, hidden_states, position_embeddings, attention_mask, past_key_values=None, **kwargs): + del kwargs + if past_key_values is not None: + raise ValueError("The fixed-window Bonsai export does not use attention caches") + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + query, gate = torch.chunk(self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1) + query = self.q_norm(query.reshape(hidden_shape)).transpose(1, 2) + key = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + query, key = modeling_qwen3_5.apply_rotary_pos_emb(query, key, *position_embeddings) + output, weights = modeling_qwen3_5.eager_attention_forward( + self, query, key, value, attention_mask, scaling=self.scaling, dropout=0.0 + ) + output = output.reshape(*input_shape, -1) * torch.nn.functional.silu(gate.reshape(*input_shape, -1)) + return self.o_proj(output), weights + + +class PackedLinear(nn.Module): + def __init__( + self, + packed_weight: torch.Tensor, + scales: torch.Tensor, + zero_points: torch.Tensor, + signs: torch.Tensor, + transform_block_size: int, + ): + super().__init__() + self.register_buffer("packed_weight", packed_weight) + self.register_buffer("scales", scales) + self.register_buffer("zero_points", zero_points) + self.register_buffer("signs", signs) + self.transform_block_size = transform_block_size + self.out_features = packed_weight.shape[0] + self.in_features = packed_weight.shape[1] * QUANT_BLOCK_SIZE + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.transform_block_size: + x = _fwht(x * self.signs, self.transform_block_size) + return _MatMulNBits.apply( + x, + self.packed_weight, + self.scales, + self.zero_points, + self.in_features, + self.out_features, + ) + + +class T8Linear(nn.Module): + def __init__( + self, + packed_weight: torch.Tensor, + scales: torch.Tensor, + signs: torch.Tensor, + transform_block_size: int, + ): + super().__init__() + self.out_features = packed_weight.shape[0] + self.in_features = packed_weight.shape[1] * QUANT_BLOCK_SIZE + planes, t8_scales = encode_q2(packed_weight.numpy(), scales.numpy(), verify=True) + # The legacy Torch exporter cannot serialize uint32; restore the ONNX type after export. + self.register_buffer("planes", torch.from_numpy(planes.view(np.int32))) + self.register_buffer("scales", torch.from_numpy(t8_scales)) + self.register_buffer("signs", signs) + self.transform_block_size = transform_block_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.transform_block_size: + x = _fwht(x * self.signs, self.transform_block_size) + return _MatMulTernary.apply(x, self.planes, self.scales, self.in_features, self.out_features) + + +class PackedEmbedding(nn.Module): + def __init__( + self, + packed_weight: torch.Tensor, + scales: torch.Tensor, + zero_points: torch.Tensor, + signs: torch.Tensor, + transform_block_size: int, + ): + super().__init__() + packed_weight = packed_weight.reshape(packed_weight.shape[0], -1) + self.register_buffer("packed_weight", packed_weight) + self.register_buffer("scales", scales) + self.register_buffer("zero_points", zero_points) + self.register_buffer("signs", signs) + self.transform_block_size = transform_block_size + self.hidden_size = packed_weight.shape[1] * (8 // QUANT_BITS) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + x = _GatherBlockQuantized.apply( + self.packed_weight, + indices, + self.scales, + self.zero_points, + self.hidden_size, + ) + if self.transform_block_size: + x = _fwht(x, self.transform_block_size) * self.signs + return x + + +class BonsaiForNextToken(nn.Module): + def __init__(self, model: Qwen3_5ForCausalLM): + super().__init__() + self.model = model + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + logits_index: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.model.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, + return_dict=False, + )[0] + selected = torch.index_select(hidden_states, 1, logits_index) + return self.model.lm_head(selected) + + +def _resolve_module(root: nn.Module, path: str) -> tuple[nn.Module, str]: + parts = path.split(".") + parent = root + for part in parts[:-1]: + parent = parent[int(part)] if part.isdigit() else getattr(parent, part) + return parent, parts[-1] + + +def _packed_tensors(weights, path: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + words = weights.get_tensor(path + ".weight") + scales = weights.get_tensor(path + ".scales").astype(np.float32) + biases = weights.get_tensor(path + ".biases").astype(np.float32) + signs = weights.get_tensor(path + ".signs").astype(np.float32) + if not np.array_equal(biases, -scales): + raise ValueError(f"{path}: expected ternary affine bias to equal -scale") + + out_features, words_per_row = words.shape + in_features = words_per_row * 16 + if in_features % QUANT_BLOCK_SIZE: + raise ValueError(f"{path}: input width {in_features} is not divisible by {QUANT_BLOCK_SIZE}") + + packed = words.view(np.uint8).reshape( + out_features, + in_features // QUANT_BLOCK_SIZE, + QUANT_BLOCK_SIZE * QUANT_BITS // 8, + ) + if any(np.any(((packed >> shift) & 3) == 3) for shift in (0, 2, 4, 6)): + raise ValueError(f"{path}: non-ternary affine code 3") + num_blocks = in_features // QUANT_BLOCK_SIZE + zero_points = np.full((out_features, (num_blocks + 3) // 4), 0x55, dtype=np.uint8) + return ( + torch.from_numpy(np.ascontiguousarray(packed)), + torch.from_numpy(np.ascontiguousarray(scales)), + torch.from_numpy(zero_points), + torch.from_numpy(np.ascontiguousarray(signs)), + ) + + +def _assign_parameter(model: nn.Module, name: str, value: np.ndarray) -> None: + parent, attribute = _resolve_module(model, name) + current = getattr(parent, attribute) + if isinstance(parent, Qwen3_5RMSNorm): + value = value - 1.0 + tensor = torch.from_numpy(np.ascontiguousarray(value)) + if tuple(tensor.shape) != tuple(current.shape): + raise ValueError(f"{name}: source shape {tuple(tensor.shape)} != target shape {tuple(current.shape)}") + parent._parameters[attribute] = nn.Parameter(tensor, requires_grad=False) + + +def load_bonsai_model(model_dir: Path, weight_format: str = "q2") -> Qwen3_5ForCausalLM: + if weight_format not in ("q2", "t8"): + raise ValueError(f"Unsupported weight format: {weight_format}") + config_json = json.loads((model_dir / "config.json").read_text(encoding="utf-8")) + if config_json.get("model_type") != "prism_hadamard_qwen35": + raise ValueError("Expected a prism_hadamard_qwen35 model") + + config = Qwen3_5TextConfig.from_dict(config_json["text_config"]) + config.use_cache = False + config._attn_implementation = "eager" + with torch.device("meta"): + model = Qwen3_5ForCausalLM(config) + output_gate = config_json["text_config"].get("output_gate_type", "sigmoid") + if output_gate not in ("sigmoid", "swish"): + raise ValueError(f"Unsupported attention output gate: {output_gate}") + if output_gate == "swish": + for index, layer in enumerate(model.model.layers): + if hasattr(layer, "self_attn"): + layer.self_attn = BonsaiSwishAttention(config, index) + model.model.rotary_emb = Qwen3_5TextRotaryEmbedding(config) + + packed_paths = {record["path"]: record for record in config_json["modules"]} + embedding_path = "model.embed_tokens" + if embedding_path not in packed_paths: + raise ValueError("Packed embedding is missing from config.json") + + weights_path = model_dir / "model.safetensors" + with safe_open(weights_path, framework="np") as weights: + for path, record in packed_paths.items(): + source_path = "language_model." + path + packed, scales, zero_points, signs = _packed_tensors(weights, source_path) + if record["embedding"]: + module = PackedEmbedding(packed, scales, zero_points, signs, record["block"]) + elif weight_format == "t8": + module = T8Linear(packed, scales, signs, record["block"]) + else: + module = PackedLinear(packed, scales, zero_points, signs, record["block"]) + parent, attribute = _resolve_module(model, path) + setattr(parent, attribute, module) + + if config.tie_word_embeddings: + embedding_record = packed_paths[embedding_path] + packed, scales, zero_points, signs = _packed_tensors(weights, "language_model." + embedding_path) + model.lm_head = ( + T8Linear(packed, scales, signs, embedding_record["block"]) + if weight_format == "t8" + else PackedLinear(packed, scales, zero_points, signs, embedding_record["block"]) + ) + elif "lm_head" not in packed_paths: + raise ValueError("The untied output head is missing from config.json") + + packed_prefixes = tuple("language_model." + path + "." for path in packed_paths) + source_names = weights.keys() + for source_name in source_names: + if source_name.startswith(packed_prefixes): + continue + if not source_name.startswith("language_model."): + continue + target_name = source_name.removeprefix("language_model.") + value = weights.get_tensor(source_name) + if target_name.endswith(".conv1d.weight"): + value = value.transpose(0, 2, 1) + _assign_parameter(model, target_name, value) + + meta_parameters = [name for name, parameter in model.named_parameters() if parameter.is_meta] + if meta_parameters: + raise ValueError(f"Parameters were not loaded: {meta_parameters}") + model.eval() + return model + + +def _finalize_t8_weights(model: onnx.ModelProto) -> None: + initializers = {tensor.name: tensor for tensor in model.graph.initializer} + aliases = {node.output[0]: node.input[0] for node in model.graph.node if node.op_type == "Identity"} + plane_names = set() + for node in model.graph.node: + if node.domain != "com.microsoft" or node.op_type != "MatMulTernary": + continue + name = node.input[1] + seen = set() + while name in aliases: + if name in seen: + raise ValueError(f"T8 weight Identity cycle at {name}") + seen.add(name) + plane_names.add(name) + name = aliases[name] + plane_names.add(name) + if name not in initializers: + raise ValueError(f"Expected constant T8 planes for {node.name}") + tensor = initializers[name] + if tensor.data_type not in (onnx.TensorProto.INT32, onnx.TensorProto.UINT32): + raise ValueError(f"Expected 32-bit T8 planes for {node.name}") + if not tensor.raw_data and not tensor.external_data: + raise ValueError(f"Expected raw or external T8 plane bytes for {node.name}") + tensor.data_type = onnx.TensorProto.UINT32 + for value in model.graph.value_info: + if value.name in plane_names: + value.type.tensor_type.elem_type = onnx.TensorProto.UINT32 + + +def export_model(model_dir: Path, output_path: Path, sequence_length: int, weight_format: str = "q2") -> None: + modeling_qwen3_5.torch_chunk_gated_delta_rule = _gated_delta_rule + model = BonsaiForNextToken(load_bonsai_model(model_dir, weight_format)) + input_ids = torch.zeros((1, sequence_length), dtype=torch.int64) + attention_mask = torch.ones((1, sequence_length), dtype=torch.int64) + position_ids = torch.arange(sequence_length, dtype=torch.int64).unsqueeze(0) + logits_index = torch.tensor([sequence_length - 1], dtype=torch.int64) + + output_path.parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export( + model, + (input_ids, attention_mask, position_ids, logits_index), + str(output_path), + input_names=["input_ids", "attention_mask", "position_ids", "logits_index"], + output_names=["logits"], + opset_version=21, + custom_opsets={"com.microsoft": 1}, + dynamo=False, + external_data=True, + do_constant_folding=False, + ) + + exported = onnx.load(output_path, load_external_data=False) + for node in exported.graph.node: + if node.domain == "com.microsoft" and node.op_type == "LinearAttention": + node.input[3] = "" + if weight_format == "t8": + _finalize_t8_weights(exported) + metadata = {item.key: item.value for item in exported.metadata_props} + metadata.update( + { + "t8.format_version": "1", + "t8.layout": LAYOUT, + "t8.source": str(model_dir.resolve()), + "t8.export": "direct_checkpoint", + "t8.round_trip_verified": "true", + } + ) + onnx.helper.set_model_props(exported, metadata) + onnx.save(exported, output_path) + onnx.checker.check_model(output_path) + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Convert a PrismML Bonsai 2 MLX pack to ONNX.") + parser.add_argument("--model-dir", type=Path, required=True, help="Directory containing the Bonsai MLX pack.") + parser.add_argument("--output", type=Path, required=True, help="Output ONNX model path.") + parser.add_argument( + "--weight-format", choices=("q2", "t8"), default="q2", help="Linear weight format; embeddings remain Q2." + ) + parser.add_argument( + "--sequence-length", + type=int, + default=64, + help="Fixed sequence length for prompt and decode recomputation.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_arguments() + if args.sequence_length <= 0: + raise ValueError("--sequence-length must be positive") + export_model(args.model_dir.resolve(), args.output.resolve(), args.sequence_length, args.weight_format) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py new file mode 100644 index 0000000000000..49980637e9d30 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse +import json +from pathlib import Path + +import numpy as np +import onnx +from onnx import external_data_helper, helper, numpy_helper +from t8_codec import LAYOUT, encode_q2 + + +def convert_model(source: Path, output: Path) -> dict: + source, output = source.resolve(), output.resolve() + data_path = output.with_name(output.name + ".data") + if source == output or output.exists() or data_path.exists(): + raise ValueError("Choose a new output path; conversion never overwrites an existing model or data file") + model = onnx.load(source, load_external_data=False) + initializers = {tensor.name: tensor for tensor in model.graph.initializer} + identities = {node.output[0]: node.input[0] for node in model.graph.node if node.op_type == "Identity"} + + def resolve(name: str) -> onnx.TensorProto: + seen = set() + while name in identities: + if name in seen: + raise ValueError(f"Identity cycle at {name}") + seen.add(name) + name = identities[name] + if name not in initializers: + raise ValueError(f"Expected constant initializer for {name}") + return initializers[name] + + def array(tensor: onnx.TensorProto) -> np.ndarray: + return numpy_helper.to_array(tensor, base_dir=str(source.parent)) + + output.parent.mkdir(parents=True, exist_ok=True) + matrices = {} + original_bytes = t8_bytes = count = 0 + with data_path.open("xb") as data: + + def store(tensor: onnx.TensorProto) -> None: + raw = tensor.raw_data if tensor.HasField("raw_data") else array(tensor).tobytes() + offset = data.tell() + data.write(raw) + tensor.raw_data = b"" + external_data_helper.set_external_data(tensor, data_path.name, offset=offset, length=len(raw)) + tensor.ClearField("raw_data") + for field in ("float_data", "int32_data", "int64_data", "double_data", "uint64_data"): + tensor.ClearField(field) + + for node in model.graph.node: + if node.domain != "com.microsoft" or node.op_type != "MatMulNBits": + continue + attrs = {attr.name: helper.get_attribute_value(attr) for attr in node.attribute} + if attrs.get("bits") != 2 or attrs.get("block_size") != 128 or attrs.get("accuracy_level", 0) not in (0, 1): + raise ValueError(f"{node.name}: expected Q2/group128 with float32 computation") + if len(node.input) < 4 or not node.input[3] or any(node.input[4:]): + raise ValueError(f"{node.name}: expected explicit zero point and no g_idx or bias") + 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) + 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") + for name, value in zip(names, (planes, t8_scales), strict=True): + tensor = numpy_helper.from_array(value, name) + store(tensor) + initializers[name] = tensor + matrices[key] = names + original_bytes += packed.nbytes + s.nbytes + zp.nbytes + t8_bytes += planes.nbytes + t8_scales.nbytes + print(f"[{len(matrices)}] {node.name}: {n}x{k}", flush=True) + del packed, s, zp, planes, t8_scales + names = matrices[key] + node.op_type = "MatMulTernary" + del node.input[1:] + node.input.extend(names) + del node.attribute[:] + node.attribute.extend( + [ + helper.make_attribute("K", int(attrs["K"])), + helper.make_attribute("N", int(attrs["N"])), + helper.make_attribute("format_version", 1), + helper.make_attribute("layout", LAYOUT), + ] + ) + count += 1 + if count == 0: + raise ValueError("No compatible MatMulNBits nodes found") + live = {value.name for value in model.graph.output} + kept = [] + for node in reversed(model.graph.node): + if any(name in live for name in node.output): + kept.append(node) + live.update(name for name in node.input if name) + del model.graph.node[:] + model.graph.node.extend(reversed(kept)) + tensors = [tensor for name, tensor in initializers.items() if name in live] + for tensor in tensors: + if not tensor.external_data or external_data_helper.ExternalDataInfo(tensor).location != data_path.name: + store(tensor) + del model.graph.initializer[:] + model.graph.initializer.extend(tensors) + metadata = {item.key: item.value for item in model.metadata_props} + metadata.update({"t8.format_version": "1", "t8.layout": LAYOUT, "t8.source": str(source)}) + helper.set_model_props(model, metadata) + onnx.save_model(model, output) + onnx.checker.check_model(str(output)) + summary = { + "source": str(source), + "output": str(output), + "layout": LAYOUT, + "nodes_converted": count, + "unique_matrices": len(matrices), + "round_trip_verified": True, + "q2_matrix_bytes": original_bytes, + "t8_matrix_bytes": t8_bytes, + "matrix_byte_ratio": t8_bytes / original_bytes, + "output_bytes": output.stat().st_size + data_path.stat().st_size, + } + output.with_suffix(".validation.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description="Losslessly re-encode a Bonsai Q2 ONNX model as T8.") + parser.add_argument("--input", type=Path, required=True, help="Existing Q2 ONNX model.") + parser.add_argument("--output", type=Path, required=True, help="New T8 ONNX model.") + args = parser.parse_args() + print(json.dumps(convert_model(args.input, args.output), indent=2)) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/run_onnx.py b/onnxruntime/python/tools/transformers/models/bonsai/run_onnx.py new file mode 100644 index 0000000000000..16a00c034c6d1 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/run_onnx.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np +from transformers import AutoTokenizer + +import onnxruntime as ort + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run greedy text generation with a converted Bonsai 2 ONNX model.") + parser.add_argument("--model", type=Path, required=True, help="Converted ONNX model.") + parser.add_argument( + "--model-dir", type=Path, required=True, help="Bonsai source directory containing the tokenizer." + ) + parser.add_argument("--prompt", required=True, help="User prompt.") + parser.add_argument("--max-new-tokens", type=int, default=1, help="Maximum number of tokens to generate.") + return parser.parse_args() + + +def main() -> None: + args = parse_arguments() + if args.max_new_tokens <= 0: + raise ValueError("--max-new-tokens must be positive") + if hasattr(sys.stdout, "reconfigure"): + sys.stdout.reconfigure(errors="backslashreplace") + + tokenizer = AutoTokenizer.from_pretrained( + args.model_dir, + trust_remote_code=False, + fix_mistral_regex=True, + ) + rendered_prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": args.prompt}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + token_ids = tokenizer(rendered_prompt, add_special_tokens=False).input_ids + + session_options = ort.SessionOptions() + session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + start = time.perf_counter() + session = ort.InferenceSession( + str(args.model.resolve()), + session_options, + providers=["CPUExecutionProvider"], + ) + session_seconds = time.perf_counter() - start + + input_shape = session.get_inputs()[0].shape + if len(input_shape) != 2 or not isinstance(input_shape[1], int): + raise ValueError(f"Expected a fixed-rank input_ids shape, got {input_shape}") + sequence_length = input_shape[1] + if len(token_ids) + args.max_new_tokens > sequence_length: + raise ValueError( + f"Prompt has {len(token_ids)} tokens, leaving room for only " + f"{sequence_length - len(token_ids)} generated tokens" + ) + + eos_token_ids = set(tokenizer.all_special_ids) + generated: list[int] = [] + run_seconds = 0.0 + for _ in range(args.max_new_tokens): + active_ids = token_ids + generated + input_ids = np.zeros((1, sequence_length), dtype=np.int64) + input_ids[0, : len(active_ids)] = active_ids + attention_mask = np.zeros((1, sequence_length), dtype=np.int64) + attention_mask[0, : len(active_ids)] = 1 + position_ids = np.arange(sequence_length, dtype=np.int64)[None, :] + logits_index = np.array([len(active_ids) - 1], dtype=np.int64) + + start = time.perf_counter() + logits = session.run( + ["logits"], + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "logits_index": logits_index, + }, + )[0] + run_seconds += time.perf_counter() - start + next_token = int(np.argmax(logits[0, 0])) + if next_token in eos_token_ids: + break + generated.append(next_token) + + print(tokenizer.decode(generated), flush=True) + print(f"session_seconds={session_seconds:.3f}", flush=True) + print(f"run_seconds={run_seconds:.3f}", flush=True) + print(f"generated_token_ids={generated}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/t8_codec.py b/onnxruntime/python/tools/transformers/models/bonsai/t8_codec.py new file mode 100644 index 0000000000000..8a5728cfa351b --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/t8_codec.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import numpy as np + +LAYOUT = "unit5x128-planes-v1" +GROUP_SIZE = 128 +ROWS_PER_UNIT = 5 +POW3 = np.array([1, 3, 9, 27, 81], dtype=np.uint16) + + +def encode_q2(packed: np.ndarray, scales: np.ndarray, *, verify: bool = True) -> tuple[np.ndarray, np.ndarray]: + """Losslessly re-encode MatMulNBits W2, zero point 1, into the documented T8 layout.""" + if packed.dtype != np.uint8 or packed.ndim != 3 or packed.shape[2] != 32: + raise ValueError("Expected uint8 Q2 weights [N, K/128, 32]") + rows, groups, _ = packed.shape + if rows <= 0 or groups <= 0 or scales.shape not in ((rows, groups), (rows * groups,)): + raise ValueError("Invalid Q2 dimensions or scales shape") + scales = scales.reshape(rows, groups) + s16 = scales.astype(np.float16) + if not np.all(np.isfinite(scales)) or not np.array_equal(s16.astype(scales.dtype), scales): + raise ValueError("Scales must be finite and exactly representable as float16") + units, width = (rows + 4) // 5, groups * GROUP_SIZE + planes = np.empty((units, 8, width // 32), dtype="> np.array([0, 2, 4, 6], dtype=np.uint8)) & 3).reshape(-1, width) + if np.any(codes == 3): + raise ValueError(f"Non-ternary Q2 code 3 in rows {row_start}..{row_end - 1}") + trits = np.zeros(((end - start) * 5, width), dtype=np.uint8) + trits[: row_end - row_start] = (codes + 2) % 3 + values = (trits.reshape(end - start, 5, width) * POW3[None, :, None]).sum(axis=1).astype(np.uint8) + bits = (values.reshape(end - start, -1, 32, 1) >> np.arange(8, dtype=np.uint8)) & 1 + bits = bits.transpose(0, 3, 1, 2) + chunk = np.ascontiguousarray(np.packbits(bits, axis=-1, bitorder="little")).view(" np.ndarray: + """Expand a T8 matrix back to the original affine Q2 codes (not floating-point weights).""" + if planes.dtype != np.uint32 or planes.ndim != 3 or planes.shape[1] != 8: + raise ValueError("Expected uint32 T8 bitplanes [ceil(N/5), 8, K/32]") + if rows <= 0 or (rows + 4) // 5 != planes.shape[0] or planes.shape[2] % 4: + raise ValueError("Invalid T8 dimensions") + units, _, sets = planes.shape + raw = np.ascontiguousarray(planes, dtype=" 242): + raise ValueError("Invalid T8 byte (243..255)") + trits = ((values[:, None, :] // POW3[None, :, None]) % 3).reshape(units * 5, -1)[:rows] + codes = ((trits + 1) % 3).astype(np.uint8).reshape(rows, -1, 4) + packed = (codes << np.array([0, 2, 4, 6], dtype=np.uint8)).sum(axis=-1).astype(np.uint8) + return packed.reshape(rows, sets // 4, 32) diff --git a/onnxruntime/python/tools/transformers/models/bonsai/test_bonsai_export.py b/onnxruntime/python/tools/transformers/models/bonsai/test_bonsai_export.py new file mode 100644 index 0000000000000..a9335f16abbe2 --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/test_bonsai_export.py @@ -0,0 +1,155 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import onnx +import torch +from convert_to_onnx import BonsaiSwishAttention, T8Linear, _finalize_t8_weights, _fwht, _gated_delta_rule +from t8_codec import decode_q2 +from torch import nn +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig +from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Attention, torch_chunk_gated_delta_rule + +import onnxruntime as ort + + +class RecurrentProjection(nn.Module): + def __init__(self, normalize=False): + super().__init__() + self.normalize = normalize + + def forward(self, query, key, value, decay, beta): + return _gated_delta_rule(query, key, value, decay, beta, use_qk_l2norm_in_kernel=self.normalize)[0] + + +class BonsaiExportTest(unittest.TestCase): + def test_direct_t8_linear_export_and_execution(self): + rng = np.random.default_rng(7) + rows, width = 7, 256 + codes = rng.integers(0, 3, (rows, width), dtype=np.uint8) + packed = (codes.reshape(rows, -1, 4) << np.array([0, 2, 4, 6], dtype=np.uint8)).sum(axis=-1) + packed = packed.astype(np.uint8).reshape(rows, width // 128, 32) + scales = rng.integers(1, 32, (rows, width // 128)).astype(np.float32) / 128 + signs = torch.from_numpy(rng.choice(np.array([-1, 1], dtype=np.float32), width)) + activations = torch.from_numpy(rng.normal(size=(1, 3, width)).astype(np.float32)) + dense = (codes.astype(np.float32).reshape(rows, -1, 128) - 1) * scales[:, :, None] + for block in (0, 128): + with self.subTest(block=block), tempfile.TemporaryDirectory() as directory: + module = T8Linear(torch.from_numpy(packed), torch.from_numpy(scales), signs, block) + np.testing.assert_array_equal(decode_q2(module.planes.numpy().view(np.uint32), rows), packed) + output = Path(directory) / "direct-t8.onnx" + torch.onnx.export( + module, + (activations,), + str(output), + input_names=["A"], + output_names=["Y"], + opset_version=21, + custom_opsets={"com.microsoft": 1}, + dynamo=False, + ) + exported = onnx.load(output) + _finalize_t8_weights(exported) + onnx.save(exported, output) + onnx.checker.check_model(exported) + self.assertEqual(sum(node.op_type == "MatMulTernary" for node in exported.graph.node), 1) + self.assertFalse(any(node.op_type == "MatMulNBits" for node in exported.graph.node)) + options = ort.SessionOptions() + options.intra_op_num_threads = 2 + session = ort.InferenceSession(str(output), options, providers=["CPUExecutionProvider"]) + actual = session.run(None, {"A": activations.numpy()})[0] + transformed = _fwht(activations * signs, block) if block else activations + expected = transformed.numpy() @ dense.reshape(rows, width).T + np.testing.assert_allclose(actual, expected, rtol=3e-4, atol=3e-5) + + def test_recurrent_nonzero_outputs_match_torch_for_both_head_counts(self): + torch.manual_seed(7) + for heads in (16, 48): + with self.subTest(heads=heads), tempfile.TemporaryDirectory() as directory: + shape, gate_shape = (1, 4, heads, 128), (1, 4, heads) + query, key, value = (torch.randn(shape) * 0.1 for _ in range(3)) + decay, beta = -torch.rand(gate_shape), torch.sigmoid(torch.randn(gate_shape)) + inputs = query, key, value, decay, beta + expected = torch_chunk_gated_delta_rule(*inputs, use_qk_l2norm_in_kernel=True)[0] + output = Path(directory) / "recurrent.onnx" + torch.onnx.export( + RecurrentProjection(normalize=True), + inputs, + output, + input_names=["query", "key", "value", "decay", "beta"], + output_names=["output"], + opset_version=21, + custom_opsets={"com.microsoft": 1}, + dynamo=False, + ) + model = onnx.load(output) + next(node for node in model.graph.node if node.op_type == "LinearAttention").input[3] = "" + options = ort.SessionOptions() + options.intra_op_num_threads = 3 + options.add_session_config_entry("session.intra_op.allow_spinning", "0") + session = ort.InferenceSession(model.SerializeToString(), options, providers=["CPUExecutionProvider"]) + feeds = { + name: tensor.numpy() + for name, tensor in zip(["query", "key", "value", "decay", "beta"], inputs, strict=True) + } + actual = torch.from_numpy(session.run(None, feeds)[0]) + torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5) + + def test_recurrent_head_count_for_both_models(self): + for heads in (16, 48): + with self.subTest(heads=heads), tempfile.TemporaryDirectory() as directory: + output = Path(directory) / "recurrent.onnx" + qkv = torch.zeros((1, 2, heads, 128)) + gates = torch.zeros((1, 2, heads)) + torch.onnx.export( + RecurrentProjection(), + (qkv, qkv, qkv, gates, gates), + output, + opset_version=21, + custom_opsets={"com.microsoft": 1}, + dynamo=False, + ) + model = onnx.load(output) + graph = model.graph + node = next(node for node in graph.node if node.op_type == "LinearAttention") + attributes = { + attribute.name: onnx.helper.get_attribute_value(attribute) for attribute in node.attribute + } + self.assertEqual(attributes["q_num_heads"], heads) + self.assertEqual(attributes["kv_num_heads"], heads) + node.input[3] = "" + session = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + feeds = {value.name: np.zeros(value.shape, dtype=np.float32) for value in session.get_inputs()} + self.assertEqual(session.run(None, feeds)[0].shape, (1, 2, heads, 128)) + + def test_27b_swish_gate_preserves_attention_and_changes_only_gate(self): + config = Qwen3_5TextConfig( + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + num_hidden_layers=1, + intermediate_size=16, + ) + config._attn_implementation = "eager" + torch.manual_seed(7) + sigmoid = Qwen3_5Attention(config, 0).eval() + swish = BonsaiSwishAttention(config, 0).eval() + swish.load_state_dict(sigmoid.state_dict()) + sigmoid.o_proj = swish.o_proj = nn.Identity() + hidden = torch.randn((1, 3, 8)) + positions = (torch.ones((1, 3, 4)), torch.zeros((1, 3, 4))) + mask = torch.zeros((1, 1, 3, 3)) + base, base_weights = sigmoid(hidden, positions, mask) + actual, actual_weights = swish(hidden, positions, mask) + gate = sigmoid.q_proj(hidden).view(1, 3, 2, 8).chunk(2, dim=-1)[1].reshape(1, 3, 8) + torch.testing.assert_close(actual, base * gate, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(actual_weights, base_weights, rtol=0, atol=0) + + +if __name__ == "__main__": + unittest.main() diff --git a/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py b/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py new file mode 100644 index 0000000000000..c584084f65b9a --- /dev/null +++ b/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py @@ -0,0 +1,174 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import tempfile +import unittest +from pathlib import Path +from unittest.mock import Mock + +import numpy as np +import onnx +from benchmark_t8 import compare +from benchmark_t8_layouts import compare_variants +from convert_to_t8 import convert_model +from onnx import helper, numpy_helper +from t8_codec import decode_q2, encode_q2 + + +def pack_codes(codes): + rows, width = codes.shape + return ( + (codes.reshape(rows, -1, 4) << np.array([0, 2, 4, 6], dtype=np.uint8)) + .sum(axis=-1) + .astype(np.uint8) + .reshape(rows, width // 128, 32) + ) + + +class T8CodecTest(unittest.TestCase): + def test_all_bytes_and_lanes(self): + values = np.arange(256, dtype=np.uint16) % 243 + trits = ((values[None, :] // np.array([1, 3, 9, 27, 81], dtype=np.uint16)[:, None]) % 3).astype(np.uint8) + packed = pack_codes((trits + 1) % 3) + scales = np.ones((5, 2), dtype=np.float32) + planes, s = encode_q2(packed, scales) + self.assertEqual(planes.shape, (1, 8, 8)) + self.assertEqual(s.dtype, np.float16) + np.testing.assert_array_equal(decode_q2(planes, 5), packed) + for bit in range(8): + for group in range(8): + expected = sum(int((values[group * 32 + lane] >> bit) & 1) << lane for lane in range(32)) + self.assertEqual(int(planes[0, bit, group]), expected) + + def test_padded_rows_multiple_chunks_and_scales(self): + rng = np.random.default_rng(7) + rows, width = 647, 256 + packed = pack_codes(rng.integers(0, 3, (rows, width), dtype=np.uint8)) + scales = rng.integers(1, 32, (rows, width // 128)).astype(np.float32) / 128 + planes, s = encode_q2(packed, scales) + np.testing.assert_array_equal(decode_q2(planes, rows), packed) + restored_scales = s.transpose(0, 2, 1).reshape(-1, width // 128) + np.testing.assert_array_equal(restored_scales[:rows], scales) + np.testing.assert_array_equal(restored_scales[rows:], 0) + self.assertLess(planes.nbytes + s.nbytes, packed.nbytes + scales.nbytes) + + def test_rejects_nonternary_and_inexact_scales(self): + packed = np.full((5, 1, 32), 0x55, dtype=np.uint8) + scales = np.ones((5, 1), dtype=np.float32) + with self.assertRaisesRegex(ValueError, "code 3"): + encode_q2(packed | 0x03, scales) + with self.assertRaisesRegex(ValueError, "exactly representable"): + encode_q2(packed, scales * np.float32(0.1)) + with self.assertRaisesRegex(ValueError, "finite"): + encode_q2(packed, scales * np.nan) + + +class T8BenchmarkTest(unittest.TestCase): + def test_layout_benchmark_checks_all_variants(self): + sessions = {name: Mock() for name in ("current", "tiled", "masks")} + for session in sessions.values(): + session.run.return_value = [np.array([[1.0, -0.5]], dtype=np.float32)] + result = compare_variants(sessions, {}, warmup=0, iterations=3) + self.assertTrue(result["bit_identical"]) + self.assertEqual(set(result["median_ms"]), set(sessions)) + for samples in result["samples_seconds"].values(): + self.assertEqual(len(samples), 3) + sessions["masks"].run.return_value = [np.array([[1.0000001, -0.5]], dtype=np.float32)] + with self.assertRaises(AssertionError): + compare_variants(sessions, {}, warmup=0, iterations=1) + + def test_layout_benchmark_rejects_nonfinite_outputs(self): + sessions = {name: Mock() for name in ("current", "masks")} + sessions["current"].run.return_value = [np.array([[1.0]], dtype=np.float32)] + sessions["masks"].run.return_value = [np.array([[np.nan]], dtype=np.float32)] + with self.assertRaisesRegex(ValueError, "non-finite output"): + compare_variants(sessions, {}, warmup=0, iterations=1) + + def test_rejects_matching_nonfinite_outputs(self): + for value in (np.nan, np.inf, -np.inf): + with self.subTest(value=value): + session = Mock() + session.run.return_value = [np.array([[value]], dtype=np.float32)] + with self.assertRaisesRegex(ValueError, "non-finite benchmark output"): + compare(session, session, {}, warmup=0, iterations=1) + + def test_finite_outputs_and_timing_samples(self): + session = Mock() + session.run.return_value = [np.array([[0.0, 1.5, -0.75]], dtype=np.float32)] + result = compare(session, session, {}, warmup=0, iterations=1) + self.assertEqual(result["max_abs_error"], 0.0) + self.assertEqual(len(result["samples_seconds"]["q2"]), 1) + self.assertEqual(len(result["samples_seconds"]["t8"]), 1) + + +class T8ConversionTest(unittest.TestCase): + def test_external_data_and_shared_initializers(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + packed = np.full((7, 1, 32), 0x24, dtype=np.uint8) + scales = np.arange(1, 8, dtype=np.float32)[:, None] / 16 + tensors = [ + numpy_helper.from_array(packed, "B"), + numpy_helper.from_array(scales, "S"), + numpy_helper.from_array(np.full((7, 1), 0x55, dtype=np.uint8), "Z"), + ] + nodes = [ + helper.make_node("Identity", ["B"], ["B_alias"]), + helper.make_node( + "MatMulNBits", + ["A", "B_alias", "S", "Z"], + ["Y1"], + domain="com.microsoft", + K=128, + N=7, + bits=2, + block_size=128, + ), + helper.make_node( + "MatMulNBits", + ["A", "B", "S", "Z"], + ["Y2"], + domain="com.microsoft", + K=128, + N=7, + bits=2, + block_size=128, + ), + helper.make_node("Add", ["Y1", "Y2"], ["Y"]), + ] + graph = helper.make_graph( + nodes, + "test", + [helper.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [1, 128])], + [helper.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1, 7])], + tensors, + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ir_version=10, + ) + source, output = root / "q2.onnx", root / "t8.onnx" + onnx.save_model( + model, + source, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="q2.data", + size_threshold=0, + ) + summary = convert_model(source, output) + self.assertEqual(summary["unique_matrices"], 1) + self.assertEqual(summary["nodes_converted"], 2) + converted = onnx.load(output) + self.assertEqual([node.op_type for node in converted.graph.node], ["MatMulTernary", "MatMulTernary", "Add"]) + self.assertEqual(len(converted.graph.initializer), 2) + values = {tensor.name: numpy_helper.to_array(tensor) for tensor in converted.graph.initializer} + np.testing.assert_array_equal(decode_q2(values["B.t8"], 7), packed) + onnx.checker.check_model(str(output)) + with self.assertRaisesRegex(ValueError, "never overwrites"): + convert_model(source, output) + + +if __name__ == "__main__": + unittest.main() diff --git a/onnxruntime/test/contrib_ops/matmul_ternary_test.cc b/onnxruntime/test/contrib_ops/matmul_ternary_test.cc new file mode 100644 index 0000000000000..2c3517d2c35d5 --- /dev/null +++ b/onnxruntime/test/contrib_ops/matmul_ternary_test.cc @@ -0,0 +1,202 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" +#include "core/session/onnxruntime_session_options_config_keys.h" + +namespace onnxruntime::test { +namespace { + +void RunT8(int64_t m, int64_t n, int64_t k, bool initializer, bool identity = false, bool batched = false, + const char* packing = "current", const char* row_tile = "1", bool disable_prepacking = false) { + const int64_t units = (n + 4) / 5, sets = k / 32, groups = k / 128; + std::vector planes(static_cast(units * 8 * sets), 0); + std::vector scales(static_cast(units * groups * 5), MLFloat16(0.0f)); + std::vector weights(static_cast(n * k)); + for (int64_t q = 0; q < units; ++q) { + for (int64_t i = 0; i < k; ++i) { + uint32_t byte = 0, power = 1; + // Across the first 243 inputs, exercise every possible five-trit byte. + uint32_t digits = static_cast((i + 17 * q) % 243); + for (int64_t j = 0; j < 5; ++j, power *= 3, digits /= 3) { + const int64_t row = 5 * q + j; + const uint32_t trit = row < n ? digits % 3 : 0; + byte += power * trit; + if (row < n) { + const float scale = static_cast(1 + j + i / 128) * 0.125f; + scales[static_cast((q * groups + i / 128) * 5 + j)] = MLFloat16(scale); + weights[static_cast(row * k + i)] = + (trit == 1 ? 1.0f : trit == 2 ? -1.0f + : 0.0f) * + scale; + } + } + for (int64_t b = 0; b < 8; ++b) { + planes[static_cast((q * 8 + b) * sets + i / 32)] |= ((byte >> b) & 1u) << (i % 32); + } + } + } + std::vector a(static_cast(m * k), 0.0f); + for (int64_t row = 0; row < m; ++row) { + for (int64_t i = 0; i < k; ++i) { + a[static_cast(row * k + i)] = identity ? (row == i ? 1.0f : 0.0f) + : static_cast((i + row * 3) % 17 - 8) * 0.125f; + } + } + std::vector expected(static_cast(m * n), 0.0f); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < n; ++col) { + for (int64_t i = 0; i < k; ++i) { + expected[static_cast(row * n + col)] += + a[static_cast(row * k + i)] * weights[static_cast(col * k + i)]; + } + } + } + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", k); + test.AddAttribute("N", n); + test.AddInput("A", batched ? std::vector{1, m, k} : std::vector{m, k}, a); + test.AddInput("B", {units, 8, sets}, planes, initializer); + test.AddInput("scales", {units, groups, 5}, scales, initializer); + test.AddOutput("Y", batched ? std::vector{1, m, n} : std::vector{m, n}, expected); + std::vector> eps; + eps.push_back(DefaultCpuExecutionProvider()); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsConfigTernaryPrepack, packing)); + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsConfigTernaryRowTile, row_tile)); + if (disable_prepacking) ASSERT_STATUS_OK(options.config_options.AddConfigEntry(kOrtSessionOptionsConfigDisablePrepacking, "1")); + size_t prepacked_count = 0; + test.Run(options, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &eps, {}, &prepacked_count); + EXPECT_EQ(prepacked_count, initializer && std::string(packing) != "current" && !disable_prepacking ? 1u : 0u); +} + +TEST(MatMulTernaryTest, EveryByteAndLaneExactly) { RunT8(256, 5, 256, true, true); } +TEST(MatMulTernaryTest, DecodePaddedRowsAndMultipleGroups) { RunT8(1, 7, 512, true); } +TEST(MatMulTernaryTest, DynamicWeightsAndBatchedPrefill) { RunT8(9, 17, 256, false, false, true); } +TEST(MatMulTernaryTest, EmptyActivations) { RunT8(0, 7, 128, true); } + +TEST(MatMulTernaryTest, PackedLayoutsAndMaskReuse) { + for (const char* packing : {"tiled", "masks"}) { + for (const char* rows : {"1", "2"}) { + SCOPED_TRACE(std::string(packing) + "/" + rows); + RunT8(256, 5, 256, true, true, false, packing, rows); + RunT8(1, 7, 512, true, false, false, packing, rows); + RunT8(9, 17, 256, true, false, true, packing, rows); + RunT8(9, 17, 256, false, false, true, packing, rows); + RunT8(3, 7, 128, true, false, false, packing, rows, true); + RunT8(0, 7, 128, true, false, false, packing, rows); + } + } +} + +TEST(MatMulTernaryTest, RejectsInvalidPackingOptions) { + for (const auto& entry : std::vector>{ + {kOrtSessionOptionsConfigTernaryPrepack, "invalid"}, + {kOrtSessionOptionsConfigTernaryRowTile, "3"}, + {kOrtSessionOptionsConfigTernaryRowTile, "2"}}) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + test.AddInput("A", {1, 128}, std::vector(128)); + test.AddInput("B", {1, 8, 4}, std::vector(32), true); + test.AddInput("scales", {1, 1, 5}, std::vector(5)); + test.AddOutput("Y", {1, 5}, std::vector(5)); + SessionOptions options; + ASSERT_STATUS_OK(options.config_options.AddConfigEntry(entry.first.c_str(), entry.second.c_str())); + test.Run(options, OpTester::ExpectResult::kExpectFailure, + entry.second == "2" ? "current T8 kernel requires row_tile 1" + : entry.first + " must be"); + } +} + +TEST(MatMulTernaryTest, VectorInput) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 1); + test.AddInput("A", {128}, std::vector(128, 1.0f)); + test.AddInput("B", {1, 8, 4}, std::vector(32, 0)); + test.AddInput("scales", {1, 1, 5}, std::vector(5, MLFloat16(1.0f))); + test.AddOutput("Y", {1}, {0.0f}); + test.Run(); +} + +TEST(MatMulTernaryTest, RejectsAllInvalidBytes) { + for (uint32_t byte = 243; byte < 256; ++byte) { + for (bool initializer : {false, true}) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + std::vector planes(32, 0); + for (size_t bit = 0; bit < 8; ++bit) planes[bit * 4] = ((byte >> bit) & 1u) << 31; + test.AddInput("A", {1, 128}, std::vector(128)); + test.AddInput("B", {1, 8, 4}, planes, initializer); + test.AddInput("scales", {1, 1, 5}, std::vector(5, MLFloat16(1.0f))); + test.AddOutput("Y", {1, 5}, std::vector(5)); + test.Run(OpTester::ExpectResult::kExpectFailure, "invalid T8 byte"); + } + } +} + +TEST(MatMulTernaryTest, RejectsUnsupportedVersion) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + test.AddAttribute("format_version", 2); + test.AddInput("A", {1, 128}, std::vector(128)); + test.AddInput("B", {1, 8, 4}, std::vector(32)); + test.AddInput("scales", {1, 1, 5}, std::vector(5)); + test.AddOutput("Y", {1, 5}, std::vector(5)); + test.Run(OpTester::ExpectResult::kExpectFailure, "Unsupported MatMulTernary"); +} + +TEST(MatMulTernaryTest, RejectsPackedShapeMismatch) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + test.AddInput("A", {1, 128}, std::vector(128)); + test.AddInput("B", {1, 8, 3}, std::vector(24)); + test.AddInput("scales", {1, 1, 5}, std::vector(5)); + test.AddOutput("Y", {1, 5}, std::vector(5)); + test.Run(OpTester::ExpectResult::kExpectFailure, "invalid packed tensor dimension"); +} + +TEST(MatMulTernaryTest, RejectsActivationAndScaleShapes) { + for (int invalid = 0; invalid < 3; ++invalid) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + if (invalid == 0) { + test.AddInput("A", {}, {0.0f}); + } else { + const int64_t width = invalid == 1 ? 129 : 128; + test.AddInput("A", {1, width}, std::vector(static_cast(width))); + } + test.AddInput("B", {1, 8, 4}, std::vector(32)); + if (invalid == 2) { + test.AddInput("scales", {1, 5}, std::vector(5)); + } else { + test.AddInput("scales", {1, 1, 5}, std::vector(5)); + } + test.AddOutput("Y", {1, 5}, std::vector(5)); + test.Run(OpTester::ExpectResult::kExpectFailure, + invalid == 0 ? "A must have rank >= 1" : invalid == 1 ? "A last dimension must equal K" + : "B and scales must have rank 3"); + } +} + +TEST(MatMulTernaryTest, RejectsUnsupportedLayout) { + OpTester test("MatMulTernary", 1, kMSDomain); + test.AddAttribute("K", 128); + test.AddAttribute("N", 5); + test.AddAttribute("layout", "invalid"); + test.AddInput("A", {1, 128}, std::vector(128)); + test.AddInput("B", {1, 8, 4}, std::vector(32)); + test.AddInput("scales", {1, 1, 5}, std::vector(5)); + test.AddOutput("Y", {1, 5}, std::vector(5)); + test.Run(OpTester::ExpectResult::kExpectFailure, "Unsupported MatMulTernary"); +} + +} // namespace +} // namespace onnxruntime::test From 9ff62d745fe3b31f8bc64e63f7eebbc499a8029c Mon Sep 17 00:00:00 2001 From: mirounga Date: Mon, 2 Nov 2026 10:25:33 -0800 Subject: [PATCH 2/3] [MLAS] Add AVX2 and AVX-512F T8 ternary GEMM kernels 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. --- cmake/onnxruntime_mlas.cmake | 6 + onnxruntime/core/mlas/lib/mlasi.h | 32 +++ onnxruntime/core/mlas/lib/platform.cpp | 2 + onnxruntime/core/mlas/lib/t8gemm.cpp | 141 +++++++---- onnxruntime/core/mlas/lib/t8gemm_kernel.h | 149 +++++++++++ .../core/mlas/lib/t8gemm_kernel_avx2.cpp | 227 +++++++++++++++++ .../core/mlas/lib/t8gemm_kernel_avx512f.cpp | 233 ++++++++++++++++++ .../test/mlas/unittest/test_t8gemm.cpp | 178 +++++++++++++ 8 files changed, 924 insertions(+), 44 deletions(-) create mode 100644 onnxruntime/core/mlas/lib/t8gemm_kernel.h create mode 100644 onnxruntime/core/mlas/lib/t8gemm_kernel_avx2.cpp create mode 100644 onnxruntime/core/mlas/lib/t8gemm_kernel_avx512f.cpp create mode 100644 onnxruntime/test/mlas/unittest/test_t8gemm.cpp diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 746995d8cacea..302e5c3d0b373 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -62,6 +62,7 @@ onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/qnbitgemm.h ${MLAS_SRC_DIR}/qnbitgemm.cpp ${MLAS_SRC_DIR}/t8gemm.cpp + ${MLAS_SRC_DIR}/t8gemm_kernel.h ${MLAS_SRC_DIR}/qlutgemm.h ${MLAS_SRC_DIR}/qlutgemm.cpp ${MLAS_SRC_DIR}/sqnbitgemm_q8_block.h @@ -261,6 +262,7 @@ function(setup_mlas_source_for_windows) set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "/arch:AVX2") set_source_files_properties(${MLAS_SRC_DIR}/layernorm_kernel_avx2.cpp PROPERTIES COMPILE_FLAGS "/arch:AVX2") + set_source_files_properties(${MLAS_SRC_DIR}/t8gemm_kernel_avx2.cpp PROPERTIES COMPILE_FLAGS "/arch:AVX2") set(mlas_platform_srcs_avx512 ${MLAS_SRC_DIR}/intrinsics/avx512/gelu_avx512f.cpp @@ -268,6 +270,7 @@ function(setup_mlas_source_for_windows) ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp + ${MLAS_SRC_DIR}/t8gemm_kernel_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_depthwise_avx512f.cpp ) @@ -279,6 +282,7 @@ function(setup_mlas_source_for_windows) ${mlas_platform_srcs_avx} ${mlas_platform_srcs_avx2} ${MLAS_SRC_DIR}/layernorm_kernel_avx2.cpp + ${MLAS_SRC_DIR}/t8gemm_kernel_avx2.cpp ${MLAS_SRC_DIR}/rotary_embedding_kernel_avx2.h ${MLAS_SRC_DIR}/rotary_embedding_kernel_avx2.cpp ${MLAS_SRC_DIR}/rotary_embedding_kernel_avx2.cpp @@ -941,6 +945,7 @@ else() ${MLAS_SRC_DIR}/qkv_quant_common.h ${MLAS_SRC_DIR}/qkv_quant_kernel_avx2.cpp ${MLAS_SRC_DIR}/layernorm_kernel_avx2.cpp + ${MLAS_SRC_DIR}/t8gemm_kernel_avx2.cpp ) if(CMAKE_CXX_COMPILER_VERSION GREATER_EQUAL 13.1 AND NOT(APPLE)) set(mlas_platform_srcs_avx2 @@ -1013,6 +1018,7 @@ else() ${MLAS_SRC_DIR}/intrinsics/avx512/quantize_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_nchw_depthwise_multiplier_greater_than_1_avx512f.cpp ${MLAS_SRC_DIR}/linear_attention_kernel_avx512f.cpp + ${MLAS_SRC_DIR}/t8gemm_kernel_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/reorder_avx512f.cpp ${MLAS_SRC_DIR}/intrinsics/avx512/sconv_depthwise_avx512f.cpp ) diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 33a5bb69c0805..cc2f2c3a1527d 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -778,6 +778,32 @@ typedef void(MLASCALL MLAS_LAYERNORM_F16_KERNEL)( bool Simplified ); +// +// T8 ternary GEMM. Weight layouts: T8 v1 bit planes, MlasT8Pack tiled words, +// or MlasT8Pack decoded nonzero/sign masks. The kernel computes output units +// [UnitBegin, UnitEnd) (five columns each) for Rows (1 or 2) rows of A, which +// are contiguous with stride K; C rows have stride N. +// + +enum class MLAS_T8_WEIGHT_LAYOUT { + Planes, + Tiled, + Masks, +}; + +typedef void(MLASCALL MLAS_T8_GEMM_KERNEL)( + size_t Rows, + size_t N, + size_t K, + size_t UnitBegin, + size_t UnitEnd, + const float* A, + const uint32_t* B, + const MLAS_FP16* Scales, + float* C, + MLAS_T8_WEIGHT_LAYOUT Layout +); + typedef void (MLASCALL MLAS_QLINEAR_BINARY_OP_S8_KERNEL)( @@ -1730,6 +1756,11 @@ extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchNeon; extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchSve; extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchRvv; +#if defined(MLAS_TARGET_AMD64) +MLAS_T8_GEMM_KERNEL MlasT8GemmKernelAvx2; +MLAS_T8_GEMM_KERNEL MlasT8GemmKernelAvx512F; +#endif + // // Quantized depthwise convolution kernels. // @@ -1987,6 +2018,7 @@ MLAS_COMPUTE_TANH_FP16_KERNEL* TanhFP16KernelRoutine = nullptr; const MLAS_ELTWISE_DISPATCH* EltwiseDispatch{nullptr}; const MLAS_KV_QUANT_GEMM_DISPATCH* KVQuantGemmDispatch{nullptr}; const MLAS_LINEAR_ATTENTION_DISPATCH* LinearAttentionDispatch{nullptr}; + MLAS_T8_GEMM_KERNEL* T8GemmKernel{nullptr}; }; inline diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index c5f6f3adbf710..3f28132393fee 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -582,6 +582,7 @@ Return Value: // TODO(vraspar): check if this really goes here or if there are other platform reqs that we need to fulfill this->LutGenKernel = &MlasLutGenKernelAvx2; this->LayerNormF32Kernel = &MlasLayerNormKernelAvx2; + this->T8GemmKernel = MlasT8GemmKernelAvx2; if ((Cpuid1[2] & (1u << 29)) != 0) { this->LayerNormF16Kernel = &MlasLayerNormKernelF16Avx2; } @@ -637,6 +638,7 @@ Return Value: this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelAvx512F; this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelAvx512F; this->LinearAttentionDispatch = &MlasLinearAttentionDispatchAvx512F; + this->T8GemmKernel = MlasT8GemmKernelAvx512F; this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelAvx512F; this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelAvx512F; this->NchwcBlockSize = 16; diff --git a/onnxruntime/core/mlas/lib/t8gemm.cpp b/onnxruntime/core/mlas/lib/t8gemm.cpp index cc4f49a54ea5d..a7a389f2d1943 100644 --- a/onnxruntime/core/mlas/lib/t8gemm.cpp +++ b/onnxruntime/core/mlas/lib/t8gemm.cpp @@ -3,6 +3,9 @@ #include "core/mlas/inc/mlas_t8.h" #include "mlasi.h" +#include "t8gemm_kernel.h" + +#include #if defined(MLAS_TARGET_ARM64) #include @@ -165,6 +168,55 @@ ComputeUnit(size_t N, size_t K, size_t unit, const float* A, const uint32_t* Pla } } #else +struct T8OpsScalar { + using V = uint32_t; + + static MLAS_FORCEINLINE V Nt(V a, V b, V c) { return (a ^ b) & (b ^ c); } + static MLAS_FORCEINLINE V Pt(V a, V b, V c, V nt) { return (a ^ b ^ c) & ~nt; } + static MLAS_FORCEINLINE V Bq(V a, V b, V c) { return (a ^ b) & (a | c); } + static MLAS_FORCEINLINE V Neg(V a, V b, V c) { return b ^ ((a ^ b) & c); } + static MLAS_FORCEINLINE V Pos(V a, V b, V c) { return (a ^ c) & ~b; } + static MLAS_FORCEINLINE V Zero(V a, V b, V c) { return ~(a | (b ^ c)); } + static MLAS_FORCEINLINE V Carry(V a, V b, V c) { return a | (b & c); } + static MLAS_FORCEINLINE V And(V a, V b) { return a & b; } + static MLAS_FORCEINLINE V NotOr(V a, V b) { return ~(a | b); } +}; + +// Decodes the four words of one group, whose eight planes are Stride words +// apart, into nonzero masks Nz[j * 8 + w] and negative masks Nz[j * 8 + 4 + w] +// (the MlasT8Pack masks layout). +void +DecodeGroup(const uint32_t* Planes, size_t Stride, uint32_t* Nz) +{ + for (size_t w = 0; w < 4; ++w) { + uint32_t z[5], p[5]; + T8Decode(Planes[w], Planes[Stride + w], Planes[2 * Stride + w], Planes[3 * Stride + w], Planes[4 * Stride + w], Planes[5 * Stride + w], Planes[6 * Stride + w], Planes[7 * Stride + w], z, p); + for (size_t j = 0; j < 5; ++j) { + Nz[j * 8 + w] = ~z[j]; + Nz[j * 8 + 4 + w] = ~(z[j] | p[j]); + } + } +} + +// Adds the signed inputs of one 128-input group selected by the masks +// (masks layout) into Sum, one input at a time in input order. +void +AccumulateGroup(const float* A, const uint32_t* Masks, float* Sum) +{ + for (size_t i = 0; i < 128; ++i) { + uint32_t x; + std::memcpy(&x, A + i, sizeof(x)); + for (size_t j = 0; j < 5; ++j) { + const uint32_t nonzero = (Masks[j * 8 + i / 32] >> (i % 32)) & 1u; + const uint32_t negative = (Masks[j * 8 + 4 + i / 32] >> (i % 32)) & 1u; + const uint32_t bits = (x ^ (negative << 31)) & (0u - nonzero); + float value; + std::memcpy(&value, &bits, sizeof(value)); + Sum[j] += value; + } + } +} + void ComputeUnit(size_t N, size_t K, size_t unit, const float* A, const uint32_t* Planes, const MLAS_FP16* Scales, float* C) { @@ -172,19 +224,10 @@ ComputeUnit(size_t N, size_t K, size_t unit, const float* A, const uint32_t* Pla const uint32_t* planes = Planes + unit * 8 * sets; float total[5]{}; for (size_t g = 0; g < K / 128; ++g) { + uint32_t masks[40]; + DecodeGroup(planes + g * 4, sets, masks); float sum[5]{}; - for (size_t i = 0; i < 128; ++i) { - const size_t input = g * 128 + i; - uint32_t byte = 0; - for (size_t b = 0; b < 8; ++b) { - byte |= ((planes[b * sets + input / 32] >> (input % 32)) & 1u) << b; - } - for (size_t j = 0; j < 5; ++j, byte /= 3) { - const uint32_t trit = byte % 3; - if (trit == 1) sum[j] += A[input]; - if (trit == 2) sum[j] -= A[input]; - } - } + AccumulateGroup(A + g * 128, masks, sum); for (size_t j = 0; j < 5; ++j) { total[j] += sum[j] * Scales[(unit * (K / 128) + g) * 5 + j].ToFloat(); } @@ -253,35 +296,44 @@ ComputePackedUnit(size_t M, size_t N, size_t K, size_t unit, const float* A, con } } #else - for (size_t r = 0; r < std::min(RowTile, M); ++r) { - float total[5]{}; - for (size_t g = 0; g < groups; ++g) { + const size_t rows = std::min(RowTile, M); + float total[RowTile][5]{}; + for (size_t g = 0; g < groups; ++g) { + const uint32_t* block = packed + g * words; + uint32_t decoded[40]; + if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { + DecodeGroup(block, 4, decoded); + block = decoded; + } + for (size_t r = 0; r < rows; ++r) { float sum[5]{}; - const uint32_t* block = packed + g * words; - for (size_t i = 0; i < 128; ++i) { - uint32_t byte = 0; - if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { - for (size_t b = 0; b < 8; ++b) byte |= ((block[b * 4 + i / 32] >> (i % 32)) & 1u) << b; - } - for (size_t j = 0; j < 5; ++j, byte /= 3) { - bool nonzero, negative; - if (Layout == MLAS_T8_PACKED_LAYOUT::Tiled) { - nonzero = byte % 3 != 0; - negative = byte % 3 == 2; - } else { - nonzero = ((block[j * 8 + i / 32] >> (i % 32)) & 1u) != 0; - negative = ((block[j * 8 + 4 + i / 32] >> (i % 32)) & 1u) != 0; - } - if (nonzero) sum[j] += negative ? -A[r * K + g * 128 + i] : A[r * K + g * 128 + i]; - } - } - for (size_t j = 0; j < 5; ++j) total[j] += sum[j] * scales[g * 5 + j].ToFloat(); + AccumulateGroup(A + r * K + g * 128, block, sum); + for (size_t j = 0; j < 5; ++j) total[r][j] += sum[j] * scales[g * 5 + j].ToFloat(); } - for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) C[r * N + unit * 5 + j] = total[j]; + } + for (size_t r = 0; r < rows; ++r) { + for (size_t j = 0; j < std::min(size_t{5}, N - unit * 5); ++j) C[r * N + unit * 5 + j] = total[r][j]; } #endif } +#if !defined(MLAS_TARGET_ARM64) +// Runs the platform T8 kernel over row tiles of RowTile rows and tiles of 16 units. +void +T8GemmWithKernel(size_t M, size_t N, size_t K, const float* A, const uint32_t* B, const MLAS_FP16* Scales, float* C, MLAS_THREADPOOL* ThreadPool, MLAS_T8_WEIGHT_LAYOUT Layout, size_t RowTile) +{ + MLAS_T8_GEMM_KERNEL* kernel = GetMlasPlatform().T8GemmKernel; + const size_t units = N / 5 + (N % 5 != 0); + const size_t tiles = units / 16 + (units % 16 != 0); + const size_t row_tiles = M / RowTile + (M % RowTile != 0); + MlasTrySimpleParallel(ThreadPool, static_cast(row_tiles * tiles), [&](ptrdiff_t task) { + const size_t m = (static_cast(task) / tiles) * RowTile; + const size_t begin = (static_cast(task) % tiles) * 16; + kernel(std::min(RowTile, M - m), N, K, begin, std::min(begin + 16, units), A + m * K, B, Scales, C + m * N, Layout); + }); +} +#endif + } // namespace bool MLASCALL @@ -308,6 +360,10 @@ MlasT8Gemm(size_t M, size_t N, size_t K, const float* A, const uint32_t* Planes, A = Workspace; #else MLAS_UNREFERENCED_PARAMETER(Workspace); + if (GetMlasPlatform().T8GemmKernel != nullptr) { + T8GemmWithKernel(M, N, K, A, Planes, Scales, C, ThreadPool, MLAS_T8_WEIGHT_LAYOUT::Planes, 1); + return; + } #endif const size_t units = N / 5 + (N % 5 != 0); const size_t tiles = units / 16 + (units % 16 != 0); @@ -340,15 +396,7 @@ MlasT8Pack(size_t N, size_t K, const uint32_t* Planes, uint32_t* Packed, MLAS_T8 vst1q_u32(dst + j * 8 + 4, vmvnq_u32(vorrq_u32(z[j], p[j]))); } #else - std::fill_n(dst, 40, 0u); - for (size_t i = 0; i < 128; ++i) { - uint32_t byte = 0; - for (size_t b = 0; b < 8; ++b) byte |= ((src[b * sets + i / 32] >> (i % 32)) & 1u) << b; - for (size_t j = 0; j < 5; ++j, byte /= 3) { - if (byte % 3 != 0) dst[j * 8 + i / 32] |= 1u << (i % 32); - if (byte % 3 == 2) dst[j * 8 + 4 + i / 32] |= 1u << (i % 32); - } - } + DecodeGroup(src, sets, dst); #endif } } @@ -364,6 +412,11 @@ MlasT8GemmPacked(size_t M, size_t N, size_t K, const float* A, const uint32_t* P A = Workspace; #else MLAS_UNREFERENCED_PARAMETER(Workspace); + if (GetMlasPlatform().T8GemmKernel != nullptr) { + const auto layout = Layout == MLAS_T8_PACKED_LAYOUT::Masks ? MLAS_T8_WEIGHT_LAYOUT::Masks : MLAS_T8_WEIGHT_LAYOUT::Tiled; + T8GemmWithKernel(M, N, K, A, Packed, Scales, C, ThreadPool, layout, RowTile); + return; + } #endif const size_t units = N / 5 + (N % 5 != 0); const size_t tiles = units / 16 + (units % 16 != 0); diff --git a/onnxruntime/core/mlas/lib/t8gemm_kernel.h b/onnxruntime/core/mlas/lib/t8gemm_kernel.h new file mode 100644 index 0000000000000..8583769d48270 --- /dev/null +++ b/onnxruntime/core/mlas/lib/t8gemm_kernel.h @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/*++ + +Module Name: + + t8gemm_kernel.h + +Abstract: + + Shared pieces of the portable, AVX2 and AVX-512F T8 ternary GEMM kernels. + + Each T8 byte holds five trits in base 3 and is stored as eight bit planes, + so one 32-bit plane word covers 32 consecutive inputs. The decode below is + the bitsliced gate network used by the NEON kernel: it turns the eight + plane words (or vectors of them) of a group directly into per-trit zero + and positive masks. The kernels then use those masks to add or subtract A + in natural input order. + + Everything here has internal linkage: the header is included by + translation units compiled for different instruction sets. + +--*/ + +#pragma once + +#include "mlasi.h" + +namespace +{ + +constexpr size_t T8GroupSize = 128; +constexpr size_t T8WordsPerGroup = 4; +constexpr size_t T8TritsPerUnit = 5; +constexpr size_t T8TiledWordsPerGroup = 32; +constexpr size_t T8MaskWordsPerGroup = 40; + +// +// Bitsliced decode of eight T8 bit planes into five trits. Ops supplies the +// three-input gates for one vector type; Z receives the "trit is 0" masks and +// P the "trit is 1" masks. +// + +template +MLAS_FORCEINLINE void +T8Decode( + typename Ops::V b0, + typename Ops::V b1, + typename Ops::V b2, + typename Ops::V b3, + typename Ops::V b4, + typename Ops::V b5, + typename Ops::V b6, + typename Ops::V b7, + typename Ops::V* z, + typename Ops::V* p +) +{ + using V = typename Ops::V; + + V nt = Ops::Nt(b7, b6, b5), pt = Ops::Pt(b7, b6, b5, nt); + b5 = Ops::Bq(b7, b6, b5); + b6 = Ops::And(b6, b7); + V n = Ops::Neg(nt, pt, b4), q = Ops::Pos(nt, pt, b4); + b4 = Ops::Carry(nt, pt, b4); + nt = Ops::Neg(n, q, b3); + pt = Ops::Pos(n, q, b3); + b3 = Ops::Carry(n, q, b3); + n = Ops::Neg(nt, pt, b2); + q = Ops::Pos(nt, pt, b2); + b2 = Ops::Carry(nt, pt, b2); + nt = Ops::Neg(n, q, b1); + pt = Ops::Pos(n, q, b1); + b1 = Ops::Carry(n, q, b1); + z[0] = Ops::Zero(nt, pt, b0); + p[0] = Ops::Pos(nt, pt, b0); + b0 = Ops::Carry(nt, pt, b0); + + nt = Ops::Nt(b6, b5, b4); + pt = Ops::Pt(b6, b5, b4, nt); + b4 = Ops::Bq(b6, b5, b4); + b5 = Ops::And(b5, b6); + n = Ops::Neg(nt, pt, b3); + q = Ops::Pos(nt, pt, b3); + b3 = Ops::Carry(nt, pt, b3); + nt = Ops::Neg(n, q, b2); + pt = Ops::Pos(n, q, b2); + b2 = Ops::Carry(n, q, b2); + n = Ops::Neg(nt, pt, b1); + q = Ops::Pos(nt, pt, b1); + b1 = Ops::Carry(nt, pt, b1); + z[1] = Ops::Zero(n, q, b0); + p[1] = Ops::Pos(n, q, b0); + b0 = Ops::Carry(n, q, b0); + + nt = Ops::Nt(b4, b3, b2); + pt = Ops::Pt(b4, b3, b2, nt); + b2 = Ops::Bq(b4, b3, b2); + b3 = Ops::And(b3, b4); + n = Ops::Neg(nt, pt, b1); + q = Ops::Pos(nt, pt, b1); + b1 = Ops::Carry(nt, pt, b1); + z[2] = Ops::Zero(n, q, b0); + p[2] = Ops::Pos(n, q, b0); + b0 = Ops::Carry(n, q, b0); + + nt = Ops::Nt(b3, b2, b1); + pt = Ops::Pt(b3, b2, b1, nt); + V n4 = Ops::Bq(b3, b2, b1); + z[3] = Ops::Zero(nt, pt, b0); + p[3] = Ops::Pos(nt, pt, b0); + V p4 = Ops::Carry(nt, pt, b0); + z[4] = Ops::NotOr(n4, p4); + p[4] = p4; +} + +// +// Truth-table immediates for vpternlogd, with the operands in order a, b, c. +// + +constexpr uint8_t T8TernA = 0xF0; +constexpr uint8_t T8TernB = 0xCC; +constexpr uint8_t T8TernC = 0xAA; + +// +// Returns the address of word 0 of group g of plane b, for the two layouts +// that store bit planes. Each group's four words for one plane are adjacent. +// + +MLAS_FORCEINLINE const uint32_t* +T8PlaneWords( + const uint32_t* B, + MLAS_T8_WEIGHT_LAYOUT Layout, + size_t K, + size_t unit, + size_t g, + size_t b +) +{ + const size_t groups = K / T8GroupSize; + if (Layout == MLAS_T8_WEIGHT_LAYOUT::Planes) { + const size_t sets = K / 32; + return B + (unit * 8 + b) * sets + g * T8WordsPerGroup; + } + return B + (unit * groups + g) * T8TiledWordsPerGroup + b * T8WordsPerGroup; +} + +} // namespace diff --git a/onnxruntime/core/mlas/lib/t8gemm_kernel_avx2.cpp b/onnxruntime/core/mlas/lib/t8gemm_kernel_avx2.cpp new file mode 100644 index 0000000000000..ffcff07fefd96 --- /dev/null +++ b/onnxruntime/core/mlas/lib/t8gemm_kernel_avx2.cpp @@ -0,0 +1,227 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/*++ + +Module Name: + + t8gemm_kernel_avx2.cpp + +Abstract: + + AVX2 T8 ternary GEMM kernel. + + Two groups (eight plane words) are decoded per 256-bit vector. To use a + decoded word on eight consecutive inputs, the word is broadcast and each + lane shifts its own bit into the sign position: an arithmetic shift turns + the nonzero bit into a lane mask, and the negative bit becomes the sign + flipped into A. + +--*/ + +#include "t8gemm_kernel.h" + +namespace +{ + +struct T8OpsAvx2 { + using V = __m256i; + + static MLAS_FORCEINLINE V Not(V a) { return _mm256_xor_si256(a, _mm256_set1_epi32(-1)); } + + static MLAS_FORCEINLINE V Nt(V a, V b, V c) + { + return _mm256_and_si256(_mm256_xor_si256(a, b), _mm256_xor_si256(b, c)); + } + static MLAS_FORCEINLINE V Pt(V a, V b, V c, V nt) + { + return _mm256_andnot_si256(nt, _mm256_xor_si256(_mm256_xor_si256(a, b), c)); + } + static MLAS_FORCEINLINE V Bq(V a, V b, V c) + { + return _mm256_and_si256(_mm256_xor_si256(a, b), _mm256_or_si256(a, c)); + } + static MLAS_FORCEINLINE V Neg(V a, V b, V c) + { + return _mm256_xor_si256(b, _mm256_and_si256(_mm256_xor_si256(a, b), c)); + } + static MLAS_FORCEINLINE V Pos(V a, V b, V c) { return _mm256_andnot_si256(b, _mm256_xor_si256(a, c)); } + static MLAS_FORCEINLINE V Zero(V a, V b, V c) { return Not(_mm256_or_si256(a, _mm256_xor_si256(b, c))); } + static MLAS_FORCEINLINE V Carry(V a, V b, V c) { return _mm256_or_si256(a, _mm256_and_si256(b, c)); } + static MLAS_FORCEINLINE V And(V a, V b) { return _mm256_and_si256(a, b); } + static MLAS_FORCEINLINE V NotOr(V a, V b) { return Not(_mm256_or_si256(a, b)); } +}; + +constexpr size_t GroupsPerBlock = 2; + +// +// Loads plane b of groups [g0, g0 + ng) of one unit; lanes of a missing group +// are zero. +// + +MLAS_FORCEINLINE __m256i +LoadPlane(const uint32_t* B, MLAS_T8_WEIGHT_LAYOUT Layout, size_t K, size_t unit, size_t g0, size_t ng, size_t b) +{ + const uint32_t* p = T8PlaneWords(B, Layout, K, unit, g0, b); + if (Layout == MLAS_T8_WEIGHT_LAYOUT::Planes && ng == GroupsPerBlock) { + return _mm256_loadu_si256(reinterpret_cast(p)); + } + const __m256i lo = _mm256_zextsi128_si256(_mm_loadu_si128(reinterpret_cast(p))); + if (ng == 1) { + return lo; + } + const size_t stride = Layout == MLAS_T8_WEIGHT_LAYOUT::Planes ? T8WordsPerGroup : T8TiledWordsPerGroup; + return _mm256_inserti128_si256(lo, _mm_loadu_si128(reinterpret_cast(p + stride)), 1); +} + +MLAS_FORCEINLINE float +ReduceAdd(__m256 v) +{ + __m128 s = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); + s = _mm_add_ps(s, _mm_movehl_ps(s, s)); + s = _mm_add_ss(s, _mm_movehdup_ps(s)); + return _mm_cvtss_f32(s); +} + +// +// Accumulates one 128-input group for Rows rows. Nz[j] and Neg[j] point to the +// four nonzero and negative mask words of trit j. +// + +template +MLAS_FORCEINLINE void +ComputeGroup( + const float* A, + size_t K, + const uint32_t* const* Nz, + const uint32_t* const* Neg, + const MLAS_FP16* Scales, + __m256 (&Totals)[Rows][T8TritsPerUnit] +) +{ + // Lane l of chunk c holds input 8 * c + l of a word; shifting by + // 31 - (8 * c + l) moves that input's bit into the sign position. + const __m256i shifts[4] = { + _mm256_setr_epi32(31, 30, 29, 28, 27, 26, 25, 24), + _mm256_setr_epi32(23, 22, 21, 20, 19, 18, 17, 16), + _mm256_setr_epi32(15, 14, 13, 12, 11, 10, 9, 8), + _mm256_setr_epi32(7, 6, 5, 4, 3, 2, 1, 0), + }; + const __m256i sign = _mm256_set1_epi32(static_cast(0x80000000u)); + + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + __m256 sums[Rows][2]; + for (size_t r = 0; r < Rows; ++r) { + sums[r][0] = _mm256_setzero_ps(); + sums[r][1] = _mm256_setzero_ps(); + } + for (size_t w = 0; w < T8WordsPerGroup; ++w) { + const __m256i nonzero_bits = _mm256_set1_epi32(static_cast(Nz[j][w])); + const __m256i negative_bits = _mm256_set1_epi32(static_cast(Neg[j][w])); + for (size_t c = 0; c < 4; ++c) { + const __m256 nonzero = _mm256_castsi256_ps(_mm256_srai_epi32(_mm256_sllv_epi32(nonzero_bits, shifts[c]), 31)); + const __m256 negative = _mm256_castsi256_ps(_mm256_and_si256(_mm256_sllv_epi32(negative_bits, shifts[c]), sign)); + for (size_t r = 0; r < Rows; ++r) { + const __m256 x = _mm256_loadu_ps(A + r * K + w * 32 + c * 8); + sums[r][c & 1] = _mm256_add_ps(sums[r][c & 1], _mm256_and_ps(_mm256_xor_ps(x, negative), nonzero)); + } + } + } + const __m256 scale = _mm256_cvtph_ps(_mm_set1_epi16(static_cast(Scales[j].val))); + for (size_t r = 0; r < Rows; ++r) { + Totals[r][j] = _mm256_fmadd_ps(_mm256_add_ps(sums[r][0], sums[r][1]), scale, Totals[r][j]); + } + } +} + +template +void +T8GemmUnits( + size_t N, + size_t K, + size_t UnitBegin, + size_t UnitEnd, + const float* A, + const uint32_t* B, + const MLAS_FP16* Scales, + float* C, + MLAS_T8_WEIGHT_LAYOUT Layout +) +{ + const size_t groups = K / T8GroupSize; + alignas(32) uint32_t nz_words[T8TritsPerUnit][GroupsPerBlock * T8WordsPerGroup]; + alignas(32) uint32_t neg_words[T8TritsPerUnit][GroupsPerBlock * T8WordsPerGroup]; + + for (size_t unit = UnitBegin; unit < UnitEnd; ++unit) { + __m256 totals[Rows][T8TritsPerUnit]; + for (size_t r = 0; r < Rows; ++r) { + for (size_t j = 0; j < T8TritsPerUnit; ++j) totals[r][j] = _mm256_setzero_ps(); + } + const MLAS_FP16* scales = Scales + unit * groups * T8TritsPerUnit; + + for (size_t g0 = 0; g0 < groups; g0 += GroupsPerBlock) { + const size_t ng = std::min(GroupsPerBlock, groups - g0); + if (Layout != MLAS_T8_WEIGHT_LAYOUT::Masks) { + __m256i z[T8TritsPerUnit], p[T8TritsPerUnit]; + T8Decode( + LoadPlane(B, Layout, K, unit, g0, ng, 0), LoadPlane(B, Layout, K, unit, g0, ng, 1), + LoadPlane(B, Layout, K, unit, g0, ng, 2), LoadPlane(B, Layout, K, unit, g0, ng, 3), + LoadPlane(B, Layout, K, unit, g0, ng, 4), LoadPlane(B, Layout, K, unit, g0, ng, 5), + LoadPlane(B, Layout, K, unit, g0, ng, 6), LoadPlane(B, Layout, K, unit, g0, ng, 7), z, p + ); + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + const __m256i nonzero = T8OpsAvx2::Not(z[j]); + _mm256_store_si256(reinterpret_cast<__m256i*>(nz_words[j]), nonzero); + _mm256_store_si256(reinterpret_cast<__m256i*>(neg_words[j]), _mm256_andnot_si256(p[j], nonzero)); + } + } + for (size_t gg = 0; gg < ng; ++gg) { + const size_t g = g0 + gg; + const uint32_t* nz[T8TritsPerUnit]; + const uint32_t* neg[T8TritsPerUnit]; + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + if (Layout == MLAS_T8_WEIGHT_LAYOUT::Masks) { + const uint32_t* block = B + (unit * groups + g) * T8MaskWordsPerGroup; + nz[j] = block + j * 8; + neg[j] = block + j * 8 + 4; + } else { + nz[j] = nz_words[j] + gg * T8WordsPerGroup; + neg[j] = neg_words[j] + gg * T8WordsPerGroup; + } + } + ComputeGroup(A + g * T8GroupSize, K, nz, neg, scales + g * T8TritsPerUnit, totals); + } + } + + const size_t count = std::min(T8TritsPerUnit, N - unit * T8TritsPerUnit); + for (size_t r = 0; r < Rows; ++r) { + for (size_t j = 0; j < count; ++j) { + C[r * N + unit * T8TritsPerUnit + j] = ReduceAdd(totals[r][j]); + } + } + } +} + +} // namespace + +void + MLASCALL + MlasT8GemmKernelAvx2( + size_t Rows, + size_t N, + size_t K, + size_t UnitBegin, + size_t UnitEnd, + const float* A, + const uint32_t* B, + const MLAS_FP16* Scales, + float* C, + MLAS_T8_WEIGHT_LAYOUT Layout + ) +{ + if (Rows == 2) { + T8GemmUnits<2>(N, K, UnitBegin, UnitEnd, A, B, Scales, C, Layout); + } else { + T8GemmUnits<1>(N, K, UnitBegin, UnitEnd, A, B, Scales, C, Layout); + } +} diff --git a/onnxruntime/core/mlas/lib/t8gemm_kernel_avx512f.cpp b/onnxruntime/core/mlas/lib/t8gemm_kernel_avx512f.cpp new file mode 100644 index 0000000000000..59813a2ecafb3 --- /dev/null +++ b/onnxruntime/core/mlas/lib/t8gemm_kernel_avx512f.cpp @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +/*++ + +Module Name: + + t8gemm_kernel_avx512f.cpp + +Abstract: + + AVX-512F T8 ternary GEMM kernel. + + Four groups (sixteen plane words) are decoded per 512-bit vector with + vpternlogd gates. Each 16-bit half of a decoded word is used directly as an + opmask: the sign mask flips A with a masked xor and the nonzero mask selects + the lanes added into the accumulator. + +--*/ + +#include "t8gemm_kernel.h" + +namespace +{ + +struct T8OpsAvx512 { + using V = __m512i; + + template + static MLAS_FORCEINLINE V Tern(V a, V b, V c) + { + return _mm512_ternarylogic_epi32(a, b, c, Imm); + } + + static MLAS_FORCEINLINE V Nt(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V Pt(V a, V b, V c, V nt) + { + return _mm512_andnot_si512(nt, Tern(a, b, c)); + } + static MLAS_FORCEINLINE V Bq(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V Neg(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V Pos(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V Zero(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V Carry(V a, V b, V c) + { + return Tern(a, b, c); + } + static MLAS_FORCEINLINE V And(V a, V b) { return _mm512_and_si512(a, b); } + static MLAS_FORCEINLINE V NotOr(V a, V b) { return Tern(a, b, b); } +}; + +constexpr size_t GroupsPerBlock = 4; + +// +// Loads plane b of groups [g0, g0 + ng) of one unit; lanes of missing groups +// are zero. +// + +MLAS_FORCEINLINE __m512i +LoadPlane(const uint32_t* B, MLAS_T8_WEIGHT_LAYOUT Layout, size_t K, size_t unit, size_t g0, size_t ng, size_t b) +{ + const uint32_t* p = T8PlaneWords(B, Layout, K, unit, g0, b); + if (Layout == MLAS_T8_WEIGHT_LAYOUT::Planes) { + if (ng == GroupsPerBlock) { + return _mm512_loadu_si512(p); + } + return _mm512_maskz_loadu_epi32(static_cast<__mmask16>((1u << (ng * T8WordsPerGroup)) - 1), p); + } + __m512i v = _mm512_zextsi128_si512(_mm_loadu_si128(reinterpret_cast(p))); + if (ng > 1) v = _mm512_inserti32x4(v, _mm_loadu_si128(reinterpret_cast(p + T8TiledWordsPerGroup)), 1); + if (ng > 2) v = _mm512_inserti32x4(v, _mm_loadu_si128(reinterpret_cast(p + 2 * T8TiledWordsPerGroup)), 2); + if (ng > 3) v = _mm512_inserti32x4(v, _mm_loadu_si128(reinterpret_cast(p + 3 * T8TiledWordsPerGroup)), 3); + return v; +} + +// +// Accumulates one 128-input group for Rows rows. Nz[j] and Neg[j] point to the +// four nonzero and negative mask words of trit j. +// + +template +MLAS_FORCEINLINE void +ComputeGroup( + const float* A, + size_t K, + const uint32_t* const* Nz, + const uint32_t* const* Neg, + const MLAS_FP16* Scales, + __m512 (&Totals)[Rows][T8TritsPerUnit] +) +{ + // Two accumulators per trit hide the add latency. Every row tile uses the + // same split, so results do not depend on the row tile. + constexpr size_t Chains = 2; + const __m512i sign = _mm512_set1_epi32(static_cast(0x80000000u)); + __m512 sums[Rows][T8TritsPerUnit][Chains]; + for (size_t r = 0; r < Rows; ++r) { + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + for (size_t c = 0; c < Chains; ++c) sums[r][j][c] = _mm512_setzero_ps(); + } + } + for (size_t w = 0; w < T8WordsPerGroup; ++w) { + for (size_t h = 0; h < 2; ++h) { + __m512i x[Rows]; + for (size_t r = 0; r < Rows; ++r) { + x[r] = _mm512_castps_si512(_mm512_loadu_ps(A + r * K + w * 32 + h * 16)); + } + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + const __mmask16 nonzero = static_cast<__mmask16>(Nz[j][w] >> (16 * h)); + const __mmask16 negative = static_cast<__mmask16>(Neg[j][w] >> (16 * h)); + for (size_t r = 0; r < Rows; ++r) { + const __m512 xs = _mm512_castsi512_ps(_mm512_mask_xor_epi32(x[r], negative, x[r], sign)); + __m512& sum = sums[r][j][h]; + sum = _mm512_mask_add_ps(sum, nonzero, sum, xs); + } + } + } + } + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + const __m512 scale = _mm512_cvtph_ps(_mm256_set1_epi16(static_cast(Scales[j].val))); + for (size_t r = 0; r < Rows; ++r) { + const __m512 sum = _mm512_add_ps(sums[r][j][0], sums[r][j][1]); + Totals[r][j] = _mm512_fmadd_ps(sum, scale, Totals[r][j]); + } + } +} + +template +void +T8GemmUnits( + size_t N, + size_t K, + size_t UnitBegin, + size_t UnitEnd, + const float* A, + const uint32_t* B, + const MLAS_FP16* Scales, + float* C, + MLAS_T8_WEIGHT_LAYOUT Layout +) +{ + const size_t groups = K / T8GroupSize; + alignas(64) uint32_t nz_words[T8TritsPerUnit][GroupsPerBlock * T8WordsPerGroup]; + alignas(64) uint32_t neg_words[T8TritsPerUnit][GroupsPerBlock * T8WordsPerGroup]; + + for (size_t unit = UnitBegin; unit < UnitEnd; ++unit) { + __m512 totals[Rows][T8TritsPerUnit]; + for (size_t r = 0; r < Rows; ++r) { + for (size_t j = 0; j < T8TritsPerUnit; ++j) totals[r][j] = _mm512_setzero_ps(); + } + const MLAS_FP16* scales = Scales + unit * groups * T8TritsPerUnit; + + for (size_t g0 = 0; g0 < groups; g0 += GroupsPerBlock) { + const size_t ng = std::min(GroupsPerBlock, groups - g0); + if (Layout != MLAS_T8_WEIGHT_LAYOUT::Masks) { + __m512i z[T8TritsPerUnit], p[T8TritsPerUnit]; + T8Decode( + LoadPlane(B, Layout, K, unit, g0, ng, 0), LoadPlane(B, Layout, K, unit, g0, ng, 1), + LoadPlane(B, Layout, K, unit, g0, ng, 2), LoadPlane(B, Layout, K, unit, g0, ng, 3), + LoadPlane(B, Layout, K, unit, g0, ng, 4), LoadPlane(B, Layout, K, unit, g0, ng, 5), + LoadPlane(B, Layout, K, unit, g0, ng, 6), LoadPlane(B, Layout, K, unit, g0, ng, 7), z, p + ); + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + const __m512i nonzero = _mm512_ternarylogic_epi32(z[j], z[j], z[j], uint8_t(~T8TernA)); + _mm512_store_si512(nz_words[j], nonzero); + _mm512_store_si512(neg_words[j], _mm512_andnot_si512(p[j], nonzero)); + } + } + for (size_t gg = 0; gg < ng; ++gg) { + const size_t g = g0 + gg; + const uint32_t* nz[T8TritsPerUnit]; + const uint32_t* neg[T8TritsPerUnit]; + for (size_t j = 0; j < T8TritsPerUnit; ++j) { + if (Layout == MLAS_T8_WEIGHT_LAYOUT::Masks) { + const uint32_t* block = B + (unit * groups + g) * T8MaskWordsPerGroup; + nz[j] = block + j * 8; + neg[j] = block + j * 8 + 4; + } else { + nz[j] = nz_words[j] + gg * T8WordsPerGroup; + neg[j] = neg_words[j] + gg * T8WordsPerGroup; + } + } + ComputeGroup(A + g * T8GroupSize, K, nz, neg, scales + g * T8TritsPerUnit, totals); + } + } + + const size_t count = std::min(T8TritsPerUnit, N - unit * T8TritsPerUnit); + for (size_t r = 0; r < Rows; ++r) { + for (size_t j = 0; j < count; ++j) { + C[r * N + unit * T8TritsPerUnit + j] = _mm512_reduce_add_ps(totals[r][j]); + } + } + } +} + +} // namespace + +void + MLASCALL + MlasT8GemmKernelAvx512F( + size_t Rows, + size_t N, + size_t K, + size_t UnitBegin, + size_t UnitEnd, + const float* A, + const uint32_t* B, + const MLAS_FP16* Scales, + float* C, + MLAS_T8_WEIGHT_LAYOUT Layout + ) +{ + if (Rows == 2) { + T8GemmUnits<2>(N, K, UnitBegin, UnitEnd, A, B, Scales, C, Layout); + } else { + T8GemmUnits<1>(N, K, UnitBegin, UnitEnd, A, B, Scales, C, Layout); + } +} diff --git a/onnxruntime/test/mlas/unittest/test_t8gemm.cpp b/onnxruntime/test/mlas/unittest/test_t8gemm.cpp new file mode 100644 index 0000000000000..847e32456736e --- /dev/null +++ b/onnxruntime/test/mlas/unittest/test_t8gemm.cpp @@ -0,0 +1,178 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include + +#include "test_util.h" +#include "core/mlas/inc/mlas_t8.h" +#include "core/mlas/lib/mlasi.h" + +// +// Checks MlasT8Gemm and MlasT8GemmPacked against a double-precision reference +// over shapes that cover padded output units, odd group counts (the x86 +// kernels decode several groups at once) and both row tiles. On x86 each +// available SIMD kernel is also called directly, so the AVX2 kernel is tested +// on AVX-512 hosts too. +// + +class MlasT8GemmTest : public MlasTestBase { + private: + struct Problem { + size_t M, N, K; + std::vector A; + std::vector Planes; + std::vector Scales; + std::vector Reference; + std::vector Magnitude; // sum of |a * w|, bounds the rounding error + }; + + std::mt19937 Rng{1234}; + + Problem Generate(size_t M, size_t N, size_t K) { + Problem p{M, N, K, {}, {}, {}, {}, {}}; + const size_t units = (N + 4) / 5, sets = K / 32, groups = K / 128; + std::uniform_int_distribution trit_dist(0, 2); + std::uniform_real_distribution scale_dist(0.01f, 0.1f); + std::uniform_real_distribution a_dist(-1.0f, 1.0f); + + std::vector trits(units * 5 * K, 0); // [unit*5 + j][k], padded rows stay 0 + for (size_t row = 0; row < N; ++row) { + for (size_t k = 0; k < K; ++k) trits[row * K + k] = trit_dist(Rng); + } + p.Planes.assign(units * 8 * sets, 0); + for (size_t q = 0; q < units; ++q) { + for (size_t k = 0; k < K; ++k) { + uint32_t byte = 0; + for (size_t j = 5; j-- > 0;) byte = byte * 3 + static_cast(trits[(q * 5 + j) * K + k]); + for (size_t b = 0; b < 8; ++b) p.Planes[(q * 8 + b) * sets + k / 32] |= ((byte >> b) & 1u) << (k % 32); + } + } + p.Scales.resize(units * groups * 5); + for (auto& s : p.Scales) s = MLAS_FP16(scale_dist(Rng)); + p.A.resize(M * K); + for (auto& a : p.A) a = a_dist(Rng); + + p.Reference.assign(M * N, 0.0); + p.Magnitude.assign(M * N, 0.0); + for (size_t m = 0; m < M; ++m) { + for (size_t n = 0; n < N; ++n) { + for (size_t k = 0; k < K; ++k) { + const int t = trits[n * K + k]; + const double w = (t == 1 ? 1.0 : t == 2 ? -1.0 + : 0.0) * + static_cast(p.Scales[((n / 5) * groups + k / 128) * 5 + n % 5].ToFloat()); + p.Reference[m * N + n] += static_cast(p.A[m * K + k]) * w; + p.Magnitude[m * N + n] += std::abs(static_cast(p.A[m * K + k]) * w); + } + } + } + return p; + } + + static void Check(const Problem& p, const std::vector& C, const char* what) { + for (size_t i = 0; i < p.M * p.N; ++i) { + const double tolerance = 1e-5 * p.Magnitude[i] + 1e-6; + ASSERT_NEAR(C[i], p.Reference[i], tolerance) + << what << " M=" << p.M << " N=" << p.N << " K=" << p.K << " index=" << i; + } + } + + void Test(size_t M, size_t N, size_t K) { + const Problem p = Generate(M, N, K); + TestDispatched(p); +#if defined(MLAS_TARGET_AMD64) + // Clearing the platform kernel routes the public entry points to the + // portable fallback. + MLAS_T8_GEMM_KERNEL* kernel = GetMlasPlatform().T8GemmKernel; + if (kernel != nullptr) { + GetMlasPlatform().T8GemmKernel = nullptr; + TestDispatched(p); + GetMlasPlatform().T8GemmKernel = kernel; + } + TestKernels(p); +#endif + } + + void TestDispatched(const Problem& p) { + const size_t M = p.M, N = p.N, K = p.K; + const size_t units = (N + 4) / 5, groups = K / 128; + std::vector workspace(M * K); + std::vector c(M * N); + + MlasT8Gemm(M, N, K, p.A.data(), p.Planes.data(), p.Scales.data(), c.data(), workspace.data(), nullptr); + Check(p, c, "MlasT8Gemm"); + + std::vector tiled(units * groups * 32), masks(units * groups * 40); + MlasT8Pack(N, K, p.Planes.data(), tiled.data(), MLAS_T8_PACKED_LAYOUT::Tiled); + MlasT8Pack(N, K, p.Planes.data(), masks.data(), MLAS_T8_PACKED_LAYOUT::Masks); + for (size_t row_tile : {size_t{1}, size_t{2}}) { + std::fill(c.begin(), c.end(), 0.0f); + MlasT8GemmPacked(M, N, K, p.A.data(), tiled.data(), p.Scales.data(), c.data(), workspace.data(), nullptr, + MLAS_T8_PACKED_LAYOUT::Tiled, row_tile); + Check(p, c, row_tile == 1 ? "Tiled/1" : "Tiled/2"); + std::fill(c.begin(), c.end(), 0.0f); + MlasT8GemmPacked(M, N, K, p.A.data(), masks.data(), p.Scales.data(), c.data(), workspace.data(), nullptr, + MLAS_T8_PACKED_LAYOUT::Masks, row_tile); + Check(p, c, row_tile == 1 ? "Masks/1" : "Masks/2"); + } + } + +#if defined(MLAS_TARGET_AMD64) + void TestKernels(const Problem& p) { + const size_t M = p.M, N = p.N, K = p.K; + const size_t units = (N + 4) / 5, groups = K / 128; + std::vector c(M * N); + std::vector tiled(units * groups * 32), masks(units * groups * 40); + MlasT8Pack(N, K, p.Planes.data(), tiled.data(), MLAS_T8_PACKED_LAYOUT::Tiled); + MlasT8Pack(N, K, p.Planes.data(), masks.data(), MLAS_T8_PACKED_LAYOUT::Masks); + + // Each SIMD kernel must give the same bits for every layout and row tile. + std::vector> kernels; + if (GetMlasPlatform().T8GemmKernel != nullptr) kernels.emplace_back(MlasT8GemmKernelAvx2, "Avx2"); + if (GetMlasPlatform().T8GemmKernel == MlasT8GemmKernelAvx512F) kernels.emplace_back(MlasT8GemmKernelAvx512F, "Avx512F"); + for (const auto& [kernel, name] : kernels) { + std::vector first; + for (const auto layout : {MLAS_T8_WEIGHT_LAYOUT::Planes, MLAS_T8_WEIGHT_LAYOUT::Tiled, MLAS_T8_WEIGHT_LAYOUT::Masks}) { + const uint32_t* b = layout == MLAS_T8_WEIGHT_LAYOUT::Planes ? p.Planes.data() + : layout == MLAS_T8_WEIGHT_LAYOUT::Tiled ? tiled.data() + : masks.data(); + for (size_t rows : {size_t{1}, size_t{2}}) { + std::fill(c.begin(), c.end(), 0.0f); + for (size_t m = 0; m < M; m += rows) { + kernel(std::min(rows, M - m), N, K, 0, units, p.A.data() + m * K, b, p.Scales.data(), c.data() + m * N, layout); + } + Check(p, c, name); + if (first.empty()) { + first = c; + } else { + ASSERT_EQ(std::memcmp(first.data(), c.data(), c.size() * sizeof(float)), 0) + << name << " layout=" << static_cast(layout) << " rows=" << rows << " M=" << M << " N=" << N << " K=" << K; + } + } + } + } + } +#endif + + public: + static const char* GetTestSuiteName() { + static const std::string suite_name("T8Gemm"); + return suite_name.c_str(); + } + + void ExecuteShort(void) override { + for (size_t M : {size_t{1}, size_t{2}, size_t{3}, size_t{5}}) { + for (size_t N : {size_t{1}, size_t{5}, size_t{7}, size_t{23}, size_t{84}}) { + for (size_t K : {size_t{128}, size_t{256}, size_t{384}, size_t{640}, size_t{1152}}) { + Test(M, N, K); + } + } + } + } +}; + +static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { + return is_short_execute ? MlasDirectShortExecuteTests::RegisterShortExecute() : 0; +}); From b8139381d6e21f1e807ac22d37c5c11a5c587c39 Mon Sep 17 00:00:00 2001 From: mirounga Date: Tue, 3 Nov 2026 15:15:21 -0800 Subject: [PATCH 3/3] [Bonsai] Validate zero points of every node before reusing a cached T8 matrix convert_to_t8.py keys converted matrices by (weight, scales, N, K) and checked zero points only when a matrix was converted for the first time. A second MatMulNBits node sharing B and scales but carrying a different zero point (e.g. 0) was then silently rewritten as if its zero point were 1, changing its results despite the lossless-conversion report. Move the zero-point dtype/shape/value checks ahead of the cache lookup so they run for every node, and add a regression test with two nodes that share B and scales but differ in zero points. Addresses review comment r4446345129. --- .../models/bonsai/convert_to_t8.py | 26 ++++++++------- .../transformers/models/bonsai/test_t8.py | 32 +++++++++++++++++++ 2 files changed, 47 insertions(+), 11 deletions(-) diff --git a/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py index 49980637e9d30..1a7de9600e8d7 100644 --- a/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py +++ b/onnxruntime/python/tools/transformers/models/bonsai/convert_to_t8.py @@ -60,21 +60,25 @@ def store(tensor: onnx.TensorProto) -> None: raise ValueError(f"{node.name}: expected explicit zero point and no g_idx or bias") weight, scales, zeros = (resolve(name) for name in node.input[1:4]) key = (weight.name, scales.name, int(attrs["N"]), int(attrs["K"])) + n, k = key[2:] + if k <= 0 or k % 128: + raise ValueError(f"{node.name}: Q2 shape disagrees with K/N") + # Zero points are not part of the matrix cache key, so validate them for every node. + zp = array(zeros) + 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") if key not in matrices: - packed, s, zp = array(weight), array(scales), array(zeros) + packed, s = array(weight), array(scales) 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: + if packed.shape != (n, k // 128, 32): 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) names = (weight.name + ".t8", scales.name + ".t8") if any(name in initializers for name in names): diff --git a/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py b/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py index c584084f65b9a..0f45ff7ed27b1 100644 --- a/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py +++ b/onnxruntime/python/tools/transformers/models/bonsai/test_t8.py @@ -169,6 +169,38 @@ def test_external_data_and_shared_initializers(self): with self.assertRaisesRegex(ValueError, "never overwrites"): convert_model(source, output) + def test_validates_zero_points_of_cached_matrices(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + tensors = [ + numpy_helper.from_array(np.full((7, 1, 32), 0x24, dtype=np.uint8), "B"), + numpy_helper.from_array(np.arange(1, 8, dtype=np.float32)[:, None] / 16, "S"), + numpy_helper.from_array(np.full((7, 1), 0x55, dtype=np.uint8), "Z1"), + numpy_helper.from_array(np.zeros((7, 1), dtype=np.uint8), "Z0"), + ] + attrs = {"K": 128, "N": 7, "bits": 2, "block_size": 128} + nodes = [ + helper.make_node("MatMulNBits", ["A", "B", "S", "Z1"], ["Y1"], domain="com.microsoft", **attrs), + helper.make_node("MatMulNBits", ["A", "B", "S", "Z0"], ["Y2"], domain="com.microsoft", **attrs), + helper.make_node("Add", ["Y1", "Y2"], ["Y"]), + ] + graph = helper.make_graph( + nodes, + "test", + [helper.make_tensor_value_info("A", onnx.TensorProto.FLOAT, [1, 128])], + [helper.make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1, 7])], + tensors, + ) + model = helper.make_model( + graph, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ir_version=10, + ) + source = root / "q2.onnx" + onnx.save_model(model, source) + with self.assertRaisesRegex(ValueError, "every zero point must equal 1"): + convert_model(source, root / "t8.onnx") + if __name__ == "__main__": unittest.main()