Skip to content

[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels #29619

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

Open
blazingphoenix7 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels #29619

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
359378a
[MLAS] Add AVX2 (+VNNI) 2-bit weight CPU kernels
blazingphoenix7 Jul 8, 2026
2d86379
Build the W2 pack helper TU without AVX-512 flags
blazingphoenix7 Jul 9, 2026
4ab49f5
Review feedback: explicit include, shared effective-K helper, comments
blazingphoenix7 Jul 10, 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
12 changes: 6 additions & 6 deletions cmake/onnxruntime_mlas.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -871,13 +871,13 @@ endif()
)
set_source_files_properties(${mlas_platform_srcs_avx512core} PROPERTIES COMPILE_FLAGS "-mfma -mavx512vnni -mavx512bw -mavx512dq -mavx512vl")

# Strip -mavx512vnni from the W2 scalar oracle / pack-helper TU so the
# compiler cannot autovectorize its int8 dot-product loops to vpdpbusd.
# Needed because this helper runs at model load on AVX-512-only
# (non-VNNI) hosts via the AVX-512 W2 dispatch. TU is pure C++ -- no
# AVX-512 intrinsics inside.
# Build the W2 scalar oracle / pack-helper TU without any AVX-512

@hariharans29 hariharans29 Jul 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this change needed ? The previous state had it to just strip away avx512 vnni flags so that auto-vectorization (if the compiler fdoes do it) didn't end up using avx512 vnni but we can use avx512 ? Or does the code there get dispatched on avx2 only hosts with this change and hence necessaitates avx512 flags strip off ?

@blazingphoenix7 blazingphoenix7 Jul 9, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, the latter. With this PR the AVX2 and AVX2-VNNI dispatch tables point at the pack helpers in this TU, so it now runs at model load on hosts with no AVX-512 at all. The first CI round caught exactly that: the three Linux GPU pipelines run on AVX2-only CPUs, and onnxruntime_provider_test died with an illegal instruction there, since with -mavx512bw/dq/vl still on the file the compiler is free to autovectorize the scalar pack loops with EVEX encodings. It is the same failure mode the existing -mavx512vnni strip was guarding against, one ISA level down.

The TU is intrinsic-free scalar and the pack is one-time load work, so baseline codegen should not cost anything measurable, and ARM64 already builds this same file with no x86 flags. If you would rather keep avx512 codegen for the avx512 dispatch path, I can give the avx2 side its own copy in a separate TU instead, stripping the flags on the shared file just seemed like the smaller change. The second commit should turn those three jobs green once CI reruns.

# flags so the compiler cannot autovectorize its scalar loops into
# EVEX instructions. This helper runs at model load on AVX-512-only
# (non-VNNI) hosts via the AVX-512 W2 dispatch and on AVX2-only hosts
# via the AVX2 W2 dispatch. TU is pure C++ -- no intrinsics inside.
set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512_2bit.cpp PROPERTIES
COMPILE_FLAGS "-mfma -mavx512bw -mavx512dq -mavx512vl")
COMPILE_FLAGS "")

set(mlas_platform_srcs_avx512vnni
${MLAS_SRC_DIR}/sqnbitgemm_kernel_avx512vnni.cpp
Expand Down
104 changes: 104 additions & 0 deletions onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ Module Name:
#include "sqnbitgemm_kernel_avx2_int8_blklen32.h"
#include "sqnbitgemm_kernel_avx2_int8_blklen64.h"

#include "sqnbitgemm_kernel_avx2_2bit.h"
#include "sqnbitgemm_kernel_avx2_2bit_blklen32.h"
#include "sqnbitgemm_kernel_avx2_2bit_blklen64.h"
#include "sqnbitgemm_kernel_avx2_2bit_blklen128.h"

Comment on lines +30 to +34

@blazingphoenix7 blazingphoenix7 Jul 10, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the explicit include in 4ab49f5.

#include "sqnbitgemm_m1_sym_kernel_avx2_int8_blklen32.h"
#include "sqnbitgemm_m1_sym_kernel_avx2_int8_blklen64.h"

Expand Down Expand Up @@ -1439,6 +1444,80 @@ SQ8BitGemmPackQuantBDataAndBlkSum(
HasZeroPoint, QuantBZPBegin, PackedQuantB, ThreadPool);
}

