Skip to content

Route fp16 HQNBIT_CompInt8 (4-bit and 8-bit) through fp32 MLAS path in MatMulNBits #27820

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jambayk merged 5 commits into from
Mar 25, 2026
Merged

Route fp16 HQNBIT_CompInt8 (4-bit and 8-bit) through fp32 MLAS path in MatMulNBits #27820

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c0c51be
Route fp16 HQNBIT_CompInt8 through fp32 MLAS path for 4-bit and 8-bit
jambayk Mar 24, 2026
a217afb
Remove dead HQ4BitGemm_CompInt8 and HQ8BitGemm_CompInt8 MLAS code
jambayk Mar 24, 2026
0931a5f
lint
jambayk Mar 24, 2026
316b13f
Fix HQNBIT_CompInt8 PrePack bugs for 4-bit and 8-bit
jambayk Mar 24, 2026
138318a
Address review: ORT_ENFORCE for scales, move SQNBIT check to GetCompu…
jambayk Mar 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
192 changes: 172 additions & 20 deletions onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc
View file
Open in desktop
Comment thread
vraspar marked this conversation as resolved.
Original file line number Diff line number Diff line change
Expand Up @@ -258,20 +258,50 @@
prepacked_weights->buffer_sizes_.push_back(packed_b_size_);
}
} else {
packed_b_size_ = MlasQNBitGemmPackQuantBDataSize(N_, K_, nbits_, block_size_, has_zp_input_, compute_type_, &mlas_backend_kernel_selector_config_);
// For HQNBIT_CompInt8 4-bit, route through SQNBIT_CompInt8 for sizing and packing.
// This gets KleidiAI-sized buffer when available and packs B+scales in KleidiAI format.
// The same approach is already used for 8-bit HQNBIT_CompInt8 (size function treats them identically).
const auto effective_compute_type = (compute_type_ == HQNBIT_CompInt8 && nbits_ == 4)
? SQNBIT_CompInt8
: compute_type_;

packed_b_size_ = MlasQNBitGemmPackQuantBDataSize(N_, K_, nbits_, block_size_, has_zp_input_, effective_compute_type, &mlas_backend_kernel_selector_config_);
Comment thread
jambayk marked this conversation as resolved.
Outdated
if (packed_b_size_ == 0) {
return Status::OK();
}

auto qptr = tensor.DataRaw();
// For HQNBIT compute types, scales are fp16 and cannot be passed directly
// to packing functions that expect float*. Pass nullptr here; scales will
// be properly converted and packed in a subsequent PrePack call.
auto scale_ptr = (scales && compute_type_ != HQNBIT_CompInt8 && compute_type_ != HQNBIT_CompFp16)
? scales->DataRaw()
: nullptr;
const void* scale_ptr = nullptr;

// For HQNBIT_CompInt8 4-bit: convert constant fp16 scales to fp32 for packing.
// KleidiAI bakes scales into packed B, so they must be available now.
if (compute_type_ == HQNBIT_CompInt8 && nbits_ == 4 && scales) {
auto sptr_fp16 = scales->Data<MLFloat16>();
auto scales_size = static_cast<size_t>(scales->Shape().Size());
scales_fp32_ = IAllocator::MakeUniquePtr<float>(alloc, scales_size, true);
MlasConvertHalfToFloatBuffer(sptr_fp16, scales_fp32_.get(), scales_size);
scale_ptr = scales_fp32_.get();
} else if (scales && compute_type_ != HQNBIT_CompInt8 && compute_type_ != HQNBIT_CompFp16) {
// For non-HQNBIT compute types, scales are already float.
scale_ptr = scales->DataRaw();
}
Comment thread
jambayk marked this conversation as resolved.

packed_b_ = IAllocator::MakeUniquePtr<void>(alloc, packed_b_size_, true);
MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, compute_type_, qptr, packed_b_.get(), scale_ptr,
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_);

