Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
162 changes: 118 additions & 44 deletions onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

Large diffs are not rendered by default.

79 changes: 75 additions & 4 deletions onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand Down
28 changes: 26 additions & 2 deletions onnxruntime/core/mlas/lib/kai_ukernel_interface.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -26,14 +28,30 @@ struct KaiMatmulKernel {
UkernelFn ukernel;
};

enum class KaiQ4RhsPackLayout {
SymmetricNxK,
AsymmetricNxK,
AsymmetricNxKInterleavedNrx4,
};

template <typename UkernelFn>
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<kai_matmul_clamp_f32_f32p_f32p_ukernel>;

// Wrapper for FP32 kernels used for GEMV-style workloads (typically a single-row/skinny-M use case).
using KaiF32SgemvKernel = KaiMatmulKernel<kai_matmul_clamp_f32_f32_f32p_ukernel>;

// Wrapper for Qnbit GEMM kernels producing FP32 output.
using KaiQnbitGemmKernel = KaiMatmulKernel<kai_matmul_clamp_f32_qai8dxp_qsi4c32p_ukernel>;
using KaiQnbitGemmKernel = KaiQ4MatmulKernel<kai_matmul_clamp_f32_qai8dxp_qsi4c32p_ukernel>;

// Wrapper for Qnbit Asymmetric-quantized GEMM kernels producing FP32 output.
using KaiQnbitAsymGemmKernel = KaiQ4MatmulKernel<kai_matmul_clamp_f32_qsi8d32p_qai4c32p_ukernel>;

// Wrapper for dynamic-quantized GEMM kernels producing FP32 output.
using KaiDynamicQGemmKernel = KaiMatmulKernel<kai_matmul_clamp_f32_qai8dxp_qsi8cxp_ukernel>;
Expand All @@ -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();

Expand All @@ -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();
const KaiBF16SBgemmKernel& GetKleidiAISBGemmUKernel();
73 changes: 55 additions & 18 deletions onnxruntime/core/mlas/lib/qnbitgemm.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<const std::byte*>(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;
Expand Down Expand Up @@ -1128,7 +1159,8 @@ InitializeWorkspace_CompInt8<float>(
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(
Expand All @@ -1140,11 +1172,12 @@ InitializeWorkspace_CompInt8<float>(

const float* ARowPtr = data.A;
std::byte* QuantARowPtr = static_cast<std::byte*>(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);
Expand Down Expand Up @@ -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);

Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand All @@ -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
);
}
}
Expand Down
Loading
Loading