A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
[Feat] dnnl build for AVX2 W8A8 Int8 (#41318)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Tianmu Li committed
May 6, 2026 at 00:28 UTC
e87e09a50a41db3ddd1cc0403e8d1ebe0b0f522c
4 files changed
+146
-20
cmake/cpu_extension.cmake
+10
-11
@@ -326,14 +326,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
326
set(ONEDNN_VERBOSE "ON")
327
set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
328
329
- # TODO: Refactor this
330
- if (ENABLE_X86_ISA)
331
- # Note: only enable oneDNN for AVX512
332
- list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
333
- else()
334
- list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
335
- endif()
336
-
329
set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
330
set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
331
FetchContent_MakeAvailable(oneDNN)
@@ -346,8 +338,14 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
338
PRIVATE ${oneDNN_SOURCE_DIR}/src
339
)
340
target_link_libraries(dnnl_ext dnnl torch)
349
- target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
341
+ if (ENABLE_X86_ISA)
342
+ target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS_AVX2} -fPIC)
343
+ else()
344
+ target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
345
+ endif()
346
list(APPEND LIBS dnnl_ext)
347
+
348
+
349
set(USE_ONEDNN ON)
350
else()
351
set(USE_ONEDNN OFF)
@@ -435,10 +433,11 @@ if (ENABLE_X86_ISA)
433
"csrc/cpu/pos_encoding.cpp"
434
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
435
438
- set(VLLM_EXT_SRC_AVX2
436
+ set(VLLM_EXT_SRC_AVX2
437
"csrc/cpu/utils.cpp"
438
"csrc/cpu/spec_decode_utils.cpp"
439
"csrc/cpu/cpu_attn.cpp"
440
+ "csrc/cpu/dnnl_kernels.cpp"
441
"csrc/cpu/torch_bindings.cpp"
442
# TODO: Remove these files
443
"csrc/cpu/activation.cpp"
@@ -453,7 +452,7 @@ if (ENABLE_X86_ISA)
452
453
set(_C_LIBS numa dnnl_ext)
454
set(_C_AVX512_LIBS numa dnnl_ext)
456
- set(_C_AVX2_LIBS numa)
455
+ set(_C_AVX2_LIBS numa dnnl_ext)
456
457
# AMX + AVX512F + AVX512BF16 + AVX512VNNI
458
define_extension_target(
csrc/cpu/cpu_types_x86.hpp
+133
-6
@@ -122,9 +122,17 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
122
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
123
124
void save(void* ptr, const int elem_num) const {
125
+#ifdef __AVX512BW__
126
constexpr uint32_t M = 0xFFFFFFFF;
127
__mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
128
_mm256_mask_storeu_epi16(ptr, mask, reg);
129
+#else
130
+ // Fallback for lack of 16-bit masked store
131
+ int16_t tmp[VEC_ELEM_NUM];
132
+ _mm256_storeu_si256((__m256i*)tmp, reg);
133
+ for (int i = 0; i < elem_num; ++i)
134
+ reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
135
+#endif
136
}
137
};
138
@@ -161,9 +169,17 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
169
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
170
171
void save(void* ptr, const int elem_num) const {
172
+#ifdef __AVX512BW__
173
constexpr uint32_t M = 0xFFFFFFFF;
174
__mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
175
_mm256_mask_storeu_epi16(ptr, mask, reg);
176
+#else
177
+ // Fallback for lack of 16-bit masked store
178
+ int16_t tmp[VEC_ELEM_NUM];
179
+ _mm256_storeu_si256((__m256i*)tmp, reg);
180
+ for (int i = 0; i < elem_num; ++i)
181
+ reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
182
+#endif
183
}
184
};
185
@@ -247,13 +263,12 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
263
explicit BF16Vec32(__m256i low, __m256i high)
264
: reg_low(low), reg_high(high) {}
265
266
+ explicit BF16Vec32()
267
+ : reg_low(_mm256_setzero_si256()), reg_high(_mm256_setzero_si256()) {}
268
+
269
explicit BF16Vec32(BF16Vec8& vec8_data)
251
- : reg_low((__m256i)_mm256_inserti32x4(
252
- _mm256_castsi128_si256((__m128i)vec8_data.reg),
253
- (__m128i)vec8_data.reg, 1)),
254
- reg_high((__m256i)_mm256_inserti32x4(
255
- _mm256_castsi128_si256((__m128i)vec8_data.reg),
256
- (__m128i)vec8_data.reg, 1)) {}
270
+ : reg_low(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)),
271
+ reg_high(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)) {}
272
273
// E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
274
// above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
@@ -674,6 +689,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
689
_mm256_sub_ps(reg_high, b.reg_high));
690
}
691
692
+ FP32Vec16 operator-() const {
693
+ const __m256 neg = _mm256_set1_ps(-0.0f);
694
+ return FP32Vec16(_mm256_xor_ps(reg_low, neg), _mm256_xor_ps(reg_high, neg));
695
+ }
696
+
697
FP32Vec16 operator/(const FP32Vec16& b) const {
698
return FP32Vec16(_mm256_div_ps(reg_low, b.reg_low),
699
_mm256_div_ps(reg_high, b.reg_high));
@@ -739,6 +759,85 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
759
_mm256_storeu_ps(ptr, reg_low);
760
_mm256_storeu_ps(ptr + 8, reg_high);
761
}
762
+
763
+ void save(float* ptr, const int elem_num) const {
764
+ // Partial store: cmpgt produces a sign-bit mask (0xFFFFFFFF/0 per lane)
765
+ // for the first elem_num lanes, applied across the two 8-wide halves.
766
+ if (elem_num <= 8) {
767
+ __m256i mask =
768
+ _mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num),
769
+ _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
770
+ _mm256_maskstore_ps(ptr, mask, reg_low);
771
+ } else {
772
+ _mm256_storeu_ps(ptr, reg_low);
773
+ __m256i mask =
774
+ _mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num - 8),
775
+ _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
776
+ _mm256_maskstore_ps(ptr + 8, mask, reg_high);
777
+ }
778
+ }
779
+
780
+ FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
781
+ return FP32Vec16(
782
+ _mm256_min_ps(max.reg_low, _mm256_max_ps(min.reg_low, reg_low)),
783
+ _mm256_min_ps(max.reg_high, _mm256_max_ps(min.reg_high, reg_high)));
784
+ }
785
+
786
+ FP32Vec16 abs() const {
787
+ const __m256 sign_mask = _mm256_set1_ps(-0.0f);
788
+ return FP32Vec16(_mm256_andnot_ps(sign_mask, reg_low),
789
+ _mm256_andnot_ps(sign_mask, reg_high));
790
+ }
791
+
792
+ FP32Vec16 min(const FP32Vec16& b) const {
793
+ return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
794
+ _mm256_min_ps(reg_high, b.reg_high));
795
+ }
796
+
797
+ // Partial element-wise min over the first elem_num lanes only (tail path).
798
+ // Scalar via AliasReg: AVX2 has no masked vminps, so we spill, loop, reload.
799
+ FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
800
+ AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
801
+ ar_this_low.reg = reg_low;
802
+ ar_this_high.reg = reg_high;
803
+ ar_b_low.reg = b.reg_low;
804
+ ar_b_high.reg = b.reg_high;
805
+ for (int i = 0; i < elem_num && i < 8; ++i)
806
+ ar_this_low.values[i] =
807
+ std::min(ar_this_low.values[i], ar_b_low.values[i]);
808
+ for (int i = 0; i < elem_num - 8 && i < 8; ++i)
809
+ ar_this_high.values[i] =
810
+ std::min(ar_this_high.values[i], ar_b_high.values[i]);
811
+ return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
812
+ }
813
+
814
+ // Partial element-wise max over the first elem_num lanes only (tail path).
815
+ // Scalar via AliasReg: AVX2 has no masked vmaxps, so we spill, loop, reload.
816
+ FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
817
+ AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
818
+ ar_this_low.reg = reg_low;
819
+ ar_this_high.reg = reg_high;
820
+ ar_b_low.reg = b.reg_low;
821
+ ar_b_high.reg = b.reg_high;
822
+ for (int i = 0; i < elem_num && i < 8; ++i)
823
+ ar_this_low.values[i] =
824
+ std::max(ar_this_low.values[i], ar_b_low.values[i]);
825
+ for (int i = 0; i < elem_num - 8 && i < 8; ++i)
826
+ ar_this_high.values[i] =
827
+ std::max(ar_this_high.values[i], ar_b_high.values[i]);
828
+ return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
829
+ }
830
+
831
+ float reduce_min() const {
832
+ __m256 v = _mm256_min_ps(reg_low, reg_high);
833
+ __m256 v_shuffled = _mm256_permute_ps(v, 0b00001011);
834
+ __m256 v_min = _mm256_min_ps(v, v_shuffled);
835
+ v_shuffled = _mm256_permute_ps(v_min, 0b00000001);
836
+ v_min = _mm256_min_ps(v_min, v_shuffled);
837
+ v_shuffled = _mm256_permute2f128_ps(v_min, v_min, 0b00000001);
838
+ v_min = _mm256_min_ps(v_min, v_shuffled);
839
+ return _mm256_cvtss_f32(v_min);
840
+ }
841
};
842
#endif
843
@@ -791,6 +890,34 @@ struct INT8Vec64 : public Vec<INT8Vec64> {
890
// non-temporal save
891
void nt_save(int8_t* ptr) { _mm512_stream_si512((__m512i*)ptr, reg); }
892
};
893
+#else
894
+struct INT8Vec16 : public Vec<INT8Vec16> {
895
+ constexpr static int VEC_ELEM_NUM = 16;
896
+ union AliasReg {
897
+ __m128i reg;
898
+ int8_t values[VEC_ELEM_NUM];
899
+ };
900
+
901
+ __m128i reg;
902
+
903
+ explicit INT8Vec16(const FP32Vec16& vec) {
904
+ __m256i lo_i32 = _mm256_cvtps_epi32(vec.reg_low);
905
+ __m256i hi_i32 = _mm256_cvtps_epi32(vec.reg_high);
906
+ __m256i packed16 = _mm256_packs_epi32(lo_i32, hi_i32);
907
+ packed16 = _mm256_permute4x64_epi64(packed16, 0xD8);
908
+ __m256i packed8 = _mm256_packs_epi16(packed16, _mm256_setzero_si256());
909
+ packed8 = _mm256_permute4x64_epi64(packed8, 0xD8);
910
+ reg = _mm256_castsi256_si128(packed8);
911
+ }
912
+
913
+ void save(int8_t* ptr) const { _mm_storeu_si128((__m128i*)ptr, reg); }
914
+
915
+ void save(int8_t* ptr, const int elem_num) const {
916
+ AliasReg ar;
917
+ ar.reg = reg;
918
+ for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i];
919
+ }
920
+};
921
#endif
922
923
template <typename T>
csrc/cpu/dnnl_kernels.cpp
+1
-1
@@ -215,7 +215,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
215
float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
216
token_zp_scale_vec = cvt_vec_t(zp_scale_val);
217
}
218
- for (; j < hidden_size - vec_elem_num; ++j) {
218
+ for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
219
cvt_vec_t elems_fp32(input_ptr + j);
220
elems_fp32 = elems_fp32 * token_scale_vec;
221
if constexpr (AZP) {
csrc/cpu/torch_bindings.cpp
+2
-2
@@ -278,8 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
278
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
279
280
// Quantization
281
-#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
282
- defined(__powerpc64__)
281
+#if defined(__AVX512F__) || defined(__AVX2__) || \
282
+ (defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
283
// Helper function to release oneDNN handlers
284
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
285
&release_dnnl_matmul_handler);