#if defined(MLAS_TARGET_ARM64)
// For KleidiAI asymmetric 4-bit path: compute BZpCorr now while scales and zero_points are accessible.
if (compute_type_ == HQNBIT_CompInt8 && nbits_ == 4 && has_zp_input_ && scales_fp32_) {
const Tensor* zp_tensor = nullptr;
OpKernel::Info().TryGetConstantInput(InputIndex::zero_points, &zp_tensor);
if (zp_tensor != nullptr) {
auto zptr = zp_tensor->Data<uint8_t>();
MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(),
scales_fp32_.get(), has_zp_input_, zptr, nullptr, &mlas_backend_kernel_selector_config_);
}
}
#endif // MLAS_TARGET_ARM64
}
is_packed = true;
} else if (compute_type_ == SQNBIT_CompInt8) {
Expand Down Expand Up @@ -337,17 +367,51 @@
}
}
#endif // MLAS_TARGET_ARM64
} else if (compute_type_ == HQNBIT_CompInt8 && nbits_ == 8) {
// For 8-bit HQNBIT_CompInt8, scales are fp16 but the SQ8 packing functions expect float.
} 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<float> with SQNBIT_CompInt8.
if (input_idx == InputIndex::scales && packed_b_ != nullptr) {
auto sptr_fp16 = tensor.Data<MLFloat16>();
std::vector<float> scales_fp32(static_cast<size_t>(tensor.Shape().Size()));
MlasConvertHalfToFloatBuffer(sptr_fp16, scales_fp32.data(), scales_fp32.size());
MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(),
scales_fp32.data(), has_zp_input_, nullptr, nullptr,
&mlas_backend_kernel_selector_config_);
is_packed = false;
#if defined(MLAS_TARGET_ARM64)
// For 4-bit on ARM64: check if KleidiAI packs scales into B (scales already packed during B packing).
if (nbits_ == 4 &&
MlasQNBitGemmScalesPacked(K_, nbits_, block_size_, SQNBIT_CompInt8,
has_zp_input_, &mlas_backend_kernel_selector_config_)) {
// For asymmetric quantization, require zero_points to be constant for KleidiAI.
if (has_zp_input_) {
const Tensor* zp_tensor = nullptr;
OpKernel::Info().TryGetConstantInput(InputIndex::zero_points, &zp_tensor);
if (zp_tensor == nullptr) {
// zero_points is dynamic: fall back to non-KleidiAI path.
// Convert scales to fp32 for use at compute time.
auto sptr_fp16 = tensor.Data<MLFloat16>();
auto tensor_size = static_cast<size_t>(tensor.Shape().Size());
if (!scales_fp32_) {
scales_fp32_ = IAllocator::MakeUniquePtr<float>(alloc, tensor_size, true);
MlasConvertHalfToFloatBuffer(sptr_fp16, scales_fp32_.get(), tensor_size);
}
return Status::OK();
}
}

// BZpCorr was already computed during B packing in Step 1 (if applicable).
scales_are_packed_ = true;
is_packed = true;
} else

Check warning on line 400 in onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 If/else bodies with multiple statements require braces  [readability/braces] [4]

Raw Output:
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc:400:  If/else bodies with multiple statements require braces  [readability/braces] [4]

Check warning on line 400 in onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 If an else has a brace on one side, it should have it on both  [readability/braces] [5]

Raw Output:
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc:400:  If an else has a brace on one side, it should have it on both  [readability/braces] [5]
#endif // MLAS_TARGET_ARM64
{
// Non-KleidiAI path (or 8-bit): convert fp16 scales to fp32 and pack into workspace.
auto sptr_fp16 = tensor.Data<MLFloat16>();
auto tensor_size = static_cast<size_t>(tensor.Shape().Size());
if (!scales_fp32_) {
scales_fp32_ = IAllocator::MakeUniquePtr<float>(alloc, tensor_size, true);
MlasConvertHalfToFloatBuffer(sptr_fp16, scales_fp32_.get(), tensor_size);
}
MlasQNBitGemmPackQuantBData(N_, K_, nbits_, block_size_, SQNBIT_CompInt8, nullptr, packed_b_.get(),
scales_fp32_.get(), has_zp_input_, nullptr, nullptr,
&mlas_backend_kernel_selector_config_);
is_packed = false;
}
}

if (input_idx == InputIndex::zero_points && packed_b_ != nullptr) {
Expand All @@ -356,6 +420,15 @@
has_zp_input_, zptr, nullptr, &mlas_backend_kernel_selector_config_);
is_packed = false;
}

