saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-12 17:57:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

mlas nbit matmul requires packed_b (#20482)

Browse source 
### Description

mlas matmul nbits implementation requires packed b. have a condition for
this.

need to update this logic if it changes.


### Motivation and Context

---------

Signed-off-by: Liqun Fu <liqfu@microsoft.com>
This commit is contained in:
liqun Fu 2024-04-27 00:18:53 +00:00 committed by GitHub
parent 619ceeed9e
commit 2f5fe4500d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc
View file

@ -322,7 +322,8 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const {
if (has_single_b_matrix) {
const auto compute_type = static_cast<MLAS_SQNBIT_GEMM_COMPUTE_TYPE>(accuracy_level_);
if (MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type)) {
// mlas nbits implementation requires packed b. update this logic if it changes.
if (MlasIsSQNBitGemmAvailable(nbits_, block_size_, compute_type) && packed_b_) {
IAllocatorUniquePtr<std::byte> workspace{};
if (const size_t workspace_size = MlasSQNBitGemmBatchWorkspaceSize(M, N, K, batch_count,
nbits_, block_size_, compute_type);
@ -332,20 +333,11 @@ Status MatMulNBits::Compute(OpKernelContext* ctx) const {
workspace = IAllocator::MakeUniquePtr<std::byte>(allocator, workspace_size);
}
const void* b_data = [&]() -> const void* {
if (packed_b_) {
return packed_b_.get();
}
const Tensor* b = ctx->Input<Tensor>(1);
return b->DataRaw();
}();
InlinedVector<MLAS_SQNBIT_GEMM_DATA_PARAMS> data(batch_count);
for (size_t i = 0; i < batch_count; ++i) {
data[i].A = a_data + helper.LeftOffsets()[i];
data[i].lda = lda;
data[i].QuantBData = b_data;
data[i].QuantBData = packed_b_.get();
data[i].QuantBScale = scales_data;
data[i].QuantBZeroPoint = zero_points_data;
data[i].C = y_data + helper.OutputOffsets()[i];