diff --git a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc index 52cff884d6ec4..5897c680cc6cd 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc @@ -144,7 +144,7 @@ class MatMulNBits final : public OpKernel { /*out*/ PrePackedWeights* prepacked_weights) override; Status UseSharedPrePackedBuffers(std::vector& prepacked_buffers, - gsl::span /*prepacked_buffer_sizes*/, + gsl::span prepacked_buffer_sizes, int input_idx, /*out*/ bool& used_shared_buffers) override; @@ -231,6 +231,21 @@ static const float* ConvertFloatZeroPointsForLutGemm( return zp_fp32_buf.data(); } +#if defined(MLAS_TARGET_ARM64) +bool RequiresDynamicZeroPointPrepackFallback( + size_t K, size_t nbits, size_t block_size, + bool has_zp_arg, bool has_zp_input, + MLAS_QNBIT_GEMM_COMPUTE_TYPE compute_type, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG& backend_kernel_selector_config) { + const auto effective_compute_type = compute_type == HQNBIT_CompInt8 ? SQNBIT_CompInt8 : compute_type; + + // KleidiAI asymmetric Q4 pack needs zero points during PrePack; dynamic zero points arrive later. + return has_zp_arg && !has_zp_input && nbits == 4 && effective_compute_type == SQNBIT_CompInt8 && + MlasQNBitGemmScalesPacked(K, nbits, block_size, effective_compute_type, + true, &backend_kernel_selector_config); +} +#endif + template Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ AllocatorPtr alloc, /*out*/ bool& is_packed, @@ -341,6 +356,13 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All return Status::OK(); } +#if defined(MLAS_TARGET_ARM64) + if (RequiresDynamicZeroPointPrepackFallback(K_, nbits_, block_size_, has_zp_arg_, has_zp_input_, + compute_type_, mlas_backend_kernel_selector_config_)) { + return Status::OK(); + } +#endif + // Create a temporary threadpool for parallel packing // This is used during model load time to speed up weight prepacking std::unique_ptr temp_threadpool; @@ -436,6 +458,15 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All scale_ptr = scales->DataRaw(); } + const void* zp_ptr = nullptr; + if (has_zp_input_) { + const Tensor* zero_points = nullptr; + ORT_IGNORE_RETURN_VALUE(OpKernel::Info().TryGetConstantInput(InputIndex::zero_points, &zero_points)); + if (zero_points != nullptr) { + zp_ptr = zero_points->DataRaw(); + } + } + packed_b_ = IAllocator::MakeUniquePtr(alloc, packed_b_size_, true); // The framework content-hashes this packed buffer to deduplicate pre-packed weights, both // within a session and across sessions (the shared container). The session-state prepack pass @@ -451,7 +482,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All std::memset(packed_b_.get(), 0, packed_b_size_); } MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, effective_compute_type, qptr, packed_b_.get(), scale_ptr, - has_zp_input_, nullptr, threadpool_ptr, &mlas_backend_kernel_selector_config_); + has_zp_input_, zp_ptr, threadpool_ptr, &mlas_backend_kernel_selector_config_); // Fold the scales and (constant) zero points into packed_b_ now, during the B PrePack, instead // of deferring them to the later scales/zero_points PrePack calls. Pre-packed weight sharing @@ -583,9 +614,8 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*out*/ All } #endif // MLAS_TARGET_ARM64 } else if (compute_type_ == HQNBIT_CompInt8) { - // For HQNBIT_CompInt8 (both 4-bit and 8-bit), scales are fp16 but packing functions expect float. - // Convert fp16 scales to float and pack using the SQNBIT_CompInt8 path. - // At compute time, we delegate to MlasQNBitGemmBatch with SQNBIT_CompInt8. + // For MLFloat16 inputs using CompInt8, scales are fp16 but MLAS SQNBIT_CompInt8 packing expects float. + // HQNBIT_CompInt8 converts fp16 data to fp32 and calls the SQNBIT_CompInt8 MLAS path. if (input_idx == InputIndex::scales && packed_b_ != nullptr) { #if defined(MLAS_TARGET_ARM64) // For 4-bit on ARM64: check if KleidiAI packs scales into B (scales already packed during B packing). @@ -725,6 +755,15 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou if (!MlasIsQNBitGemmAvailable(nbits_, block_size_, compute_type_)) { return Status::OK(); } + +#if defined(MLAS_TARGET_ARM64) + if (RequiresDynamicZeroPointPrepackFallback(K_, nbits_, block_size_, has_zp_arg_, has_zp_input_, + compute_type_, mlas_backend_kernel_selector_config_)) { + return Status::OK(); + } +#endif + + const auto effective_compute_type = compute_type_ == HQNBIT_CompInt8 ? SQNBIT_CompInt8 : compute_type_; if (input_idx == InputIndex::B) { const Tensor* scales = nullptr; OpKernel::Info().TryGetConstantInput(InputIndex::scales, &scales); @@ -733,7 +772,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou // weight sharing content-hashes the buffer right after this B PrePack returns, so for CompInt8 // everything that affects the packed bytes (the scales, and the block sum / KleidiAI BZpCorr that // depend on the zero points) must be folded in by now. - if (scales && compute_type_ == SQNBIT_CompInt8) { + if (scales && effective_compute_type == SQNBIT_CompInt8) { auto sptr = scales->Data(); auto scales_size = static_cast(scales->Shape().Size()); auto ptr = IAllocator::MakeUniquePtr(alloc, scales_size, true); @@ -742,12 +781,21 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou } packed_b_size_ = MlasQNBitGemmPackQuantBDataSize(N_, K_, nbits_, block_size_, - has_zp_input_, compute_type_, &mlas_backend_kernel_selector_config_); + has_zp_input_, effective_compute_type, + &mlas_backend_kernel_selector_config_); if (packed_b_size_ == 0) { return Status::OK(); } auto qptr = tensor.DataRaw(); packed_b_ = IAllocator::MakeUniquePtr(alloc, packed_b_size_, true); + const void* zp_ptr = nullptr; + if (has_zp_input_) { + const Tensor* zero_points = nullptr; + ORT_IGNORE_RETURN_VALUE(OpKernel::Info().TryGetConstantInput(InputIndex::zero_points, &zero_points)); + if (zero_points != nullptr) { + zp_ptr = zero_points->DataRaw(); + } + } // See the primary PrePack() above: SessionState::PrepackConstantInitializedTensors passes a // non-null prepacked_weights on both the container and the default single-session paths, so this // zero-fill runs on essentially every prepack at load (the guard only skips a caller that asks for @@ -757,8 +805,9 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou if (prepacked_weights != nullptr) { std::memset(packed_b_.get(), 0, packed_b_size_); } - MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, qptr, packed_b_.get(), - scales_fp32_.get(), has_zp_input_, nullptr, nullptr, &mlas_backend_kernel_selector_config_); + MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, effective_compute_type, qptr, packed_b_.get(), + scales_fp32_.get(), has_zp_input_, zp_ptr, nullptr, + &mlas_backend_kernel_selector_config_); // Fold the scales and (constant) zero points into packed_b_ now (see the primary PrePack above): // the CompInt8 block sum and the KleidiAI BZpCorr depend on the zero points, so they must be @@ -767,14 +816,14 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou // first's buffer. The B pack above only partially populates the buffer, so issue one more pack // call with QuantBData == nullptr to finalize it. This is byte-identical to the staged // scales + zero_points packing it replaces. - bool finalize_scale_zp_into_packed_b = compute_type_ == SQNBIT_CompInt8 && scales_fp32_ != nullptr; + bool finalize_scale_zp_into_packed_b = effective_compute_type == SQNBIT_CompInt8 && scales_fp32_ != nullptr; #if !defined(MLAS_TARGET_AMD64_IX86) // On ARM64 the scales/zero points are folded into B only for 8-bit, or for 4-bit when MLAS bakes // them in (KleidiAI). For 4-bit non-KleidiAI they are applied at compute time and must not be // passed to the packing routine, which would dereference the null QuantBData buffer. finalize_scale_zp_into_packed_b = finalize_scale_zp_into_packed_b && - (nbits_ == 8 || MlasQNBitGemmScalesPacked(K_, nbits_, block_size_, compute_type_, + (nbits_ == 8 || MlasQNBitGemmScalesPacked(K_, nbits_, block_size_, effective_compute_type, has_zp_input_, &mlas_backend_kernel_selector_config_)); #endif if (finalize_scale_zp_into_packed_b) { @@ -786,7 +835,7 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou zp_ptr = zp_tensor->Data(); } } - MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, nullptr /*QuantBData*/, + MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, effective_compute_type, nullptr /*QuantBData*/, packed_b_.get(), scales_fp32_.get(), has_zp_input_, zp_ptr, nullptr, &mlas_backend_kernel_selector_config_); packed_b_finalized_ = true; @@ -826,28 +875,47 @@ Status MatMulNBits::PrePack(const Tensor& tensor, int input_idx, /*ou template Status MatMulNBits::UseSharedPrePackedBuffers(std::vector& prepacked_buffers, - gsl::span /*prepacked_buffer_sizes*/, + gsl::span prepacked_buffer_sizes, int input_idx, /*out*/ bool& used_shared_buffers) { used_shared_buffers = false; if (input_idx == InputIndex::B && !prepacked_buffers.empty()) { - // The buffer handed back is fully finalized: the producing session folded the scales and zero - // points (block sums / KleidiAI BZpCorr) into it during its PrePack(B), which is also when this - // kernel set packed_b_finalized_ on its own (identical) B PrePack. The later scale/zero-point - // PrePack calls already skip the staged packing whenever packed_b_finalized_ is set, so simply - // adopt the shared buffer here - no extra bookkeeping is needed to avoid re-folding into it. +#if defined(MLAS_TARGET_ARM64) + ORT_RETURN_IF(RequiresDynamicZeroPointPrepackFallback(K_, nbits_, block_size_, has_zp_arg_, has_zp_input_, + compute_type_, mlas_backend_kernel_selector_config_), + "MatMulNBits cannot use shared prepacked B for KleidiAI Q4 with runtime zero_points. ", + "PrePack should have declined prepacking for this node."); +#endif + + ORT_RETURN_IF(prepacked_buffer_sizes.empty(), + "Missing MatMulNBits prepacked B buffer size metadata."); + packed_b_ = std::move(prepacked_buffers[0]); + packed_b_size_ = prepacked_buffer_sizes[0]; used_shared_buffers = true; if (prefer_lut_gemm_) { MlasInitLutGemmKernelConfig(N_, K_, nbits_, block_size_, has_zp_input_); - packed_b_size_ = MlasLutGemmPackedSize(N_, K_, nbits_, block_size_, has_zp_input_); + } else { + const Tensor* scales = nullptr; + const auto effective_compute_type = compute_type_ == HQNBIT_CompInt8 ? SQNBIT_CompInt8 : compute_type_; + bool shared_packed_b_finalized = effective_compute_type == SQNBIT_CompInt8 && + OpKernel::Info().TryGetConstantInput(InputIndex::scales, &scales); +#if defined(MLAS_TARGET_ARM64) + // 4-bit scales are folded into B only when MLAS reports backend support. + shared_packed_b_finalized = + shared_packed_b_finalized && + (nbits_ == 8 || MlasQNBitGemmScalesPacked(K_, nbits_, block_size_, effective_compute_type, + has_zp_input_, &mlas_backend_kernel_selector_config_)); +#endif + if (shared_packed_b_finalized) { + scales_are_packed_ = true; + packed_b_finalized_ = true; + } } } - // Only the quantized weight B yields a separately cached pre-packed buffer. The scales (and zero - // points) are folded into packed_b_ during the B PrePack and reported with is_packed = false, so - // the framework never asks this kernel to adopt a shared buffer for them. + // Scales and zero points are folded into packed_b_ and are not shared as separate buffers. return Status::OK(); } @@ -901,13 +969,10 @@ Status MatMulNBits::ComputeBPacked(const Tensor* a, const size_t K = static_cast(helper.K()); const size_t lda = helper.Lda(false); - // For HQNBIT_CompInt8 with fp16 inputs: delegate to fp32 MLAS path (SQNBIT_CompInt8). - // The HQ CompInt8 kernels are just wrappers that convert fp16->fp32 per-tile and call the same - // SQ fp32 kernels. By doing bulk conversion at the operator level we eliminate per-tile overhead - // and automatically get KleidiAI support for 4-bit (since SQ4BitGemm_CompInt8 checks KleidiAI). - // This matches the approach used by x64 and Apple ARM64 (non-fp16-intrinsics fallback). + // For fp16 MatMulNBits with CompInt8, run through the fp32 SQNBIT_CompInt8 MLAS path. + // Bulk conversion avoids per-tile fp16 wrappers and lets the same path use KleidiAI when selected. if constexpr (std::is_same_v) { - if (compute_type_ == HQNBIT_CompInt8) { + if (compute_type_ == SQNBIT_CompInt8 || compute_type_ == HQNBIT_CompInt8) { const auto* a_data_fp16 = a->Data(); const auto* bias_data_fp16 = bias == nullptr ? nullptr : bias->Data(); @@ -1024,7 +1089,7 @@ Status MatMulNBits::ComputeBPacked(const Tensor* a, concurrency::ThreadPool* thread_pool, const MatMulComputeHelper& helper) const { const auto* a_data = a->Data(); - const auto* scales_data = scales->Data(); + const auto* scales_data = scales == nullptr ? nullptr : scales->Data(); const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw(); const auto* bias_data = bias == nullptr ? nullptr : bias->Data(); auto* y_data = y->MutableData(); @@ -1034,10 +1099,12 @@ Status MatMulNBits::ComputeBPacked(const Tensor* a, const size_t N = static_cast(helper.N()); const size_t K = static_cast(helper.K()); const size_t lda = helper.Lda(false); + const auto effective_compute_type = compute_type_ == HQNBIT_CompInt8 ? SQNBIT_CompInt8 : compute_type_; IAllocatorUniquePtr workspace{}; const size_t workspace_size = MlasQNBitGemmBatchWorkspaceSize( - M, N, K, batch_count, nbits_, block_size_, zero_points, compute_type_, &mlas_backend_kernel_selector_config_); + M, N, K, batch_count, nbits_, block_size_, zero_points, effective_compute_type, + &mlas_backend_kernel_selector_config_); if (workspace_size > 0) { // Use reserve since no caching is needed workspace = IAllocator::MakeUniquePtr(allocator, workspace_size, true); @@ -1048,33 +1115,40 @@ Status MatMulNBits::ComputeBPacked(const Tensor* a, MlasConvertHalfToFloatBuffer(a_data, tmp_a_data_ptr.get(), a_size); float* scales_ptr = nullptr; - if (!scales_fp32_) { - auto scales_temp = IAllocator::MakeUniquePtr(allocator, static_cast(scales->Shape().Size()), true); - MlasConvertHalfToFloatBuffer(scales_data, scales_temp.get(), static_cast(scales->Shape().Size())); - scales_ptr = scales_temp.get(); - } else { - scales_ptr = scales_fp32_.get(); + IAllocatorUniquePtr scales_temp; + if (!scales_are_packed_) { + if (scales_fp32_) { + scales_ptr = scales_fp32_.get(); + } else { + ORT_ENFORCE(scales != nullptr, "scales must be provided when not packed and not pre-converted"); + const auto scales_size = static_cast(scales->Shape().Size()); + scales_temp = IAllocator::MakeUniquePtr(allocator, scales_size, true); + MlasConvertHalfToFloatBuffer(scales_data, scales_temp.get(), scales_size); + scales_ptr = scales_temp.get(); + } } float* bias_ptr = nullptr; + IAllocatorUniquePtr bias_temp; if (bias_data) { - if (!bias_fp32_) { - auto bias_temp = IAllocator::MakeUniquePtr(allocator, static_cast(bias->Shape().Size()), true); - MlasConvertHalfToFloatBuffer(bias_data, bias_temp.get(), static_cast(bias->Shape().Size())); - bias_ptr = bias_temp.get(); - } else { + if (bias_fp32_) { bias_ptr = bias_fp32_.get(); + } else { + const auto bias_size = static_cast(bias->Shape().Size()); + bias_temp = IAllocator::MakeUniquePtr(allocator, bias_size, true); + MlasConvertHalfToFloatBuffer(bias_data, bias_temp.get(), bias_size); + bias_ptr = bias_temp.get(); } } - size_t c_size = static_cast(y->Shape().Size()); + const size_t c_size = static_cast(y->Shape().Size()); std::vector c_v(c_size); InlinedVector> data(batch_count); for (size_t i = 0; i < batch_count; ++i) { data[i].A = tmp_a_data_ptr.get() + helper.LeftOffsets()[i]; data[i].lda = lda; - if (compute_type_ == SQNBIT_CompInt8) { + if (effective_compute_type == SQNBIT_CompInt8) { data[i].QuantBDataWorkspace = packed_b_.get(); } data[i].PackedQuantBData = static_cast(packed_b_.get()); @@ -1084,7 +1158,7 @@ Status MatMulNBits::ComputeBPacked(const Tensor* a, data[i].C = c_v.data() + helper.OutputOffsets()[i]; data[i].ldc = N; } - MlasQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, compute_type_, data.data(), workspace.get(), + MlasQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, effective_compute_type, data.data(), workspace.get(), thread_pool, &mlas_backend_kernel_selector_config_); MlasConvertFloatToHalfBuffer(c_v.data(), y_data, c_size); return Status::OK(); diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp index 6ee80594c6b49..2e6b177499dde 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp @@ -17,6 +17,12 @@ #include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod.h" #include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp4x4_qsi4c32p4x4_16x4_neon_dotprod.h" #include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_16x4x32_neon_i8mm.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4x4_1x4_neon_dotprod.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p1x8_qai4c32p4x8_1x4_neon_dotprod.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p4x4_qai4c32p4x4_8x4_neon_dotprod.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p4x8_qai4c32p4x8_8x4_neon_i8mm.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p1vlx4_qai4c32p4vlx4_1vlx4vl_sme2_mopa.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4vlx4_1x4vl_sme2_dot.h" // GEMV #include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p8x1biasf32_6x8x4_neon_mla.h" @@ -97,6 +103,23 @@ kai_run_##STEM} \ } +#define KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(STEM, RHS_LAYOUT) \ + { \ + "kai_run_" #STEM, \ + {kai_get_m_step_##STEM, \ + kai_get_n_step_##STEM, \ + kai_get_mr_##STEM, \ + kai_get_nr_##STEM, \ + kai_get_kr_##STEM, \ + kai_get_sr_##STEM, \ + kai_get_lhs_packed_offset_##STEM, \ + kai_get_rhs_packed_offset_##STEM, \ + kai_get_dst_offset_##STEM, \ + kai_get_dst_size_##STEM, \ + kai_run_##STEM}, \ + RHS_LAYOUT \ + } + // 7-slot packed `run_imatmul` interface shape. // // Applies to KleidiAI ukernel interface headers/structs such as: @@ -208,16 +231,44 @@ const KaiQnbitGemmKernel kai_matmul_clamp_f32_qai8dxp1x4_qsi4c32p4x4_1x4_neon_dotprod = - KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp1x4_qsi4c32p4x4_1x4_neon_dotprod); + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp1x4_qsi4c32p4x4_1x4_neon_dotprod, + KaiQ4RhsPackLayout::SymmetricNxK); const KaiQnbitGemmKernel kai_matmul_clamp_f32_qai8dxp4x4_qsi4c32p4x4_16x4_neon_dotprod = - KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp4x4_qsi4c32p4x4_16x4_neon_dotprod); + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp4x4_qsi4c32p4x4_16x4_neon_dotprod, + KaiQ4RhsPackLayout::SymmetricNxK); const KaiQnbitGemmKernel kai_matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod = - KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod); + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp1x8_qsi4c32p4x8_1x4x32_neon_dotprod, + KaiQ4RhsPackLayout::SymmetricNxK); const KaiQnbitGemmKernel kai_matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_16x4x32_neon_i8mm = - KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_16x4x32_neon_i8mm); + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qai8dxp4x8_qsi4c32p4x8_16x4x32_neon_i8mm, + KaiQ4RhsPackLayout::SymmetricNxK); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4x4_1x4_neon_dotprod = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p1x4_qai4c32p4x4_1x4_neon_dotprod, + KaiQ4RhsPackLayout::AsymmetricNxK); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p4x4_qai4c32p4x4_8x4_neon_dotprod = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p4x4_qai4c32p4x4_8x4_neon_dotprod, + KaiQ4RhsPackLayout::AsymmetricNxK); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p1x8_qai4c32p4x8_1x4_neon_dotprod = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p1x8_qai4c32p4x8_1x4_neon_dotprod, + KaiQ4RhsPackLayout::AsymmetricNxK); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p4x8_qai4c32p4x8_8x4_neon_i8mm = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p4x8_qai4c32p4x8_8x4_neon_i8mm, + KaiQ4RhsPackLayout::AsymmetricNxK); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p1vlx4_qai4c32p4vlx4_1vlx4vl_sme2_mopa = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p1vlx4_qai4c32p4vlx4_1vlx4vl_sme2_mopa, + KaiQ4RhsPackLayout::AsymmetricNxKInterleavedNrx4); + +const KaiQnbitAsymGemmKernel kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4vlx4_1x4vl_sme2_dot = + KAI_WRAP_Q4_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_qsi8d32p1x4_qai4c32p4vlx4_1x4vl_sme2_dot, + KaiQ4RhsPackLayout::AsymmetricNxKInterleavedNrx4); const KaiF32SgemmKernel sgemm_gemm_sme = KAI_WRAP_UKERNEL_RUN_MATMUL_11(matmul_clamp_f32_f32p2vlx1_f32p2vlx1b_2vlx2vl_sme_mopa); @@ -306,6 +357,26 @@ const KaiQnbitGemmKernel& GetKleidiAIGemvUKernel() { } } +const KaiQnbitAsymGemmKernel& GetKleidiAIQai4GemmUKernel() { + if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + return kai_matmul_clamp_f32_qsi8d32p1vlx4_qai4c32p4vlx4_1vlx4vl_sme2_mopa; + } else if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeon_I8MM()) { + return kai_matmul_clamp_f32_qsi8d32p4x8_qai4c32p4x8_8x4_neon_i8mm; + } else { + return kai_matmul_clamp_f32_qsi8d32p4x4_qai4c32p4x4_8x4_neon_dotprod; + } +} + +const KaiQnbitAsymGemmKernel& GetKleidiAIQai4GemvUKernel() { + if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { + return kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4vlx4_1x4vl_sme2_dot; + } else if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeon_I8MM()) { + return kai_matmul_clamp_f32_qsi8d32p1x8_qai4c32p4x8_1x4_neon_dotprod; + } else { + return kai_matmul_clamp_f32_qsi8d32p1x4_qai4c32p4x4_1x4_neon_dotprod; + } +} + const KaiF32SgemmKernel& GetKleidiAISGemmUKernel() { if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2()) { return sgemm_gemm_sme2; diff --git a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h index 155ecf1762b3b..1c1d704ab1b20 100644 --- a/onnxruntime/core/mlas/lib/kai_ukernel_interface.h +++ b/onnxruntime/core/mlas/lib/kai_ukernel_interface.h @@ -8,6 +8,8 @@ #include "kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi4c32p/kai_matmul_clamp_f32_qai8dxp_qsi4c32p_interface.h" +#include "kai/ukernels/matmul/matmul_clamp_f32_qsi8d32p_qai4c32p/kai_matmul_clamp_f32_qsi8d32p_qai4c32p_interface.h" + #include "kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p_f32p_interface.h" #include "kai/ukernels/matmul/matmul_clamp_f32_f32_f32p/kai_matmul_clamp_f32_f32_f32p_interface.h" @@ -26,6 +28,19 @@ struct KaiMatmulKernel { UkernelFn ukernel; }; +enum class KaiQ4RhsPackLayout { + SymmetricNxK, + AsymmetricNxK, + AsymmetricNxKInterleavedNrx4, +}; + +template +struct KaiQ4MatmulKernel { + const char* name; + UkernelFn ukernel; + KaiQ4RhsPackLayout rhs_layout; +}; + // Wrapper for FP32 GEMM kernels where both LHS and RHS are pre-packed (common SGEMM path). using KaiF32SgemmKernel = KaiMatmulKernel; @@ -33,7 +48,10 @@ using KaiF32SgemmKernel = KaiMatmulKernel; // Wrapper for Qnbit GEMM kernels producing FP32 output. -using KaiQnbitGemmKernel = KaiMatmulKernel; +using KaiQnbitGemmKernel = KaiQ4MatmulKernel; + +// Wrapper for Qnbit Asymmetric-quantized GEMM kernels producing FP32 output. +using KaiQnbitAsymGemmKernel = KaiQ4MatmulKernel; // Wrapper for dynamic-quantized GEMM kernels producing FP32 output. using KaiDynamicQGemmKernel = KaiMatmulKernel; @@ -49,6 +67,12 @@ const KaiQnbitGemmKernel& GetKleidiAIGemmUKernel(); // Returns the selected Qnbit kernel used for GEMV-style workloads based on runtime CPU capabilities. const KaiQnbitGemmKernel& GetKleidiAIGemvUKernel(); +// Returns the selected Qnbit Asymmetric-quantized GEMM ukernel. +const KaiQnbitAsymGemmKernel& GetKleidiAIQai4GemmUKernel(); + +// Returns the selected Qnbit Asymmetric-quantized kernel used for GEMV-style workloads. +const KaiQnbitAsymGemmKernel& GetKleidiAIQai4GemvUKernel(); + // Returns the selected dynamic-quantized GEMM ukernel based on runtime CPU capabilities and optional vendor selection. const KaiDynamicQGemmKernel& GetKleidiAIQGemmUKernel(); @@ -62,4 +86,4 @@ const KaiF32SgemvKernel& GetKleidiAISGemvUKernel(); const KaiF32IMatmulKernel& GetKleidiAIF32IMatmulUKernel(); // Returns the selected BF16 SBGEMM ukernel used by the KleidiAI based on runtime CPU capabilities. -const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); \ No newline at end of file +const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel(); diff --git a/onnxruntime/core/mlas/lib/qnbitgemm.cpp b/onnxruntime/core/mlas/lib/qnbitgemm.cpp index 441cd7f0d3c70..42ea3369afe99 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm.cpp +++ b/onnxruntime/core/mlas/lib/qnbitgemm.cpp @@ -82,6 +82,36 @@ GetQNBitGemmVariant( return SQNBitGemmVariantInvalid; } +bool RequiresPackedZpCorrection( + size_t K, + size_t BlkLen, + const void* QuantBZeroPoint, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) +{ + const bool has_zp = QuantBZeroPoint != nullptr; + if (!has_zp) { + return false; + } + + const auto* dispatch = GetMlasPlatform().QNBitGemmDispatch; + const auto fn = dispatch == nullptr ? nullptr : dispatch->NeedsPackedZpCorrection_CompInt8; + return fn != nullptr && fn(K, BlkLen, has_zp, BackendKernelSelectorConfig); +} + +size_t GetPackedQ4BitGemmNAlignment( + size_t K, + size_t BlkLen, + const void* QuantBZeroPoint, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) +{ + const auto* dispatch = GetMlasPlatform().QNBitGemmDispatch; + const auto fn = dispatch == nullptr ? nullptr : dispatch->PackedQ4BitGemmNAlignment_CompInt8; + if (fn == nullptr) { + return MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + } + return fn(K, BlkLen, QuantBZeroPoint != nullptr, BackendKernelSelectorConfig); +} + } // namespace bool MLASCALL @@ -739,17 +769,18 @@ SQ4BitGemm_CompInt8( { const auto UsePacked = GetMlasPlatform().QNBitGemmDispatch->UsePacked_CompInt8; const auto SQ4BitGemm = GetMlasPlatform().QNBitGemmDispatch->SQ4BitGemmKernel_Packed_CompInt8; - if (UsePacked && SQ4BitGemm && UsePacked(K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig)) { + const bool HasQuantBZeroPoint = DataParams->QuantBZeroPoint != nullptr; + if (UsePacked && SQ4BitGemm && UsePacked(K, BlkLen, HasQuantBZeroPoint, BackendKernelSelectorConfig)) { const std::byte* QuantA = static_cast(PerGemmWorkspace); SQ4BitGemm(BlkLen, QuantA, DataParams->PackedQuantBData, DataParams->C, RangeStartM, RangeCountM, RangeStartN, RangeCountN, K, + HasQuantBZeroPoint, BackendKernelSelectorConfig, DataParams->ldc, DataParams->Bias); - // Apply zero-point correction for asymmetric quantization (KleidiAI path only). - // BZpCorr and AFloatBlkSum are only set when KleidiAI is active with asymmetric - // quantization (has zero points). On all other paths they remain nullptr. + // Apply correction only for packed backends that cannot consume RHS zero points directly. // C += AFloatBlkSum * BZpCorr^T (for this tile's M/N ranges) - if (DataParams->BZpCorr != nullptr && DataParams->AFloatBlkSum != nullptr) { + if (RequiresPackedZpCorrection(K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig) && + DataParams->BZpCorr != nullptr && DataParams->AFloatBlkSum != nullptr) { const size_t BlockCountK = MlasDivRoundup(K, BlkLen); const size_t ldc = DataParams->ldc; const float* ABlkSum = DataParams->AFloatBlkSum + RangeStartM * BlockCountK; @@ -1128,7 +1159,8 @@ InitializeWorkspace_CompInt8( const size_t QuantAStride = BlockCountK * Q8BlkSize(BlkLen); // TODO: try parallel on BatchN * M threads because BatchN is usually 1. - if (BlkBitWidth == 4 && UsePacked && QuantizeA_Packed && UsePacked(K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig)) { + const bool has_zp_input = DataParams->QuantBZeroPoint != nullptr; + if (BlkBitWidth == 4 && UsePacked && QuantizeA_Packed && UsePacked(K, BlkLen, has_zp_input, BackendKernelSelectorConfig)) { // Compute KleidiAI packed A size (same as workspace size without zero points) const size_t kleidiAIPackedASize = GetMlasPlatform().QNBitGemmDispatch->QNBitGemmPerGemmWorkspaceSize ? GetMlasPlatform().QNBitGemmDispatch->QNBitGemmPerGemmWorkspaceSize( @@ -1140,11 +1172,12 @@ InitializeWorkspace_CompInt8( const float* ARowPtr = data.A; std::byte* QuantARowPtr = static_cast(Workspace) + gemm_idx * PerGemmWorkspaceStride; - QuantizeA_Packed(BlkLen, ARowPtr, M, K, QuantARowPtr, BackendKernelSelectorConfig); + const bool HasZp = data.QuantBZeroPoint != nullptr; + QuantizeA_Packed(BlkLen, ARowPtr, M, K, HasZp, QuantARowPtr, BackendKernelSelectorConfig); - // For asymmetric KleidiAI path, also compute float-domain A block sums - // for zero-point correction. AFloatBlkSum is stored after KleidiAI packed A. - if (data.QuantBZeroPoint != nullptr && ComputeAFloatBlkSumFn != nullptr && kleidiAIPackedASize > 0) { + // Some packed backends need A block sums for RHS zero-point correction. + if (RequiresPackedZpCorrection(K, BlkLen, data.QuantBZeroPoint, BackendKernelSelectorConfig) && + ComputeAFloatBlkSumFn != nullptr && kleidiAIPackedASize > 0) { // Align offset so AFloatBlkSum starts at a float-aligned address constexpr size_t FloatAlignment = alignof(float); const size_t alignedAOffset = (kleidiAIPackedASize + FloatAlignment - 1) & ~(FloatAlignment - 1); @@ -1357,7 +1390,7 @@ MlasQNBitGemmBatch( ); } - const bool has_zp_input = DataParams->QuantBZeroPoint; + const bool has_zp_input = DataParams->QuantBZeroPoint != nullptr; const size_t PerGemmWorkspaceStride = QNBitGemmPerGemmWorkspaceStride(M, N, K, BlkBitWidth, BlkLen, has_zp_input, ComputeType, BackendKernelSelectorConfig); @@ -1372,12 +1405,10 @@ MlasQNBitGemmBatch( const size_t BlockCountK = MlasDivRoundup(K, BlkLen); - // For KleidiAI asymmetric path: set up BZpCorr and AFloatBlkSum pointers. - // BZpCorr is stored after KleidiAI packed B data. - // AFloatBlkSum is stored after KleidiAI packed A data in each per-GEMM workspace. + // Set up correction buffers for packed backends that cannot consume RHS zero points directly. const auto UsePacked = GetMlasPlatform().QNBitGemmDispatch->UsePacked_CompInt8; - if (Variant == SQ4BitGemmVariant_CompInt8 && has_zp_input && UsePacked && - UsePacked(K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig)) { + if (Variant == SQ4BitGemmVariant_CompInt8 && RequiresPackedZpCorrection(K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig) && + UsePacked && UsePacked(K, BlkLen, has_zp_input, BackendKernelSelectorConfig)) { // Compute KleidiAI packed B size (without zero point correction space) const size_t kleidiAIPackedBSize = GetMlasPlatform().QNBitGemmDispatch->Q4BitGemmPackQuantBDataSize ? GetMlasPlatform().QNBitGemmDispatch->Q4BitGemmPackQuantBDataSize( @@ -1468,6 +1499,13 @@ MlasQNBitGemmBatch( constexpr size_t StrideM = 128; + size_t StrideNThreadAlign = MLAS_QGEMM_STRIDEN_THREAD_ALIGN; + if (Variant == SQ4BitGemmVariant_CompInt8) { + const size_t PackedNAlignment = GetPackedQ4BitGemmNAlignment( + K, BlkLen, DataParams->QuantBZeroPoint, BackendKernelSelectorConfig); + StrideNThreadAlign = std::max(StrideNThreadAlign, PackedNAlignment); + } + size_t nc = N; if (ThreadsPerGemm > 1) { // more than one thread per GEMM @@ -1476,8 +1514,7 @@ MlasQNBitGemmBatch( const size_t max_nc = MlasDivRoundup(N * BlockedM, ThreadsPerGemm); if (max_nc < nc) { nc = std::min( - nc, MlasDivRoundup(max_nc, MLAS_QGEMM_STRIDEN_THREAD_ALIGN) * - MLAS_QGEMM_STRIDEN_THREAD_ALIGN + nc, MlasDivRoundup(max_nc, StrideNThreadAlign) * StrideNThreadAlign ); } } diff --git a/onnxruntime/core/mlas/lib/qnbitgemm.h b/onnxruntime/core/mlas/lib/qnbitgemm.h index 441da585c6611..c23941382f0bd 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm.h +++ b/onnxruntime/core/mlas/lib/qnbitgemm.h @@ -377,7 +377,10 @@ struct MLAS_QNBIT_GEMM_DISPATCH { * @param RangeStartN Start of N range. * @param RangeCountN Number of columns of B and C. * @param CountK Number of columns of A and rows of B. + * @param HasQuantBZeroPoint Whether RHS zero points were present during B packing. + * @param BackendKernelSelectorConfig Optional backend kernel selector configuration. * @param ldc Number of elements between adjacent rows of C. + * @param Bias Bias vector of length N. Optional. */ typedef void(SQ4BitGemmKernel_Packed_CompInt8_Fn)( size_t BlkLen, @@ -389,6 +392,8 @@ struct MLAS_QNBIT_GEMM_DISPATCH { const size_t RangeStartN, const size_t RangeCountN, size_t CountK, + bool HasQuantBZeroPoint, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig, size_t ldc, const float* Bias ); @@ -582,16 +587,41 @@ struct MLAS_QNBIT_GEMM_DISPATCH { UsePacked_CompInt8_Fn* UsePacked_CompInt8 = nullptr; + /** + * @brief Required N-start alignment for SQ4BitGemmKernel_Packed_CompInt8 ranges. + */ + typedef size_t(PackedQ4BitGemmNAlignment_CompInt8_Fn)( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ); + + PackedQ4BitGemmNAlignment_CompInt8_Fn* PackedQ4BitGemmNAlignment_CompInt8 = nullptr; + + /** + * @brief Whether to correct for asymmetric zero-point in RHS packed CompInt8 workloads + */ + typedef bool(NeedsPackedZpCorrection_CompInt8_Fn)( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig + ); + + NeedsPackedZpCorrection_CompInt8_Fn* NeedsPackedZpCorrection_CompInt8 = nullptr; + /** * @brief Block quantize values from matrix A from floats to quantized 8-bit integers. * Used in conjunction with SQ4BitGemmKernel_Packed_CompInt8. * - * @param BlkLen Number of values in a block. - * @param A Supplies the A matrix. - * @param CountM Number of rows of A. - * @param CountK Number of columns of A. - * @param[out] QuantA Supplies the output quantized A matrix. - * Binary data containing block quantized int8 data and scale values. + * @param BlkLen Number of values in a block. + * @param A Supplies the A matrix. + * @param CountM Number of rows of A. + * @param CountK Number of columns of A. + * @param RHSHasZeroPoint Whether packing happens for asymmetric or symmetric RHS + * @param[out] QuantA Supplies the output quantized A matrix. + * Binary data containing block quantized int8 data and scale values. * @param BackendKernelSelectorConfig Optional configuration for selecting backend kernels. */ typedef void(QuantizeA_Packed_CompInt8_Fn)( @@ -599,6 +629,7 @@ struct MLAS_QNBIT_GEMM_DISPATCH { const float* A, size_t CountM, size_t CountK, + bool RHSHasZeroPoint, std::byte* QuantA, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ); diff --git a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp index 4ff9b6f1a4d6f..e5e6c887b2b7c 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp +++ b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp @@ -38,6 +38,9 @@ Module Name: #include "kai/kai_common.h" #include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0.h" #include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.h" +#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qai4c32p_qau4c32s0s1_f32_f32_f32_neon.h" +#include "kai/ukernels/matmul/pack/kai_rhs_pack_nxk_qai4c32ps1s0nrx4_qau4c32s0s1_f32_f32_f32_neon.h" +#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32pscalef32_f32_neon.h" #include "kai_ukernel_interface.h" #endif @@ -70,29 +73,51 @@ QNBitGemmPackQuantBDataSize( #endif #ifdef USE_KLEIDIAI - if (ComputeType == SQNBIT_CompInt8 && UseKleidiAI(K, BlkLen, BackendKernelSelectorConfig)) { - const auto& k = GetKleidiAIGemmUKernel(); - const auto& ukernel = k.ukernel; - const size_t nr = ukernel.get_nr(); - const size_t kr = ukernel.get_kr(); - const size_t sr = ukernel.get_sr(); - size_t packed_size = kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0(N, K, nr, kr, sr, BlkLen, kai_dt_bf16); - if (HasZeroPoint) { - // Align so that BZpCorrection starts at a float-aligned offset - constexpr size_t FloatAlignment = alignof(float); - packed_size = (packed_size + FloatAlignment - 1) & ~(FloatAlignment - 1); - // Additional space for BZpCorrection: N * BlockCountK floats - const size_t BlockCountK = MlasDivRoundup(K, BlkLen); - packed_size += N * BlockCountK * sizeof(float); + if (ComputeType == SQNBIT_CompInt8) { + switch (SelectKleidiAIQ4Backend(K, BlkLen, HasZeroPoint, BackendKernelSelectorConfig)) { + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: { + const auto& k = GetKleidiAIGemmUKernel(); + const auto& ukernel = k.ukernel; + const size_t nr = ukernel.get_nr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + size_t packed_size = kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0(N, K, nr, kr, sr, BlkLen, kai_dt_bf16); + if (HasZeroPoint) { + // Align so that BZpCorrection starts at a float-aligned offset + constexpr size_t FloatAlignment = alignof(float); + packed_size = (packed_size + FloatAlignment - 1) & ~(FloatAlignment - 1); + // Additional space for BZpCorrection: N * BlockCountK floats + const size_t BlockCountK = MlasDivRoundup(K, BlkLen); + packed_size += N * BlockCountK * sizeof(float); + } + return packed_size; + } + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: { + // Packed B is shared by GEMM and GEMV, so both kernels must use this RHS layout. + const auto& k = GetKleidiAIQai4GemmUKernel(); + const auto& ukernel = k.ukernel; + const size_t nr = ukernel.get_nr(); + const size_t kr = ukernel.get_kr(); + // Get asymmetric packed size based on which RHS layout is used by the kernel + switch (k.rhs_layout) { + case KaiQ4RhsPackLayout::AsymmetricNxK: + return kai_get_rhs_packed_size_rhs_pack_nxk_qai4c32p_qau4c32s0s1_f32_f32_f32_neon(N, K, nr, kr, BlkLen); + case KaiQ4RhsPackLayout::AsymmetricNxKInterleavedNrx4: + return kai_get_rhs_packed_size_rhs_pack_nxk_qai4c32ps1s0nrx4_qau4c32s0s1_f32_f32_f32_neon(N, K, nr, kr, BlkLen); + default: + // Invalid layout for the asymmetric Q4 backend. + assert(false); + return 0; + } + } + case KleidiAIQ4Backend::None: + break; } - return packed_size; - } else -#endif - { - const size_t BlockCountK = MlasDivRoundup(K, BlkLen); - const size_t PackedQuantBDataSize = N * BlockCountK * MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); - return PackedQuantBDataSize; } +#endif + const size_t BlockCountK = MlasDivRoundup(K, BlkLen); + const size_t PackedQuantBDataSize = N * BlockCountK * MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); + return PackedQuantBDataSize; } else { const size_t BlockCountK = MlasDivRoundup(K, BlkLen); size_t PackedQuantBDataSize = N * BlockCountK * MlasQNBitBlkDataSizeInBytes(BlkBitWidth, BlkLen); @@ -214,61 +239,137 @@ SQ4BitGemmPackQuantBDataAndBlkSum( assert(BlkLen >= 16 && BlkLen % 16 == 0); #ifdef USE_KLEIDIAI - if (UseKleidiAI(K, BlkLen, BackendKernelSelectorConfig)) { - const auto& k = GetKleidiAIGemmUKernel(); - const auto& ukernel = k.ukernel; + KleidiAIQ4Backend backend = SelectKleidiAIQ4Backend(K, BlkLen, HasZeroPoint, BackendKernelSelectorConfig); + if (backend != KleidiAIQ4Backend::None) { std::byte* PackedQuantBDataBegin = PackedQuantB.PackedQuantBData; - - const size_t nr = ukernel.get_nr(); - const size_t kr = ukernel.get_kr(); - const size_t sr = ukernel.get_sr(); const size_t BlockCountK = MlasDivRoundup(K, BlkLen); // Pack B data with KleidiAI (only when B data is provided) if (QuantBDataBegin != nullptr) { - assert(QuantBScaleBegin != nullptr); - kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0_params params; - params.lhs_zero_point = 1; - params.rhs_zero_point = 8; - params.scale_dt = kai_dt_bf16; - - const size_t scales_len = N * BlockCountK; - std::vector scales(scales_len); - for (size_t i = 0; i < scales_len; i++) { - uint32_t bits; - static_assert(sizeof(bits) == sizeof(QuantBScaleBegin[i]), "Unexpected float size"); - std::memcpy(&bits, &QuantBScaleBegin[i], sizeof(bits)); - scales[i] = static_cast(bits >> 16); - } - - kai_run_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0(1, N, K, nr, kr, sr, BlkLen, - reinterpret_cast(QuantBDataBegin), BlockCountK * BlkLen / 2, - nullptr, scales.data(), BlockCountK * sizeof(uint16_t), - PackedQuantBDataBegin, 0, ¶ms); - } - - // Compute BZpCorrection when both scales and zero points are available. - // BZpCorr[n * BlockCountK + blk] = scale_b * (8 - zp_b) - // Note: We intentionally use fp32 scales here (not bf16-truncated) for higher precision - // in the correction term. KleidiAI internally truncates scales to bf16 for the main GEMM, - // but the correction benefits from full fp32 precision to better approximate the true result. - // This may be called separately from B packing when zero points arrive later. - if (HasZeroPoint && QuantBZPBegin != nullptr && QuantBScaleBegin != nullptr) { - const size_t kleidiai_packed_size = kai_get_rhs_packed_size_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0( - N, K, nr, kr, sr, BlkLen, kai_dt_bf16); - // Align offset so BZpCorr starts at a float-aligned address - constexpr size_t FloatAlignment = alignof(float); - const size_t bzpcorr_offset = (kleidiai_packed_size + FloatAlignment - 1) & ~(FloatAlignment - 1); - float* BZpCorr = reinterpret_cast(PackedQuantBDataBegin + bzpcorr_offset); - - for (size_t n = 0; n < N; ++n) { - for (size_t blk = 0; blk < BlockCountK; ++blk) { - const size_t idx = n * BlockCountK + blk; - const size_t zp_byte_idx = blk / 2; - const uint8_t zp_byte = static_cast(QuantBZPBegin[n * MlasDivRoundup(BlockCountK, 2) + zp_byte_idx]); - const uint8_t zp = (blk & 1) == 0 ? (zp_byte & 0x0F) : (zp_byte >> 4); - BZpCorr[idx] = QuantBScaleBegin[idx] * (8.0f - static_cast(zp)); + switch (backend) { + // Symmetric Q4 block-quantized RHS path + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: { + const auto& k = GetKleidiAIGemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t nr = ukernel.get_nr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + + assert(QuantBScaleBegin != nullptr); + kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0_params params{}; + params.lhs_zero_point = 1; + params.rhs_zero_point = 8; + params.scale_dt = kai_dt_bf16; + + const size_t scales_len = N * BlockCountK; + std::vector scales(scales_len); + for (size_t i = 0; i < scales_len; i++) { + uint32_t bits; + static_assert(sizeof(bits) == sizeof(QuantBScaleBegin[i]), "Unexpected float size"); + std::memcpy(&bits, &QuantBScaleBegin[i], sizeof(bits)); + scales[i] = static_cast(bits >> 16); + } + + kai_run_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0(1, N, K, nr, kr, sr, BlkLen, + reinterpret_cast(QuantBDataBegin), BlockCountK * BlkLen / 2, + nullptr, scales.data(), BlockCountK * sizeof(uint16_t), + PackedQuantBDataBegin, 0, ¶ms); + + break; } + // Asymmetric Q4 block-quantized RHS path + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: { + // Packed B is shared by GEMM and GEMV, so both kernels must use this RHS layout. + const auto& k = GetKleidiAIQai4GemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t nr = ukernel.get_nr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + + assert(QuantBScaleBegin != nullptr); + assert(QuantBZPBegin != nullptr); + kai_rhs_pack_nxk_qai4c32p_params params{}; + params.lhs_zero_point = 1; + params.rhs_zero_point = 8; + + const size_t zp_stride = MlasDivRoundup(BlockCountK, 2); + const size_t rhs_stride = K / 2; + // Reuse one kernel-width panel to keep scratch memory independent of N. + std::vector zero_offsets(nr * BlockCountK); + std::vector rhs_for_kai(nr * rhs_stride); + const auto* rhs = reinterpret_cast(QuantBDataBegin); + + for (size_t panel_start = 0; panel_start < N; panel_start += nr) { + const size_t panel_rows = std::min(nr, N - panel_start); + + // Getting zero-points for the current scratch buffer panel from provided packed buffer + for (size_t panel_n = 0; panel_n < panel_rows; ++panel_n) { + const size_t n = panel_start + panel_n; + for (size_t blk = 0; blk < BlockCountK; ++blk) { + const uint8_t zp_byte = + static_cast(QuantBZPBegin[n * zp_stride + blk / 2]); + const uint8_t zp = (blk & 1) == 0 ? (zp_byte & 0x0F) : (zp_byte >> 4); + const size_t src_idx = n * BlockCountK + blk; + const size_t panel_idx = panel_n * BlockCountK + blk; + // ORT stores asymmetric Q4 as uint4; KAI QAI4 interprets it as signed int4 + // offset by rhs_zero_point. + const float kai_zp_offset = + static_cast(params.rhs_zero_point) - static_cast(zp); + zero_offsets[panel_idx] = kai_zp_offset * QuantBScaleBegin[src_idx]; + } + } + + // Rearrange high/low nibble in the current scratch buffer panel for KAI-expected format + const uint8_t* rhs_panel = rhs + panel_start * rhs_stride; + const size_t rhs_panel_size = panel_rows * rhs_stride; + for (size_t i = 0; i < rhs_panel_size; ++i) { + rhs_for_kai[i] = + static_cast( + ((rhs_panel[i] & 0x0F) << 4) | ((rhs_panel[i] & 0xF0) >> 4)); + } + + const float* scales_panel = QuantBScaleBegin + panel_start * BlockCountK; + + // Pack the current scratch buffer panel using layout based on what the kernel expects + switch (k.rhs_layout) { + case KaiQ4RhsPackLayout::AsymmetricNxK: { + const size_t packed_offset = + kai_get_rhs_packed_offset_rhs_pack_nxk_qai4c32p_qau4c32s0s1_f32_f32_f32_neon( + panel_start, K, nr, kr, BlkLen); + kai_run_rhs_pack_nxk_qai4c32p_qau4c32s0s1_f32_f32_f32_neon( + 1, panel_rows, K, nr, kr, sr, BlkLen, + rhs_for_kai.data(), + zero_offsets.data(), + nullptr, + scales_panel, + PackedQuantBDataBegin + packed_offset, + 0, ¶ms); + break; + } + case KaiQ4RhsPackLayout::AsymmetricNxKInterleavedNrx4: { + const size_t packed_offset = + kai_get_rhs_packed_offset_rhs_pack_nxk_qai4c32ps1s0nrx4_qau4c32s0s1_f32_f32_f32_neon( + panel_start, K, nr, kr, BlkLen); + kai_run_rhs_pack_nxk_qai4c32ps1s0nrx4_qau4c32s0s1_f32_f32_f32_neon( + 1, panel_rows, K, nr, kr, sr, BlkLen, + rhs_for_kai.data(), + zero_offsets.data(), + nullptr, + scales_panel, + PackedQuantBDataBegin + packed_offset, + 0, ¶ms); + break; + } + default: + assert(false); + } + } + break; + } + case KleidiAIQ4Backend::None: + break; } } } else @@ -468,32 +569,46 @@ QNBitGemmPerGemmWorkspaceSize( case SQNBIT_CompInt8: { // workspace buffer is used for block quantization of A to int8 #ifdef USE_KLEIDIAI - if (BlkBitWidth == 4 && UseKleidiAI(K, BlkLen, BackendKernelSelectorConfig)) { - const auto& k = (M == 1) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); - const auto& ukernel = k.ukernel; - - const size_t mr = ukernel.get_mr(); - const size_t kr = ukernel.get_kr(); - const size_t sr = ukernel.get_sr(); - size_t ws = kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32(M, K, mr, kr, sr); - if (HasZeroPoint) { - // Align so that AFloatBlkSum starts at a float-aligned offset - constexpr size_t FloatAlignment = alignof(float); - ws = (ws + FloatAlignment - 1) & ~(FloatAlignment - 1); - // Additional space for AFloatBlkSum: M * BlockCountK floats - const size_t BlockCountK = MlasDivRoundup(K, BlkLen); - ws += M * BlockCountK * sizeof(float); + if (BlkBitWidth == 4) { + switch (SelectKleidiAIQ4Backend(K, BlkLen, HasZeroPoint, BackendKernelSelectorConfig)) { + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: { + const auto& k = (M == 1) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t mr = ukernel.get_mr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + size_t ws = kai_get_lhs_packed_size_lhs_quant_pack_qai8dxp_f32(M, K, mr, kr, sr); + if (HasZeroPoint) { + // Align so that AFloatBlkSum starts at a float-aligned offset + constexpr size_t FloatAlignment = alignof(float); + ws = (ws + FloatAlignment - 1) & ~(FloatAlignment - 1); + // Additional space for AFloatBlkSum: M * BlockCountK floats + const size_t BlockCountK = MlasDivRoundup(K, BlkLen); + ws += M * BlockCountK * sizeof(float); + } + return ws; + } + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: { + const auto& k = (M == 1) ? GetKleidiAIQai4GemvUKernel() : GetKleidiAIQai4GemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t mr = ukernel.get_mr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + return kai_get_lhs_packed_size_lhs_quant_pack_qsi8d32pscalef32_f32_neon( + M, K, BlkLen, mr, kr, sr); + } + case KleidiAIQ4Backend::None: + break; } - return ws; - } else -#endif - { - // workspace buffer is used for block quantization of A to int8 - const size_t BlockCountK = MlasDivRoundup(K, BlkLen); - // QuantData + Scale + BlkSum - const size_t PerGemmWorkspaceSize = M * BlockCountK * (Q8BlkSize(BlkLen) + sizeof(float)); - return PerGemmWorkspaceSize; } +#endif + // workspace buffer is used for block quantization of A to int8 + const size_t BlockCountK = MlasDivRoundup(K, BlkLen); + // QuantData + Scale + BlkSum + const size_t PerGemmWorkspaceSize = M * BlockCountK * (Q8BlkSize(BlkLen) + sizeof(float)); + return PerGemmWorkspaceSize; } default: { return 0; @@ -522,15 +637,18 @@ QNBitGemmPerGemmWorkspaceAlignment( } // namespace bool -UseKleidiAI(size_t K, size_t BlkLen, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) +IsKleidiAIQ4ShapeSupported(size_t K, size_t BlkLen, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) { #ifdef USE_KLEIDIAI if (BackendKernelSelectorConfig != nullptr && !BackendKernelSelectorConfig->use_kleidiai) { return false; } - bool has_dotprod = MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot(); - return (BlkLen % 32) == 0 && (K % BlkLen) == 0 && has_dotprod; + if (BlkLen == 0) { + return false; + } + + return (BlkLen % 32) == 0 && (K % BlkLen) == 0; #else MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); MLAS_UNREFERENCED_PARAMETER(K); @@ -539,6 +657,28 @@ UseKleidiAI(size_t K, size_t BlkLen, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* #endif } +KleidiAIQ4Backend +SelectKleidiAIQ4Backend(size_t K, size_t BlkLen, bool HasZp, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) +{ + if (!IsKleidiAIQ4ShapeSupported(K, BlkLen, BackendKernelSelectorConfig)) { + return KleidiAIQ4Backend::None; + } + + const auto& cpuid = MLAS_CPUIDINFO::GetCPUIDInfo(); + const bool has_neon_q4 = cpuid.HasArmNeon_I8MM() || cpuid.HasArmNeonDot(); + + // If zero-points input present, try asymmetric path + if (HasZp && (cpuid.HasArm_SME2() || has_neon_q4)) { + return KleidiAIQ4Backend::Qsi8d32pQai4c32p; + } + + if (!HasZp && has_neon_q4) { + return KleidiAIQ4Backend::Qai8dxpQsi4c32p; + } + + return KleidiAIQ4Backend::None; +} + template size_t SQ8BitGemmKernel_BlkSum_CompInt8( @@ -644,6 +784,8 @@ GetMlasQNBitGemmDispatchNeon( d.SQ4BitGemmKernel_CompInt8 = sqnbitgemm_neon::SQ4BitGemmKernel_CompInt8; d.QuantizeARow_CompInt8 = sqnbitgemm_neon::QuantizeARow_CompInt8; d.UsePacked_CompInt8 = sqnbitgemm_neon::UsePacked_CompInt8; + d.NeedsPackedZpCorrection_CompInt8 = sqnbitgemm_neon::NeedsPackedZpCorrection_CompInt8; + d.PackedQ4BitGemmNAlignment_CompInt8 = sqnbitgemm_neon::PackedQ4BitGemmNAlignment_CompInt8; d.QuantizeARowComputeBlkSum_CompInt8 = sqnbitgemm_neon::QuantizeARowComputeBlkSum_CompInt8; d.SQ8BitGemmKernel_BlkSum_CompInt8 = sqnbitgemm_neon::SQ8BitGemmKernel_BlkSum_CompInt8; diff --git a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h index b211bd06ee477..892d5acf0bc54 100644 --- a/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h +++ b/onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h @@ -141,6 +141,22 @@ UsePacked_CompInt8( const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ); +bool +NeedsPackedZpCorrection_CompInt8( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig +); + +size_t +PackedQ4BitGemmNAlignment_CompInt8( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig +); + void QuantizeARow_CompInt8( size_t BlkLen, @@ -224,6 +240,7 @@ QuantizeA_Packed_CompInt8( const float* A, size_t CountM, size_t CountK, + bool HasZeroPoint, std::byte* QuantA, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ); @@ -239,6 +256,8 @@ SQ4BitGemmKernel_Packed_CompInt8( const size_t RangeStartN, const size_t RangeCountN, size_t CountK, + bool HasQuantBZeroPoint, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig, size_t ldc, const float *Bias ); @@ -266,7 +285,16 @@ ApplyBZpCorrection( #endif bool -UseKleidiAI(size_t K, size_t BlkLen, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); +IsKleidiAIQ4ShapeSupported(size_t K, size_t BlkLen, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); + +enum class KleidiAIQ4Backend { + None, + Qai8dxpQsi4c32p, // 4-bit symmetric block-quantized RHS + Qsi8d32pQai4c32p, // 4-bit asymmetric block-quantized RHS +}; + +KleidiAIQ4Backend +SelectKleidiAIQ4Backend(size_t K, size_t BlkLen, bool HasZp, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig); // // General helpers. diff --git a/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp index 5b09495436eaf..502100973c383 100644 --- a/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp +++ b/onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp @@ -27,6 +27,7 @@ Module Name: #ifdef USE_KLEIDIAI #include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qai8dxp_f32.h" +#include "kai/ukernels/matmul/pack/kai_lhs_quant_pack_qsi8d32pscalef32_f32_neon.h" #include "kai_ukernel_interface.h" #endif @@ -134,42 +135,126 @@ QuantizeBlock( bool UsePacked_CompInt8(size_t K, size_t BlkLen, bool HasZp, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) { + // Use KleidiAI packed path for both symmetric and asymmetric. + return SelectKleidiAIQ4Backend(K, BlkLen, HasZp, BackendKernelSelectorConfig) != KleidiAIQ4Backend::None; +} + +bool +NeedsPackedZpCorrection_CompInt8( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig +) +{ + KleidiAIQ4Backend backend = SelectKleidiAIQ4Backend(K, BlkLen, HasZp, BackendKernelSelectorConfig); + return HasZp && backend == KleidiAIQ4Backend::Qai8dxpQsi4c32p; +} + +#ifdef USE_KLEIDIAI +template +size_t GetKleidiAIQ4NAlignment(const KleidiAIKernel& k) +{ + const auto& ukernel = k.ukernel; + const size_t n_step = ukernel.get_n_step(); + const size_t nr = ukernel.get_nr(); + const size_t n_alignment = std::max(n_step, nr); + + // KAI packed RHS ranges must start on both the kernel step and packed-panel boundaries. + assert(n_step != 0 && nr != 0); + assert((n_alignment % n_step) == 0 && (n_alignment % nr) == 0); + + return n_alignment; +} +#endif + +size_t +PackedQ4BitGemmNAlignment_CompInt8( + size_t K, + size_t BlkLen, + bool HasZp, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) +{ +#ifdef USE_KLEIDIAI + switch (SelectKleidiAIQ4Backend(K, BlkLen, HasZp, BackendKernelSelectorConfig)) { + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: + return std::max( + GetKleidiAIQ4NAlignment(GetKleidiAIGemmUKernel()), + GetKleidiAIQ4NAlignment(GetKleidiAIGemvUKernel())); + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: + return std::max( + GetKleidiAIQ4NAlignment(GetKleidiAIQai4GemmUKernel()), + GetKleidiAIQ4NAlignment(GetKleidiAIQai4GemvUKernel())); + case KleidiAIQ4Backend::None: + break; + } +#else + MLAS_UNREFERENCED_PARAMETER(K); + MLAS_UNREFERENCED_PARAMETER(BlkLen); MLAS_UNREFERENCED_PARAMETER(HasZp); - // Use KleidiAI packed path for both symmetric and asymmetric (with ZP correction). - return UseKleidiAI(K, BlkLen, BackendKernelSelectorConfig); + MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); +#endif + return MLAS_QGEMM_STRIDEN_THREAD_ALIGN; } #ifdef USE_KLEIDIAI void QuantizeA_Packed_CompInt8( - size_t, + size_t BlkLen, const float* A, size_t CountM, size_t CountK, + bool HasZeroPoint, std::byte* QuantA, const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig ) { // Currently this routine only supports KleidiAI packed quantization of A assert(BackendKernelSelectorConfig == nullptr || BackendKernelSelectorConfig->use_kleidiai); - MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig); - const auto& k = (CountM == 1) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); - const auto& ukernel = k.ukernel; + switch (SelectKleidiAIQ4Backend(CountK, BlkLen, HasZeroPoint, BackendKernelSelectorConfig)) { + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: { + const auto& k = (CountM == 1) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t mr = ukernel.get_mr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); + + const size_t src_stride = CountK * sizeof(float); + const size_t lhs_offset = kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32(0, src_stride); + const size_t lhs_packed_offset = kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32( + 0, CountK, mr, kr, sr); + + const float* src_ptr = reinterpret_cast(reinterpret_cast(A) + lhs_offset); + void* dst_ptr = QuantA + lhs_packed_offset; - const size_t mr = ukernel.get_mr(); - const size_t kr = ukernel.get_kr(); - const size_t sr = ukernel.get_sr(); + kai_run_lhs_quant_pack_qai8dxp_f32(CountM, CountK, mr, kr, sr, 0, src_ptr, src_stride, dst_ptr); + return; + } + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: { + const auto& k = (CountM == 1) ? GetKleidiAIQai4GemvUKernel() : GetKleidiAIQai4GemmUKernel(); + const auto& ukernel = k.ukernel; + + const size_t mr = ukernel.get_mr(); + const size_t kr = ukernel.get_kr(); + const size_t sr = ukernel.get_sr(); - const size_t src_stride = CountK * sizeof(float); - const size_t lhs_offset = kai_get_lhs_offset_lhs_quant_pack_qai8dxp_f32(0, src_stride); - const size_t lhs_packed_offset = kai_get_lhs_packed_offset_lhs_quant_pack_qai8dxp_f32( - 0, CountK, mr, kr, sr); + const size_t src_stride = CountK * sizeof(float); + const size_t lhs_offset = kai_get_lhs_offset_lhs_quant_pack_qsi8d32pscalef32_f32_neon(0, src_stride); + const size_t lhs_packed_offset = kai_get_lhs_packed_offset_lhs_quant_pack_qsi8d32pscalef32_f32_neon(0, CountK, BlkLen, mr, kr, sr); - const float* src_ptr = reinterpret_cast(reinterpret_cast(A) + lhs_offset); - void* dst_ptr = QuantA + lhs_packed_offset; + const float* src_ptr = reinterpret_cast(reinterpret_cast(A) + lhs_offset); + void* dst_ptr = QuantA + lhs_packed_offset; - kai_run_lhs_quant_pack_qai8dxp_f32(CountM, CountK, mr, kr, sr, 0, src_ptr, src_stride, dst_ptr); + kai_run_lhs_quant_pack_qsi8d32pscalef32_f32_neon(CountM, CountK, BlkLen, mr, kr, sr, 0, src_ptr, src_stride, dst_ptr); + return; + } + case KleidiAIQ4Backend::None: + // None KleidiAIQ4Backend shouldn't get to this point + assert(false); + return; + } } void @@ -2502,8 +2587,31 @@ MlasQ8Int8GemmKernelNeon( } #ifdef USE_KLEIDIAI -void -SQ4BitGemmKernel_Packed_CompInt8( +size_t +GetKleidiAIQ4LhsPackedOffset( + const kai_matmul_clamp_f32_qai8dxp_qsi4c32p_ukernel& ukernel, + size_t m_idx, + size_t k, + size_t +) +{ + return ukernel.get_lhs_packed_offset(m_idx, k); +} + +size_t +GetKleidiAIQ4LhsPackedOffset( + const kai_matmul_clamp_f32_qsi8d32p_qai4c32p_ukernel& ukernel, + size_t m_idx, + size_t k, + size_t bl +) +{ + return ukernel.get_lhs_packed_offset(m_idx, k, bl); +} + +template +void RunKleidiAIQ4Packed( + const KleidiAIKernel& k, size_t BlkLen, const std::byte* QuantA, const std::byte* PackedQuantBData, @@ -2517,12 +2625,11 @@ SQ4BitGemmKernel_Packed_CompInt8( const float* Bias ) { - const auto& k = (RangeCountM == 1 && RangeStartM == 0) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); const auto& ukernel = k.ukernel; const size_t dst_stride = ldc * sizeof(float); - const size_t lhs_packed_offset = ukernel.get_lhs_packed_offset(RangeStartM, CountK); + const size_t lhs_packed_offset = GetKleidiAIQ4LhsPackedOffset(ukernel, RangeStartM, CountK, BlkLen); const size_t rhs_packed_offset = ukernel.get_rhs_packed_offset(RangeStartN, CountK, BlkLen); const size_t dst_offset = ukernel.get_dst_offset(RangeStartM, RangeStartN, dst_stride); @@ -2542,6 +2649,46 @@ SQ4BitGemmKernel_Packed_CompInt8( } } } + +void +SQ4BitGemmKernel_Packed_CompInt8( + size_t BlkLen, + const std::byte* QuantA, + const std::byte* PackedQuantBData, + float* C, + const size_t RangeStartM, + const size_t RangeCountM, + const size_t RangeStartN, + const size_t RangeCountN, + size_t CountK, + bool HasQuantBZeroPoint, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig, + size_t ldc, + const float* Bias +) +{ + switch (SelectKleidiAIQ4Backend(CountK, BlkLen, HasQuantBZeroPoint, BackendKernelSelectorConfig)) { + case KleidiAIQ4Backend::Qai8dxpQsi4c32p: { + const auto& k = (RangeCountM == 1 && RangeStartM == 0) ? GetKleidiAIGemvUKernel() : GetKleidiAIGemmUKernel(); + + RunKleidiAIQ4Packed(k, BlkLen, QuantA, PackedQuantBData, C, RangeStartM, + RangeCountM, RangeStartN, RangeCountN, + CountK, ldc, Bias); + return; + } + case KleidiAIQ4Backend::Qsi8d32pQai4c32p: { + const auto& k = (RangeCountM == 1 && RangeStartM == 0) ? GetKleidiAIQai4GemvUKernel() : GetKleidiAIQai4GemmUKernel(); + + RunKleidiAIQ4Packed(k, BlkLen, QuantA, PackedQuantBData, C, RangeStartM, + RangeCountM, RangeStartN, RangeCountN, + CountK, ldc, Bias); + return; + } + case KleidiAIQ4Backend::None: + assert(false); + return; + } +} #endif } // namespace sqnbitgemm_neon diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 51895938ce3aa..acc36656df002 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -88,6 +88,7 @@ struct TestOptions { bool has_zero_point{false}; bool zp_is_4bit{true}; + bool zero_points_are_initializers{true}; bool has_g_idx{false}; bool has_bias{false}; @@ -111,6 +112,7 @@ struct TestOptions { << ", accuracy_level:" << opts.accuracy_level << ", has_zero_point:" << opts.has_zero_point << ", zp_is_4bit:" << opts.zp_is_4bit + << ", zero_points_are_initializers:" << opts.zero_points_are_initializers << ", has_g_idx:" << opts.has_g_idx << ", has_bias:" << opts.has_bias; } @@ -215,7 +217,7 @@ void RunTest(const TestOptions& opts, if (zp_is_4bit) { auto zp_shape = opts.legacy_shape ? std::vector{N * zero_point_blob_size} : std::vector{N, zero_point_blob_size}; - test.AddInput("zero_points", zp_shape, zp, true); + test.AddInput("zero_points", zp_shape, zp, opts.zero_points_are_initializers); } else { std::vector zp_f; zp_f.reserve(q_zp_size_in_bytes * 2); @@ -230,11 +232,11 @@ void RunTest(const TestOptions& opts, } if constexpr (std::is_same_v) { - test.AddInput("zero_points", scales_shape, zp_f, true); + test.AddInput("zero_points", scales_shape, zp_f, opts.zero_points_are_initializers); } else if constexpr (std::is_same_v) { - test.AddInput("zero_points", scales_shape, FloatsToMLFloat16s(zp_f), true); + test.AddInput("zero_points", scales_shape, FloatsToMLFloat16s(zp_f), opts.zero_points_are_initializers); } else if constexpr (std::is_same_v) { - test.AddInput("zero_points", scales_shape, FloatsToBFloat16s(zp_f), true); + test.AddInput("zero_points", scales_shape, FloatsToBFloat16s(zp_f), opts.zero_points_are_initializers); } } } else { @@ -670,6 +672,46 @@ TEST(MatMulNBits, SharedPrepackedWeights_AddInitializer) { } } +// Covers backends that fold asymmetric scales and zero points into the shared B buffer. +TEST(MatMulNBits, SharedPrepackedWeights_AsymmetricPackedScales) { + auto opts = MakeSharingTestOptions(288, 1024, /*block_size*/ 128, /*accuracy_level*/ 4, + /*has_zero_point*/ true, /*has_bias*/ false, + PrepackSharingMode::kAddInitializer); + opts.M = 1; + RunTest(opts); +} + +// Uses a KleidiAI-compatible Q4 CompInt8 shape. Runtime zero points must not reuse a B pack +// generated without them. +TEST(MatMulNBits, DynamicZeroPoints_AsymmetricCompInt8) { + TestOptions opts{}; + opts.M = 1; + opts.N = 288; + opts.K = 1024; + opts.block_size = 128; + opts.accuracy_level = 4; + opts.has_zero_point = true; + opts.zero_points_are_initializers = false; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + RunTest(opts); +} + +// Same runtime-ZP case with the shared container enabled. PrePack must decline B packing, so the +// second session must not adopt a shared B buffer built without those runtime zero points. +TEST(MatMulNBits, SharedPrepackedWeights_DynamicZeroPoints_AsymmetricCompInt8) { + if (!MlasQNBitGemmScalesPacked(1024, QBits, 128, SQNBIT_CompInt8, true, nullptr)) { + GTEST_SKIP() << "KleidiAI Q4 packed-scales path is not active."; + } + + auto opts = MakeSharingTestOptions(288, 1024, /*block_size*/ 128, /*accuracy_level*/ 4, + /*has_zero_point*/ true, /*has_bias*/ false, + PrepackSharingMode::kAddInitializerExpectNoPrepack); + opts.M = 1; + opts.zero_points_are_initializers = false; + RunTest(opts); +} + // Negative control: with the shared container present but neither opt-in mechanism enabled, no // pre-packed weights are shared across sessions. TEST(MatMulNBits, SharedPrepackedWeights_NotSharedWithoutOptIn) { diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc b/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc index 97566afe02489..d4aa30b8b0881 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc +++ b/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc @@ -26,6 +26,7 @@ void CheckSharedPrepackedWeights(OpTester& test, PrepackSharingMode mode, switch (mode) { case PrepackSharingMode::kAddInitializer: + case PrepackSharingMode::kAddInitializerExpectNoPrepack: // Register B as an explicitly shared initializer (the pre-existing sharing mechanism). Tensor::InitOrtValue(DataTypeImpl::GetType(), TensorShape(b_dims), b_data.data(), OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator), b_ortvalue); @@ -61,6 +62,23 @@ void CheckSharedPrepackedWeights(OpTester& test, PrepackSharingMode mode, const auto number_of_elements_in_shared_container = test.GetNumPrePackedWeightsShared(); + if (mode == PrepackSharingMode::kAddInitializerExpectNoPrepack) { + ASSERT_EQ(number_of_pre_packed_weights_counter_session_1, static_cast(0)); + ASSERT_EQ(number_of_elements_in_shared_container, static_cast(0)); + + { + size_t number_of_pre_packed_weights_counter_session_2 = 0; + auto ep_vec = cpu_ep(); + test.Run(so, OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &ep_vec, {}, + &number_of_pre_packed_weights_counter_session_2, + &number_of_shared_pre_packed_weights_counter); + ASSERT_EQ(number_of_pre_packed_weights_counter_session_2, static_cast(0)); + ASSERT_EQ(number_of_shared_pre_packed_weights_counter, static_cast(0)); + ASSERT_EQ(test.GetNumPrePackedWeightsShared(), static_cast(0)); + } + return; + } + if (mode == PrepackSharingMode::kNoSharing) { // Without opting in, pre-packed weights must not be placed in the shared container. ASSERT_EQ(number_of_elements_in_shared_container, static_cast(0)); diff --git a/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h b/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h index 1de0bbaa4bb85..4c6064dc526fe 100644 --- a/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h +++ b/onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h @@ -15,6 +15,8 @@ enum class PrepackSharingMode { // Legacy path: the weight is explicitly registered as a shared initializer via // SessionOptions::AddInitializer. kAddInitializer, + // Shared-initializer path configured, but the kernel is expected to decline prepacking. + kAddInitializerExpectNoPrepack, // Negative control: the shared container exists but neither opt-in mechanism is used, so no // cross-session sharing must happen. kNoSharing, @@ -23,8 +25,8 @@ enum class PrepackSharingMode { // Runs the already-configured MatMulNBits OpTester in two CPU sessions that share the same // pre-packed weights container and asserts that the pre-packed weights are shared as expected. // This logic is independent of the weight bit width, so it is shared by the 4-bit and 8-bit tests. -// `b_dims`/`b_data` describe the quantized B initializer and are only needed for the -// PrepackSharingMode::kAddInitializer path (to register B as a shared initializer). +// `b_dims`/`b_data` describe the quantized B initializer and are only needed for modes that +// register B as a shared initializer. void CheckSharedPrepackedWeights(OpTester& test, PrepackSharingMode mode, const std::vector& b_dims, std::vector& b_data); diff --git a/onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp b/onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp index 342dd9ad34750..8d2eb9bf9827a 100644 --- a/onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp +++ b/onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp @@ -64,7 +64,8 @@ class MlasSQNBitGemmTest : public MlasTestBase { size_t ldc, void* Workspace, MLAS_QNBIT_GEMM_COMPUTE_TYPE ComputeType, - MLAS_THREADPOOL* Threadpool) { + MLAS_THREADPOOL* Threadpool, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig) { MLAS_QNBIT_GEMM_DATA_PARAMS params; params.A = A; params.lda = lda; @@ -81,7 +82,8 @@ class MlasSQNBitGemmTest : public MlasTestBase { params.QuantBZeroPoint = QuantBZeroPoint; params.PostProcessor = nullptr; - MlasQNBitGemmBatch(M, N, K, 1, BlkBitWidth, BlkLen, ComputeType, ¶ms, Workspace, Threadpool, nullptr); + MlasQNBitGemmBatch(M, N, K, 1, BlkBitWidth, BlkLen, ComputeType, ¶ms, Workspace, Threadpool, + BackendKernelSelectorConfig); } void QuantizeA(size_t M, size_t K, const float* A, int8_t* QuantAData, float* QuantAScale) { @@ -202,7 +204,8 @@ class MlasSQNBitGemmTest : public MlasTestBase { public: void Test(size_t M, size_t N, size_t K, MLAS_QNBIT_GEMM_COMPUTE_TYPE ComputeType, - bool WithThreadpool, bool Symmetric, bool WithBias) { + bool WithThreadpool, bool Symmetric, bool WithBias, + const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr) { MLAS_THREADPOOL* Threadpool = WithThreadpool ? GetMlasThreadPool() : nullptr; const float* A = BufferA.GetBuffer(K * M); @@ -265,19 +268,21 @@ class MlasSQNBitGemmTest : public MlasTestBase { } void* Workspace = nullptr; - if (const auto WorkspaceSize = MlasQNBitGemmBatchWorkspaceSize(M, N, K, 1, BlkBitWidth, BlkLen, !Symmetric, ComputeType, nullptr); + if (const auto WorkspaceSize = MlasQNBitGemmBatchWorkspaceSize(M, N, K, 1, BlkBitWidth, BlkLen, !Symmetric, + ComputeType, BackendKernelSelectorConfig); WorkspaceSize > 0) { Workspace = BufferWorkspace.GetBuffer(WorkspaceSize); } void* PackedQuantBDataWorkspace = nullptr; - if (const auto PackedQuantBDataSize = MlasQNBitGemmPackQuantBDataSize(N, K, BlkBitWidth, BlkLen, !Symmetric, ComputeType, nullptr); + if (const auto PackedQuantBDataSize = MlasQNBitGemmPackQuantBDataSize(N, K, BlkBitWidth, BlkLen, !Symmetric, + ComputeType, BackendKernelSelectorConfig); PackedQuantBDataSize > 0) { PackedQuantBDataWorkspace = BufferPackedQuantBData.GetBuffer(PackedQuantBDataSize); bool has_zp_input = QuantBZeroPoint != nullptr; MlasQNBitGemmPackQuantBData(N, K, BlkBitWidth, BlkLen, ComputeType, QuantBData, PackedQuantBDataWorkspace, QuantBScale, has_zp_input, QuantBZeroPoint, - GetMlasThreadPool(), nullptr); + GetMlasThreadPool(), BackendKernelSelectorConfig); } CallGemm(M, N, K, @@ -287,7 +292,8 @@ class MlasSQNBitGemmTest : public MlasTestBase { C, /* ldc */ N, Workspace, ComputeType, - Threadpool); + Threadpool, + BackendKernelSelectorConfig); if (ComputeType == SQNBIT_CompFp32) { CallReferenceGemm_CompFp32(M, N, K, A, QuantBData, QuantBScale, QuantBZeroPoint, Bias, CReference); @@ -308,6 +314,35 @@ class MlasSQNBitGemmTest : public MlasTestBase { } } + void TestAsymmetricKleidiAICompInt8(size_t M, size_t N, size_t K, bool WithThreadpool, bool WithBias) { +#if defined(MLAS_TARGET_ARM64) + MLAS_BACKEND_KERNEL_SELECTOR_CONFIG config; + config.use_kleidiai = true; + constexpr bool HasZeroPoint = true; + + if (!MlasIsQNBitGemmAvailable(BlkBitWidth, BlkLen, SQNBIT_CompInt8) || + !MlasQNBitGemmScalesPacked(K, BlkBitWidth, BlkLen, SQNBIT_CompInt8, HasZeroPoint, &config)) { + GTEST_SKIP() << "KleidiAI packed asymmetric SQ4 path is unavailable."; + } + + ASSERT_GT(MlasQNBitGemmPackQuantBDataSize(N, K, BlkBitWidth, BlkLen, HasZeroPoint, + SQNBIT_CompInt8, &config), + 0u); + ASSERT_GT(MlasQNBitGemmBatchWorkspaceSize(M, N, K, 1, BlkBitWidth, BlkLen, HasZeroPoint, + SQNBIT_CompInt8, &config), + 0u); + + Test(M, N, K, SQNBIT_CompInt8, WithThreadpool, /*Symmetric=*/false, WithBias, &config); +#else + (void)M; + (void)N; + (void)K; + (void)WithThreadpool; + (void)WithBias; + GTEST_SKIP() << "KleidiAI Q4 tests require ARM64."; +#endif + } + public: static const char* GetTestSuiteName() { static std::string suite_name = std::string("SQNBitGemm") + @@ -420,6 +455,56 @@ class SQNBitGemmShortExecuteTest : public MlasTestFixture> { + public: + explicit SQNBitGemmKleidiAIShortExecuteTest(size_t M, size_t N, size_t K, + bool WithThreadpool, bool WithBias) + : M_(M), + N_(N), + K_(K), + WithThreadpool_(WithThreadpool), + WithBias_(WithBias) { + } + + void TestBody() override { + MlasTestFixture>::mlas_tester->TestAsymmetricKleidiAICompInt8( + M_, N_, K_, WithThreadpool_, WithBias_); + } + + static size_t RegisterSingleTest(const char* test_name, size_t M, size_t N, size_t K, + bool WithThreadpool, bool WithBias) { + testing::RegisterTest( + MlasSQNBitGemmTest<4, 128>::GetTestSuiteName(), + test_name, + nullptr, + test_name, + __FILE__, + __LINE__, + [=]() -> MlasTestFixture>* { + return new SQNBitGemmKleidiAIShortExecuteTest(M, N, K, WithThreadpool, WithBias); + }); + + return 1; + } + + static size_t RegisterShortExecuteTests() { + size_t tests_registered = 0; + + tests_registered += RegisterSingleTest( + "KleidiAIAsymGemv_M1_N257_K128", 1, 257, 128, /*WithThreadpool=*/false, /*WithBias=*/true); + tests_registered += RegisterSingleTest( + "KleidiAIAsymGemm_M5_N257_K128", 5, 257, 128, /*WithThreadpool=*/true, /*WithBias=*/true); + tests_registered += RegisterSingleTest( + "KleidiAIAsymGemv_M1_N288_K1024_NoBias", 1, 288, 1024, /*WithThreadpool=*/false, /*WithBias=*/false); + + return tests_registered; + } + + private: + size_t M_, N_, K_; + bool WithThreadpool_, WithBias_; +}; + static size_t SQNBitGemmRegisterAllShortExecuteTests() { size_t count = 0; @@ -428,6 +513,7 @@ static size_t SQNBitGemmRegisterAllShortExecuteTests() { count += SQNBitGemmShortExecuteTest<4, 64>::RegisterShortExecuteTests(); count += SQNBitGemmShortExecuteTest<4, 128>::RegisterShortExecuteTests(); count += SQNBitGemmShortExecuteTest<4, 256>::RegisterShortExecuteTests(); + count += SQNBitGemmKleidiAIShortExecuteTest::RegisterShortExecuteTests(); return count; }