// Pre-convert fp16 bias to fp32 for use at compute time.
if (input_idx == InputIndex::bias) {
auto bptr_fp16 = tensor.Data<MLFloat16>();
auto tensor_size = static_cast<size_t>(tensor.Shape().Size());
bias_fp32_ = IAllocator::MakeUniquePtr<float>(alloc, tensor_size, true);
MlasConvertHalfToFloatBuffer(bptr_fp16, bias_fp32_.get(), tensor_size);
is_packed = false;
}
} else if (prefer_lut_gemm_) {
// Pack scales/zero_points for LUT GEMM if B was already packed but scales weren't available then
if (input_idx == InputIndex::scales && packed_b_ != nullptr) {
Expand Down Expand Up @@ -519,9 +592,7 @@
concurrency::ThreadPool* thread_pool,
const MatMulComputeHelper& helper) const {
const auto* a_data = a->Data<T1>();
const auto* scales_data = scales == nullptr ? nullptr : scales->Data<T1>();
const auto* zero_points_data = zero_points == nullptr ? nullptr : zero_points->DataRaw();
const auto* bias_data = bias == nullptr ? nullptr : bias->Data<T1>();
auto* y_data = y->MutableData<T1>();

const size_t batch_count = helper.OutputOffsets().size();
Expand All @@ -530,6 +601,87 @@
const size_t K = static_cast<size_t>(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).
if constexpr (std::is_same_v<T1, MLFloat16>) {
if (compute_type_ == HQNBIT_CompInt8) {
const auto* a_data_fp16 = a->Data<MLFloat16>();
const auto* bias_data_fp16 = bias == nullptr ? nullptr : bias->Data<MLFloat16>();

// Bulk convert A from fp16 to fp32.
auto a_size = static_cast<size_t>(a->Shape().Size());
auto tmp_a_data_ptr = IAllocator::MakeUniquePtr<float>(allocator, a_size, true);
Comment thread
vraspar marked this conversation as resolved.
MlasConvertHalfToFloatBuffer(a_data_fp16, tmp_a_data_ptr.get(), a_size);

// Use pre-converted fp32 scales, or nullptr if scales are baked into packed B (KleidiAI).
// For non-KleidiAI 4-bit: scales_fp32_ was set during PrePack.
// For 8-bit: scales are packed inside PackedQuantBDataStruct and extracted at dispatch.
float* scales_ptr = nullptr;
if (!scales_are_packed_ && scales_fp32_) {
scales_ptr = scales_fp32_.get();
} else if (!scales_are_packed_ && scales != nullptr) {
// Scales not pre-converted (shouldn't happen with our PrePack, but handle gracefully).
Comment thread
vraspar marked this conversation as resolved.
Outdated
auto scales_size = static_cast<size_t>(scales->Shape().Size());
auto tmp_scales = IAllocator::MakeUniquePtr<float>(allocator, scales_size, true);
Comment thread
jambayk marked this conversation as resolved.
Outdated
MlasConvertHalfToFloatBuffer(scales->Data<MLFloat16>(), tmp_scales.get(), scales_size);
scales_ptr = tmp_scales.get();
}

// Use pre-converted fp32 bias, or convert on the fly.
float* bias_ptr = nullptr;
IAllocatorUniquePtr<float> tmp_bias;
if (bias_data_fp16) {
if (bias_fp32_) {
bias_ptr = bias_fp32_.get();
} else {
auto bias_size = static_cast<size_t>(bias->Shape().Size());
tmp_bias = IAllocator::MakeUniquePtr<float>(allocator, bias_size, true);
MlasConvertHalfToFloatBuffer(bias_data_fp16, tmp_bias.get(), bias_size);
bias_ptr = tmp_bias.get();
}
}

// Allocate fp32 output buffer.
auto c_size = static_cast<size_t>(y->Shape().Size());
auto tmp_c = IAllocator::MakeUniquePtr<float>(allocator, c_size, true);

// Compute workspace sized for SQNBIT_CompInt8 (includes KleidiAI workspace when available).
IAllocatorUniquePtr<std::byte> workspace{};
const size_t workspace_size = MlasQNBitGemmBatchWorkspaceSize(
M, N, K, batch_count, nbits_, block_size_, zero_points, SQNBIT_CompInt8, &mlas_backend_kernel_selector_config_);
if (workspace_size > 0) {
workspace = IAllocator::MakeUniquePtr<std::byte>(allocator, workspace_size, true);
}

InlinedVector<MLAS_QNBIT_GEMM_DATA_PARAMS<float>> 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;
data[i].QuantBDataWorkspace = packed_b_.get();
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
data[i].QuantBScale = scales_ptr;
data[i].QuantBZeroPoint = zero_points_data;
data[i].Bias = bias_ptr;
data[i].C = tmp_c.get() + helper.OutputOffsets()[i];
data[i].ldc = N;
}

MlasQNBitGemmBatch(M, N, K, batch_count, nbits_, block_size_, SQNBIT_CompInt8, data.data(), workspace.get(),
thread_pool, &mlas_backend_kernel_selector_config_);

// Bulk convert output from fp32 to fp16.
MlasConvertFloatToHalfBuffer(tmp_c.get(), y_data, c_size);
return Status::OK();
}
}

// Standard path for non-HQNBIT_CompInt8 compute types (fp32 inputs, CompFp32, CompFp16, etc.)
const auto* scales_data = scales == nullptr ? nullptr : scales->Data<T1>();
const auto* bias_data = bias == nullptr ? nullptr : bias->Data<T1>();

IAllocatorUniquePtr<std::byte> workspace{};
const size_t workspace_size = MlasQNBitGemmBatchWorkspaceSize(
M, N, K, batch_count, nbits_, block_size_, zero_points, compute_type_, &mlas_backend_kernel_selector_config_);
Expand All @@ -542,7 +694,7 @@
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
if (compute_type_ == SQNBIT_CompInt8 || (compute_type_ == HQNBIT_CompInt8 && nbits_ == 8)) {
if (compute_type_ == SQNBIT_CompInt8) {
data[i].QuantBDataWorkspace = packed_b_.get();
}
data[i].PackedQuantBData = static_cast<std::byte*>(packed_b_.get());
Expand Down
Loading
Loading