//
// BlkLen-routing wrappers for the W2 CompInt8 AVX2 / AVX2-VNNI dispatch
// entries. Production code reaches these via the MLAS dispatch table; tests
// call them directly via the namespace. Siblings of the AVX-512 variants in
// sqnbitgemm_kernel_avx512.cpp / sqnbitgemm_kernel_avx512vnni.cpp.
//
namespace onnxruntime::mlas::sq2bit_avx2 {
size_t MLASCALL
SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch(
size_t BlkLen,
const std::byte* QuantA,
const float* QuantAScale,
const std::byte* QuantBData,
const float* QuantBScale,
const std::byte* QuantBZeroPoint,
float* C,
size_t CountM,
size_t CountN,
size_t CountK,
size_t BlockCountK,
const float* Bias,
size_t ldc,
const float* ABlockSum,
const float* QuantBBlkSum)
{
if (BlkLen == 128) {
return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen128_Avx2(
QuantA, QuantAScale, QuantBData, QuantBScale,
C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}
if (BlkLen == 32) {
return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen32_Avx2(
QuantA, QuantAScale, QuantBData, QuantBScale,
C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}
return SQ2BitGemmKernel_BlkSum_CompInt8_Avx2(
BlkLen, QuantA, QuantAScale, QuantBData, QuantBScale, QuantBZeroPoint,
C, CountM, CountN, CountK, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}

size_t MLASCALL
SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch(
size_t BlkLen,
const std::byte* QuantA,
const float* QuantAScale,
const std::byte* QuantBData,
const float* QuantBScale,
const std::byte* QuantBZeroPoint,
float* C,
size_t CountM,
size_t CountN,
size_t CountK,
size_t BlockCountK,
const float* Bias,
size_t ldc,
const float* ABlockSum,
const float* QuantBBlkSum)
{
if (BlkLen == 128) {
return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen128_Avx2Vnni(
QuantA, QuantAScale, QuantBData, QuantBScale,
C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}
if (BlkLen == 32) {
return SQ2BitGemmKernel_BlkSum_CompInt8_BlkLen32_Avx2Vnni(
QuantA, QuantAScale, QuantBData, QuantBScale,
C, CountM, CountN, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}
return SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni(
BlkLen, QuantA, QuantAScale, QuantBData, QuantBScale, QuantBZeroPoint,
C, CountM, CountN, CountK, BlockCountK, Bias, ldc, ABlockSum, QuantBBlkSum);
}
} // namespace onnxruntime::mlas::sq2bit_avx2

//
// Kernel dispatch structure definition.
//
Expand All @@ -1461,6 +1540,20 @@ const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2 = []() {
d.SQ8BitGemmKernel_BlkSum_CompInt8 = SQ8BitGemmKernel_BlkSum_CompInt8_avx2<false>;
d.QuantizeARowComputeBlkSum_CompInt8 = QuantizeARow_CompInt8_avx2;

// 2-bit native CompInt8 path. Reuses the portable block-group packer and
// effective-K helper (shared with AVX-512); the AVX2 kernel supplies the
// 256-bit compute. QuantizeARow above already provides the signed-int8 A
// layout the W2 path consumes.
static_assert(
onnxruntime::mlas::sq2bit_avx512::kBlockGroupBlks == kSq2BitAvx512WeightKBlockGroup,
"kBlockGroupBlks (kernel-internal) must match kSq2BitAvx512WeightKBlockGroup (qnbitgemm.h).");
d.Q2BitGemmPackQuantBDataSize = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512;
d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar;
d.SQ2BitGemmKernel_BlkSum_CompInt8 = onnxruntime::mlas::sq2bit_avx2::SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch;
d.Q2BitGemmEffectiveBlockCountK = [](size_t BlockCountK) {
return MlasDivRoundup(BlockCountK, kSq2BitAvx512WeightKBlockGroup) * kSq2BitAvx512WeightKBlockGroup;
};

return d;
}();

Expand All @@ -1483,5 +1576,16 @@ const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2vnni = []() {
d.SQ8BitGemmKernel_BlkSum_CompInt8 = SQ8BitGemmKernel_BlkSum_CompInt8_avx2<true>;
d.QuantizeARowComputeBlkSum_CompInt8 = QuantizeARow_CompInt8_avx2;

// 2-bit native CompInt8 path (AVX-VNNI compute). See the AVX2 table above.
static_assert(
onnxruntime::mlas::sq2bit_avx512::kBlockGroupBlks == kSq2BitAvx512WeightKBlockGroup,
"kBlockGroupBlks (kernel-internal) must match kSq2BitAvx512WeightKBlockGroup (qnbitgemm.h).");
d.Q2BitGemmPackQuantBDataSize = onnxruntime::mlas::sq2bit_avx512::Q2BitGemmPackQuantBDataSize_Avx512;
d.SQ2BitGemmPackQuantBDataAndBlkSum = onnxruntime::mlas::sq2bit_avx512::SQ2BitGemmPackQuantBDataAndBlkSum_Scalar;
d.SQ2BitGemmKernel_BlkSum_CompInt8 = onnxruntime::mlas::sq2bit_avx2::SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch;
d.Q2BitGemmEffectiveBlockCountK = [](size_t BlockCountK) {
return MlasDivRoundup(BlockCountK, kSq2BitAvx512WeightKBlockGroup) * kSq2BitAvx512WeightKBlockGroup;
};

return d;
}();
72 changes: 72 additions & 0 deletions onnxruntime/core/mlas/lib/sqnbitgemm_kernel_avx2_2bit.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
/*++

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Module Name:

sqnbitgemm_kernel_avx2_2bit.h

Abstract:

Declarations for the AVX2 / AVX2-VNNI W2 CompInt8 BlkLen-routing dispatch
forwarders. The forwarders are defined in sqnbitgemm_kernel_avx2.cpp and the
per-BlkLen kernels live in sqnbitgemm_kernel_avx2_2bit_blklen{32,64,128}.h.
Declared here so the dispatch tables and the unit tests can reference them,
mirroring the AVX-512 forwarder declarations in
sqnbitgemm_kernel_avx512_2bit.h.

--*/

#pragma once

#include <cstddef>

#include "mlas.h"

namespace onnxruntime {
namespace mlas {
namespace sq2bit_avx2 {

size_t MLASCALL
SQ2BitGemmKernel_BlkSum_CompInt8_Avx2_Dispatch(
size_t BlkLen,
const std::byte* QuantA,
const float* QuantAScale,
const std::byte* QuantBData,
const float* QuantBScale,
const std::byte* QuantBZeroPoint,
float* C,
size_t CountM,
size_t CountN,
size_t CountK,
size_t BlockCountK,
const float* Bias,
size_t ldc,
const float* ABlockSum,
const float* QuantBBlkSum
);

size_t MLASCALL
SQ2BitGemmKernel_BlkSum_CompInt8_Avx2Vnni_Dispatch(
size_t BlkLen,
const std::byte* QuantA,
const float* QuantAScale,
const std::byte* QuantBData,
const float* QuantBScale,
const std::byte* QuantBZeroPoint,
float* C,
size_t CountM,
size_t CountN,
size_t CountK,
size_t BlockCountK,
const float* Bias,
size_t ldc,
const float* ABlockSum,
const float* QuantBBlkSum
);

} // namespace sq2bit_avx2
} // namespace mlas
} // namespace onnxruntime
Loading
Loading