@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 20852c8f4c
Fork 0

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 13, 00:45 UTC

[CPU] Refactor CPU WNA16 (#28826)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Nov 19, 2025 at 10:32 UTC 20852c8f4c10d80204c47e0cb85f5b252ff51c86
22 files changed +1656 -78
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+5 -6
@@ -73,12 +73,11 @@ function cpu_tests() {
73 pytest -x -s -v \
74 tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
75
76 - # Note: disable it until supports V1
77 - # Run AWQ test
78 - # docker exec cpu-test-"$NUMA_NODE" bash -c "
79 - # set -e
80 - # pytest -x -s -v \
81 - # tests/quantization/test_ipex_quant.py"
76 + # Run AWQ/GPTQ test
77 + docker exec cpu-test-"$NUMA_NODE" bash -c "
78 + set -e
79 + pytest -x -s -v \
80 + tests/quantization/test_cpu_wna16.py"
81
82 # Run multi-lora tests
83 docker exec cpu-test-"$NUMA_NODE" bash -c "
cmake/cpu_extension.cmake
+1
@@ -375,6 +375,7 @@ set(VLLM_EXT_SRC
375 if (AVX512_FOUND AND NOT AVX512_DISABLED)
376 set(VLLM_EXT_SRC
377 "csrc/cpu/shm.cpp"
378 + "csrc/cpu/cpu_wna16.cpp"
379 ${VLLM_EXT_SRC})
380 if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
381 set(VLLM_EXT_SRC
csrc/cpu/cpu_attn_impl.hpp
+1 -1
@@ -1,7 +1,6 @@
1 #ifndef CPU_ATTN_HPP
2 #define CPU_ATTN_HPP
3
4 -#include <unistd.h>
4 #include <type_traits>
5 #include <cstddef>
6
@@ -12,6 +11,7 @@
11 #include "cpu_types.hpp"
12 #include "scratchpad_manager.h"
13 #include "cpu_attn_macros.h"
14 +#include "utils.hpp"
15
16 namespace cpu_attention {
17 enum class ISA { AMX, VEC, VEC16 };
csrc/cpu/cpu_types_x86.hpp
+39 -8
@@ -104,6 +104,8 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
104 explicit FP16Vec16(bool, void* ptr)
105 : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
106
107 + explicit FP16Vec16(const c10::Half v) : reg(_mm256_set1_epi16(v.x)) {}
108 +
109 explicit FP16Vec16(const FP32Vec16&);
110
111 void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -141,6 +143,8 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
143 explicit BF16Vec16(bool, void* ptr)
144 : reg(_mm256_stream_load_si256((__m256i*)ptr)) {}
145
146 + explicit BF16Vec16(const c10::BFloat16 v) : reg(_mm256_set1_epi16(v.x)) {}
147 +
148 explicit BF16Vec16(const FP32Vec16&);
149
150 void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
@@ -350,6 +354,22 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
354
355 explicit FP32Vec16(__m512 data) : reg(data) {}
356
357 + // de-pack 4 bit values
358 + explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
359 + int64_t mask_0 = 0x0F0F0F0F0F0F0F0F;
360 + int64_t mask_1 = 0xF0F0F0F0F0F0F0F0;
361 + int64_t value_0 = value & mask_0;
362 + int64_t value_1 = value & mask_1;
363 + __m128i vec_0 = _mm_movpi64_epi64((__m64)value_0);
364 + __m128i vec_1 = _mm_movpi64_epi64((__m64)value_1);
365 + vec_0 = _mm_cvtepu8_epi16(vec_0);
366 + vec_1 = _mm_cvtepu8_epi16(vec_1);
367 + vec_1 = _mm_slli_epi16(vec_1, 4);
368 + __m128i vec = _mm_or_si128(vec_0, vec_1);
369 + __m512i vec_i32 = _mm512_cvtepu8_epi32(vec);
370 + reg = _mm512_permutexvar_ps(vec_i32, lut.reg);
371 + }
372 +
373 explicit FP32Vec16(const FP32Vec4& data)
374 : reg((__m512)_mm512_inserti32x4(
375 _mm512_inserti32x4(
@@ -426,14 +446,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
446
447 float get_last_elem() const { return _mm512_cvtss_f32(reg); }
448
429 - template <int group_size>
430 - float reduce_sub_sum(int idx) {
431 - static_assert(VEC_ELEM_NUM % group_size == 0);
432 - constexpr uint32_t base_mask = (0xFFFF >> (16 - group_size));
433 - __mmask16 mask = _cvtu32_mask16(base_mask << (idx * group_size));
434 - return _mm512_mask_reduce_add_ps(mask, reg);
435 - }
436 -
449 void save(float* ptr) const { _mm512_storeu_ps(ptr, reg); }
450
451 void save(float* ptr, const int elem_num) const {
@@ -755,6 +767,25 @@ inline void non_temporal_save(BF16Vec16& vec, void* ptr) {
767 inline void non_temporal_save(FP32Vec16& vec, void* ptr) {
768 _mm512_stream_ps((float*)ptr, vec.reg);
769 }
770 +
771 +static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
772 + void* ptr) {
773 + __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
774 + __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
775 + vec_1 = _mm512_slli_epi32(vec_1, 16);
776 + vec_0 = _mm512_or_si512(vec_0, vec_1);
777 + _mm512_storeu_epi32(ptr, vec_0);
778 +}
779 +
780 +static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
781 + void* ptr) {
782 + __m512i vec_0 = _mm512_cvtepu16_epi32(vec0.reg);
783 + __m512i vec_1 = _mm512_cvtepu16_epi32(vec1.reg);
784 + vec_1 = _mm512_slli_epi32(vec_1, 16);
785 + vec_0 = _mm512_or_si512(vec_0, vec_1);
786 + _mm512_storeu_epi32(ptr, vec_0);
787 +}
788 +
789 #endif
790
791 inline void mem_barrier() { _mm_mfence(); }
csrc/cpu/cpu_wna16.cpp new
+402
@@ -0,0 +1,402 @@
1 +#include "cpu_types.hpp"
2 +#include "scratchpad_manager.h"
3 +#include "utils.hpp"
4 +
5 +#ifdef CPU_CAPABILITY_AMXBF16
6 + #include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
7 +#endif
8 +#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
9 +
10 +#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
11 + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
12 + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
13 +
14 +#define VLLM_DISPATCH_16B_TYPES(TYPE, NAME, ...) \
15 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_16B_TYPES(__VA_ARGS__))
16 +
17 +template <typename T>
18 +void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
19 + int32_t stride) {
20 + std::stringstream ss;
21 + ss << std::fixed << std::setprecision(5) << name << ": [\n";
22 + auto* curr_logits_buffer = ptr;
23 + for (int32_t m = 0; m < row; ++m) {
24 + for (int32_t n = 0; n < col; ++n) {
25 + ss << curr_logits_buffer[n] << ", ";
26 + }
27 + ss << "\n";
28 + curr_logits_buffer += stride;
29 + }
30 + ss << "]\n";
31 + std::printf("%s", ss.str().c_str());
32 +}
33 +
34 +namespace {
35 +using cpu_utils::ISA;
36 +using cpu_utils::VecTypeTrait;
37 +
38 +template <typename scalar_t, ISA isa, bool has_zp, bool use_desc_act>
39 +class Dequantizer4b {
40 + public:
41 + constexpr static int32_t pack_num = 32 / 4;
42 + using scalar_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
43 +
44 + public:
45 + static void dequant(int32_t* __restrict__ q_weight,
46 + scalar_t* __restrict__ weight,
47 + scalar_t* __restrict__ scales,
48 + int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
49 + const int64_t scales_stride, const int64_t zeros_stride,
50 + const int32_t k_size, const int32_t group_size) {
51 + vec_op::FP32Vec16 lut;
52 + if constexpr (has_zp) {
53 + // AWQ
54 + alignas(64) static const float LUT[16] = {
55 + 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f,
56 + 8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f};
57 + lut = vec_op::FP32Vec16(LUT);
58 + } else {
59 + // GPTQ
60 + alignas(64) static const float LUT[16] = {
61 + -8.0f, -7.0f, -6.0f, -5.0f, -4.0f, -3.0f, -2.0f, -1.0f,
62 + 0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
63 + lut = vec_op::FP32Vec16(LUT);
64 + }
65 +
66 + // per 64-bits elem contains 16 output channels
67 + int64_t* __restrict__ curr_q_weight = reinterpret_cast<int64_t*>(q_weight);
68 + int64_t* __restrict__ curr_zeros = reinterpret_cast<int64_t*>(zeros);
69 + scalar_t* __restrict__ curr_weight = weight;
70 + scalar_t* __restrict__ curr_scale = scales;
71 + vec_op::FP32Vec16 scale_0;
72 + vec_op::FP32Vec16 scale_1;
73 + vec_op::FP32Vec16 zero_0;
74 + vec_op::FP32Vec16 zero_1;
75 + int32_t group_counter = 0;
76 + for (int32_t k_idx = 0; k_idx < k_size; k_idx += 2) {
77 + int64_t qwb_0 = *curr_q_weight;
78 + int64_t qwb_1 = *(curr_q_weight + 1);
79 + vec_op::FP32Vec16 wb_0(qwb_0, lut);
80 + vec_op::FP32Vec16 wb_1(qwb_1, lut);
81 +
82 + if constexpr (!use_desc_act) {
83 + if (group_counter == 0) {
84 + scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
85 + scale_1 = vec_op::FP32Vec16(scale_0);
86 + curr_scale += scales_stride;
87 +
88 + if constexpr (has_zp) {
89 + zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
90 + zero_1 = vec_op::FP32Vec16(zero_0);
91 + curr_zeros += zeros_stride / 2;
92 + }
93 + }
94 + } else {
95 + int32_t g_idx_0 = g_idx[k_idx];
96 + int32_t g_idx_1 = g_idx[k_idx + 1];
97 + scale_0 = vec_op::FP32Vec16(
98 + scalar_vec_t(curr_scale + g_idx_0 * scales_stride));
99 + scale_1 = vec_op::FP32Vec16(
100 + scalar_vec_t(curr_scale + g_idx_1 * scales_stride));
101 + if constexpr (has_zp) {
102 + zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2),
103 + lut);
104 + zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2),
105 + lut);
106 + }
107 + }
108 +
109 + if constexpr (has_zp) {
110 + wb_0 = wb_0 - zero_0;
111 + wb_1 = wb_1 - zero_1;
112 + }
113 +
114 + wb_0 = wb_0 * scale_0;
115 + wb_1 = wb_1 * scale_1;
116 +
117 + scalar_vec_t output_vec_0(wb_0);
118 + scalar_vec_t output_vec_1(wb_1);
119 +
120 + // AMX needs to interlave K elements to pack as 32 bits
121 + if constexpr (isa == ISA::AMX) {
122 + vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
123 + } else {
124 + output_vec_0.save(curr_weight);
125 + output_vec_1.save(curr_weight + 16);
126 + }
127 +
128 + // update
129 + curr_q_weight += 2;
130 + curr_weight += 32;
131 + if constexpr (!use_desc_act) {
132 + group_counter += 2;
133 + if (group_counter == group_size) {
134 + group_counter = 0;
135 + }
136 + }
137 + }
138 + }
139 +};
140 +}; // namespace
141 +
142 +template <typename scalar_t, typename dequantizer_t, typename gemm_t>
143 +void cpu_gemm_wna16_impl(
144 + scalar_t* __restrict__ input, int32_t* __restrict__ q_weight,
145 + scalar_t* __restrict__ output, scalar_t* __restrict__ scales,
146 + int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
147 + scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size,
148 + const int32_t k_size, const int64_t input_stride,
149 + const int64_t output_stride, const int64_t scales_group_stride,
150 + const int64_t zeros_group_stride, const int32_t group_num,
151 + const int32_t group_size, const int64_t pack_factor) {
152 + constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
153 + constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
154 + constexpr int32_t n_block_size = 16;
155 + static_assert(gemm_n_tile_size % n_block_size == 0);
156 + const int32_t thread_num = omp_get_max_threads();
157 +
158 + // a simple schedule policy, just to hold more B tiles in L2 and make sure
159 + // each thread has tasks
160 + const int32_t n_partition_size = [&]() {
161 + const int64_t cache_size = cpu_utils::get_l2_size();
162 + int64_t ps_cache_limit = cache_size / (k_size * sizeof(scalar_t));
163 + int64_t ps_thread_limit = n_size / thread_num;
164 + ps_cache_limit =
165 + std::max((ps_cache_limit / gemm_n_tile_size) * gemm_n_tile_size,
166 + (int64_t)gemm_n_tile_size);
167 + ps_thread_limit =
168 + std::max((ps_thread_limit / gemm_n_tile_size) * gemm_n_tile_size,
169 + (int64_t)gemm_n_tile_size);
170 + return std::min(ps_cache_limit, ps_thread_limit);
171 + }();
172 + const int32_t task_num = (n_size + n_partition_size - 1) / n_partition_size;
173 +
174 + // get buffer size
175 + const int64_t b_buffer_size =
176 + (((n_partition_size * k_size * sizeof(scalar_t) + 63) / 64) * 64);
177 + const int64_t c_buffer_size =
178 + (((gemm_m_tile_size * gemm_n_tile_size * sizeof(float) + 63) / 64) * 64);
179 + const int64_t b_buffer_offset = 0;
180 + const int64_t c_buffer_offset = b_buffer_size;
181 + const int64_t buffer_size = b_buffer_size + c_buffer_size;
182 + DNNLScratchPadManager::get_dnnl_scratchpad_manager()->realloc(buffer_size *
183 + thread_num);
184 +
185 + alignas(64) cpu_utils::Counter counter;
186 + cpu_utils::Counter* counter_ptr = &counter;
187 +
188 +#pragma omp parallel for schedule(static, 1)
189 + for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
190 + scalar_t* __restrict__ b_buffer = nullptr;
191 + float* __restrict__ c_buffer = nullptr;
192 + {
193 + uint8_t* buffer_ptr = DNNLScratchPadManager::get_dnnl_scratchpad_manager()
194 + ->get_data<uint8_t>() +
195 + thread_id * buffer_size;
196 + b_buffer = reinterpret_cast<scalar_t*>(buffer_ptr + b_buffer_offset);
197 + c_buffer = reinterpret_cast<float*>(buffer_ptr + c_buffer_offset);
198 + }
199 +
200 + const int64_t q_weight_block_stride = n_block_size / pack_factor * k_size;
201 + const int64_t b_buffer_block_stride = n_block_size * k_size;
202 + const int32_t zeros_block_stride = n_block_size / pack_factor;
203 +
204 + gemm_t gemm;
205 +
206 + for (;;) {
207 + int32_t task_id = counter_ptr->acquire_counter();
208 +
209 + if (task_id >= task_num) {
210 + break;
211 + }
212 +
213 + const int32_t n_start_idx = task_id * n_partition_size;
214 + const int32_t n_block_start_idx = n_start_idx / n_block_size;
215 + const int32_t n_num = std::min(n_partition_size, n_size - n_start_idx);
216 + const int32_t n_block_num = n_num / n_block_size;
217 + // std::printf("thread_id: %d, task_id: %d, n_start_idx: %d, n_num: %d\n",
218 + // thread_id, task_id, n_start_idx, n_num);
219 +
220 + // dequant weight
221 + {
222 + int32_t* __restrict__ curr_q_weight =
223 + q_weight + n_block_start_idx * q_weight_block_stride;
224 + scalar_t* __restrict__ curr_b_buffer = b_buffer;
225 + scalar_t* __restrict__ curr_scales = scales + n_start_idx;
226 + int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor;
227 + for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) {
228 + dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales,
229 + curr_zeros, g_idx, scales_group_stride,
230 + zeros_group_stride, k_size, group_size);
231 +
232 + // if (block_idx == 0 && n_start_idx == 0) {
233 + // print_logits("depacked weight", curr_b_buffer, k_size,
234 + // n_block_size, n_block_size);
235 + // }
236 +
237 + // update
238 + curr_q_weight += q_weight_block_stride;
239 + curr_b_buffer += b_buffer_block_stride;
240 + curr_scales += n_block_size;
241 + curr_zeros += zeros_block_stride;
242 + }
243 + }
244 +
245 + // compute loop
246 + {
247 + const int32_t n_tile_num = n_num / gemm_n_tile_size;
248 + scalar_t* __restrict__ curr_input = input;
249 + scalar_t* __restrict__ init_bias = bias;
250 + if (bias != nullptr) {
251 + init_bias += n_start_idx;
252 + }
253 + scalar_t* __restrict__ init_output = output + n_start_idx;
254 + for (int32_t m_idx = 0; m_idx < m_size; m_idx += gemm_m_tile_size) {
255 + const int32_t curr_m_size =
256 + std::min(gemm_m_tile_size, m_size - m_idx);
257 + scalar_t* __restrict__ curr_b_buffer = b_buffer;
258 + scalar_t* __restrict__ curr_bias = init_bias;
259 + scalar_t* __restrict__ curr_output = init_output;
260 + for (int32_t n_tile_idx = 0; n_tile_idx < n_tile_num; ++n_tile_idx) {
261 + gemm.gemm(curr_input, curr_b_buffer, c_buffer, curr_m_size, k_size,
262 + input_stride, b_buffer_block_stride, gemm_n_tile_size,
263 + false);
264 +
265 + if (bias != nullptr) {
266 + cpu_micro_gemm::bias_epilogue<gemm_n_tile_size>(
267 + c_buffer, curr_output, curr_bias, curr_m_size,
268 + gemm_n_tile_size, output_stride);
269 + curr_bias += gemm_n_tile_size;
270 + } else {
271 + cpu_micro_gemm::default_epilogue<gemm_n_tile_size>(
272 + c_buffer, curr_output, curr_m_size, gemm_n_tile_size,
273 + output_stride);
274 + }
275 +
276 + curr_b_buffer +=
277 + b_buffer_block_stride * (gemm_n_tile_size / n_block_size);
278 + curr_output += gemm_n_tile_size;
279 + }
280 + curr_input += gemm_m_tile_size * input_stride;
281 + init_output += gemm_m_tile_size * output_stride;
282 + }
283 + }
284 + }
285 + }
286 +}
287 +
288 +void cpu_gemm_wna16(
289 + const torch::Tensor& input, // [M, K]
290 + const torch::Tensor&
291 + q_weight, // [N / 16, K * 16 / pack_factor], packed as int32
292 + torch::Tensor& output, // [M, N]
293 + const torch::Tensor& scales, // [group_num, N]
294 + const std::optional<torch::Tensor>&
295 + zeros, // [group_num, N / pack_factor], packed as int32
296 + const std::optional<torch::Tensor>& g_idx, // [K]
297 + const std::optional<torch::Tensor>& bias, // [N]
298 + const int64_t pack_factor, const std::string& isa_hint) {
299 + using cpu_utils::ISA;
300 + TORCH_CHECK_EQ(pack_factor, 8); // only supports 4bits
301 + const int32_t a_m_size = input.size(0);
302 + const int32_t a_k_size = input.size(1);
303 + const int64_t a_m_stride = input.stride(0);
304 + const int32_t b_n_size = q_weight.size(0) * 16;
305 + TORCH_CHECK_EQ(a_k_size % 32, 0);
306 + TORCH_CHECK_EQ(b_n_size % 32, 0);
307 + const int32_t group_num = scales.size(0);
308 + const int32_t group_size = a_k_size / group_num;
309 + TORCH_CHECK_EQ(group_size % 2, 0);
310 + const int64_t scales_group_stride = scales.stride(0);
311 + const int64_t output_m_stride = output.stride(0);
312 +
313 + bool has_zp = zeros.has_value();
314 + bool use_desc_act = g_idx.has_value();
315 + TORCH_CHECK(!(has_zp && use_desc_act));
316 +
317 + ISA isa = [&]() {
318 + if (isa_hint == "amx") {
319 + return ISA::AMX;
320 + } else if (isa_hint == "vec") {
321 + return ISA::VEC;
322 + } else {
323 + TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
324 + }
325 + }();
326 +
327 + int32_t* zeros_ptr = has_zp ? zeros->data_ptr<int32_t>() : nullptr;
328 + const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0;
329 + int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr<int32_t>() : nullptr;
330 +
331 + VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() {
332 + if (isa == ISA::AMX) {
333 + using gemm_t = cpu_micro_gemm::MicroGemm<ISA::AMX, scalar_t>;
334 + if (has_zp) {
335 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, true, false>;
336 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
337 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
338 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
339 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
340 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
341 + scales_group_stride, zeros_group_stride, group_num, group_size,
342 + pack_factor);
343 + return;
344 + }
345 + if (use_desc_act) {
346 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, true>;
347 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
348 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
349 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
350 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
351 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
352 + scales_group_stride, zeros_group_stride, group_num, group_size,
353 + pack_factor);
354 + return;
355 + } else {
356 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, false>;
357 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
358 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
359 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
360 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
361 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
362 + scales_group_stride, zeros_group_stride, group_num, group_size,
363 + pack_factor);
364 + return;
365 + }
366 + } else if (isa == ISA::VEC) {
367 + using gemm_t = cpu_micro_gemm::MicroGemm<ISA::VEC, scalar_t>;
368 + if (has_zp) {
369 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, true, false>;
370 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
371 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
372 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
373 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
374 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
375 + scales_group_stride, zeros_group_stride, group_num, group_size,
376 + pack_factor);
377 + return;
378 + }
379 + if (use_desc_act) {
380 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, true>;
381 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
382 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
383 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
384 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
385 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
386 + scales_group_stride, zeros_group_stride, group_num, group_size,
387 + pack_factor);
388 + return;
389 + } else {
390 + using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, false>;
391 + cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
392 + input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
393 + output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
394 + g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
395 + a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
396 + scales_group_stride, zeros_group_stride, group_num, group_size,
397 + pack_factor);
398 + return;
399 + }
400 + }
401 + });
402 +}
csrc/cpu/dnnl_helper.cpp
+3 -3
@@ -396,9 +396,9 @@ MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
396 : DNNLMatMulPrimitiveHandler(
397 static_cast<DNNLMatMulPrimitiveHandler::Args>(args), args.ab_type),
398 m_size_cache_(nullptr) {
399 - assert(ab_type_ == dnnl::memory::data_type::f32 ||
400 - ab_type_ == dnnl::memory::data_type::bf16 ||
401 - ab_type_ == dnnl::memory::data_type::f16);
399 + assert(b_type_ == dnnl::memory::data_type::f32 ||
400 + b_type_ == dnnl::memory::data_type::bf16 ||
401 + b_type_ == dnnl::memory::data_type::f16);
402
403 dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
404 {b_k_stride_, b_n_stride_});
csrc/cpu/micro_gemm/cpu_micro_gemm_amx.hpp new
+245
@@ -0,0 +1,245 @@
1 +#ifndef CPU_MICRO_GEMM_AMX_HPP
2 +#define CPU_MICRO_GEMM_AMX_HPP
3 +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
4 +
5 +namespace cpu_micro_gemm {
6 +namespace {
7 +// AMX specific
8 +constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
9 +constexpr static int64_t AMX_TILE_ROW_NUM = 16;
10 +constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
11 +
12 +typedef struct __tile_config {
13 + uint8_t palette_id = 1;
14 + uint8_t start_row = 0;
15 + uint8_t reserved_0[14] = {0};
16 + uint16_t colsb[16] = {0};
17 + uint8_t rows[16] = {0};
18 +} __tilecfg;
19 +
20 +// 2-2-4 pattern, for 16 < m <= 32
21 +// TILE 0, 1: load A matrix, row num should be 16, m - 16
22 +// TILE 2, 3: load B matrix, row num should be 16
23 +// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
24 +// - 16
25 +template <typename scalar_t>
26 +class TileGemm224 {
27 + public:
28 + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
29 + TORCH_CHECK(false, "Unsupported data type for TileGemm224");
30 + }
31 +
32 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
33 + TORCH_CHECK(false, "Unsupported data type for TileGemm224");
34 + }
35 +};
36 +
37 +template <>
38 +class TileGemm224<c10::BFloat16> {
39 + public:
40 + using scalar_t = c10::BFloat16;
41 + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
42 + const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
43 + c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
44 + c10::BFloat16* __restrict__ a_tile_1 = a_ptr + lda * AMX_TILE_ROW_NUM;
45 + const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
46 +
47 + // B is always packed as 16 output channels block
48 + c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
49 + c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
50 + const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
51 +
52 + float* __restrict__ c_tile_4 = c_ptr;
53 + float* __restrict__ c_tile_5 =
54 + c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
55 + float* __restrict__ c_tile_6 = c_ptr + AMX_TILE_ROW_NUM * ldc;
56 + float* __restrict__ c_tile_7 =
57 + c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
58 + const int32_t c_tile_stride = ldc * sizeof(float);
59 +
60 + if (accum_c) {
61 + _tile_loadd(4, c_tile_4, c_tile_stride);
62 + _tile_loadd(5, c_tile_5, c_tile_stride);
63 + _tile_loadd(6, c_tile_6, c_tile_stride);
64 + _tile_loadd(7, c_tile_7, c_tile_stride);
65 + } else {
66 + _tile_zero(4);
67 + _tile_zero(5);
68 + _tile_zero(6);
69 + _tile_zero(7);
70 + }
71 +
72 + for (int32_t k = 0; k < k_times; ++k) {
73 + _tile_loadd(0, a_tile_0, a_tile_stride);
74 + _tile_stream_loadd(2, b_tile_2, b_tile_stride);
75 + _tile_dpbf16ps(4, 0, 2);
76 + _tile_stream_loadd(3, b_tile_3, b_tile_stride);
77 + _tile_dpbf16ps(5, 0, 3);
78 + _tile_loadd(1, a_tile_1, a_tile_stride);
79 + _tile_dpbf16ps(6, 1, 2);
80 + _tile_dpbf16ps(7, 1, 3);
81 +
82 + // update ptrs
83 + a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
84 + a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
85 + b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
86 + b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
87 + }
88 +
89 + _tile_stored(4, c_tile_4, c_tile_stride);
90 + _tile_stored(5, c_tile_5, c_tile_stride);
91 + _tile_stored(6, c_tile_6, c_tile_stride);
92 + _tile_stored(7, c_tile_7, c_tile_stride);
93 + }
94 +
95 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
96 + const int32_t m_0 = AMX_TILE_ROW_NUM;
97 + const int32_t m_1 = m - AMX_TILE_ROW_NUM;
98 + config.rows[0] = m_0;
99 + config.rows[1] = m_1;
100 + config.rows[2] = AMX_TILE_ROW_NUM;
101 + config.rows[3] = AMX_TILE_ROW_NUM;
102 + config.rows[4] = m_0;
103 + config.rows[5] = m_0;
104 + config.rows[6] = m_1;
105 + config.rows[7] = m_1;
106 + _tile_loadconfig(&config);
107 + }
108 +};
109 +
110 +// 1-2-2 pattern, for 0 < m <= 16
111 +// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
112 +// m, m
113 +// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
114 +// num should be 16
115 +// TILE 6, 7, (6, 7): store results C matrix, row num should be
116 +// m
117 +template <typename scalar_t>
118 +class TileGemm122 {
119 + public:
120 + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
121 + TORCH_CHECK(false, "Unsupported data type for TileGemm122");
122 + }
123 +
124 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
125 + TORCH_CHECK(false, "Unsupported data type for TileGemm122");
126 + }
127 +};
128 +
129 +template <>
130 +class TileGemm122<c10::BFloat16> {
131 + public:
132 + using scalar_t = c10::BFloat16;
133 + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
134 + c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
135 + c10::BFloat16* __restrict__ a_tile_1 =
136 + a_ptr + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
137 + const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
138 +
139 + c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
140 + c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
141 + c10::BFloat16* __restrict__ b_tile_4 =
142 + b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
143 + c10::BFloat16* __restrict__ b_tile_5 =
144 + b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
145 + int64_t b_stride = AMX_TILE_ROW_BYTES;
146 +
147 + float* __restrict__ c_tile_6 = c_ptr;
148 + float* __restrict__ c_tile_7 = c_ptr + AMX_TILE_ROW_BYTES / sizeof(float);
149 + int64_t c_stride = ldc * sizeof(float);
150 +
151 + const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
152 + const int32_t k_group_times = k_times / 2;
153 + const bool has_tail = (k_times % 2 == 1);
154 +
155 + if (accum_c) {
156 + _tile_loadd(6, c_tile_6, c_stride);
157 + _tile_loadd(7, c_tile_7, c_stride);
158 + } else {
159 + _tile_zero(6);
160 + _tile_zero(7);
161 + }
162 +
163 + for (int32_t k = 0; k < k_group_times; ++k) {
164 + _tile_loadd(0, a_tile_0, a_tile_stride);
165 + _tile_stream_loadd(2, b_tile_2, b_stride);
166 + _tile_dpbf16ps(6, 0, 2);
167 + _tile_stream_loadd(3, b_tile_3, b_stride);
168 + _tile_dpbf16ps(7, 0, 3);
169 + _tile_loadd(1, a_tile_1, a_tile_stride);
170 + _tile_stream_loadd(4, b_tile_4, b_stride);
171 + _tile_dpbf16ps(6, 1, 4);
172 + _tile_stream_loadd(5, b_tile_5, b_stride);
173 + _tile_dpbf16ps(7, 1, 5);
174 +
175 + // update ptrs
176 + a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
177 + a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
178 + b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
179 + b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
180 + b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
181 + b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
182 + }
183 +
184 + if (has_tail) {
185 + _tile_loadd(0, a_tile_0, a_tile_stride);
186 + _tile_stream_loadd(2, b_tile_2, b_stride);
187 + _tile_dpbf16ps(6, 0, 2);
188 + _tile_stream_loadd(3, b_tile_3, b_stride);
189 + _tile_dpbf16ps(7, 0, 3);
190 + }
191 +
192 + _tile_stored(6, c_tile_6, c_stride);
193 + _tile_stored(7, c_tile_7, c_stride);
194 + }
195 +
196 + FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
197 + config.rows[0] = m;
198 + config.rows[1] = m;
199 + config.rows[2] = AMX_TILE_ROW_NUM;
200 + config.rows[3] = AMX_TILE_ROW_NUM;
201 + config.rows[4] = AMX_TILE_ROW_NUM;
202 + config.rows[5] = AMX_TILE_ROW_NUM;
203 + config.rows[6] = m;
204 + config.rows[7] = m;
205 + _tile_loadconfig(&config);
206 + }
207 +};
208 +} // namespace
209 +
210 +// Gemm kernel uses AMX, requires B matrix to be packed
211 +template <typename scalar_t>
212 +class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
213 + public:
214 + static constexpr int32_t MaxMSize = 32;
215 + static constexpr int32_t NSize = 32;
216 +
217 + public:
218 + MicroGemm() : curr_m_(-1) {
219 + vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
220 + }
221 +
222 + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
223 + if (m > AMX_TILE_ROW_NUM) {
224 + if (m != curr_m_) {
225 + curr_m_ = m;
226 + TileGemm224<scalar_t>::init_tile_config(m, amx_tile_config_);
227 + }
228 + TileGemm224<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
229 + } else {
230 + if (m != curr_m_) {
231 + curr_m_ = m;
232 + TileGemm122<scalar_t>::init_tile_config(m, amx_tile_config_);
233 + }
234 + TileGemm122<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
235 + }
236 + }
237 +
238 + private:
239 + alignas(64) __tilecfg amx_tile_config_;
240 + int32_t curr_m_;
241 +};
242 +
243 +} // namespace cpu_micro_gemm
244 +
245 +#endif
csrc/cpu/micro_gemm/cpu_micro_gemm_impl.hpp new
+91
@@ -0,0 +1,91 @@
1 +#ifndef CPU_MICRO_GEMM_IMPL_HPP
2 +#define CPU_MICRO_GEMM_IMPL_HPP
3 +#include "cpu/utils.hpp"
4 +#include "cpu/cpu_types.hpp"
5 +
6 +namespace cpu_micro_gemm {
7 +#define DEFINE_CPU_MICRO_GEMM_PARAMS \
8 + scalar_t *__restrict__ a_ptr, scalar_t *__restrict__ b_ptr, \
9 + float *__restrict__ c_ptr, const int32_t m, const int32_t k, \
10 + const int64_t lda, const int64_t b_n_group_stride, const int64_t ldc, \
11 + const bool accum_c
12 +
13 +#define CPU_MICRO_GEMM_PARAMS \
14 + a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
15 +
16 +template <cpu_utils::ISA isa, typename scalar_t>
17 +class MicroGemm {
18 + public:
19 + static constexpr int32_t MaxMSize = 16;
20 + static constexpr int32_t NSize = 16;
21 +
22 + public:
23 + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
24 + TORCH_CHECK(false, "Unimplemented MicroGemm.");
25 + }
26 +};
27 +
28 +template <int32_t n_size, typename scalar_t>
29 +FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
30 + scalar_t* __restrict__ d_ptr,
31 + const int32_t m, const int64_t ldc,
32 + const int64_t ldd) {
33 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
34 + static_assert(n_size % 16 == 0);
35 +
36 + float* __restrict__ curr_c = c_ptr;
37 + scalar_t* __restrict__ curr_d = d_ptr;
38 + for (int32_t i = 0; i < m; ++i) {
39 + float* __restrict__ curr_c_iter = curr_c;
40 + scalar_t* __restrict__ curr_d_iter = curr_d;
41 + vec_op::unroll_loop<int32_t, n_size / 16>([&](int32_t n_g_idx) {
42 + vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
43 + scalar_vec_t c_vec(c_vec_fp32);
44 + c_vec.save(curr_d_iter);
45 + curr_c_iter += 16;
46 + curr_d_iter += 16;
47 + });
48 + curr_c += ldc;
49 + curr_d += ldd;
50 + }
51 +}
52 +
53 +template <int32_t n_size, typename scalar_t>
54 +FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
55 + scalar_t* __restrict__ d_ptr,
56 + scalar_t* __restrict__ bias_ptr,
57 + const int32_t m, const int64_t ldc,
58 + const int64_t ldd) {
59 + using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
60 + static_assert(n_size % 16 == 0);
61 + constexpr int32_t n_group_num = n_size / 16;
62 + static_assert(n_group_num <= 16);
63 +
64 + vec_op::FP32Vec16 bias_vecs[n_group_num];
65 + scalar_t* __restrict__ curr_bias = bias_ptr;
66 + vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
67 + scalar_vec_t vec(curr_bias);
68 + bias_vecs[i] = vec_op::FP32Vec16(vec);
69 + curr_bias += 16;
70 + });
71 +
72 + float* __restrict__ curr_c = c_ptr;
73 + scalar_t* __restrict__ curr_d = d_ptr;
74 + for (int32_t i = 0; i < m; ++i) {
75 + float* __restrict__ curr_c_iter = curr_c;
76 + scalar_t* __restrict__ curr_d_iter = curr_d;
77 + vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
78 + vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
79 + c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
80 + scalar_vec_t c_vec(c_vec_fp32);
81 + c_vec.save(curr_d_iter);
82 + curr_c_iter += 16;
83 + curr_d_iter += 16;
84 + });
85 + curr_c += ldc;
86 + curr_d += ldd;
87 + }
88 +}
89 +} // namespace cpu_micro_gemm
90 +
91 +#endif
csrc/cpu/micro_gemm/cpu_micro_gemm_vec.hpp new
+115
@@ -0,0 +1,115 @@
1 +#ifndef CPU_MICRO_GEMM_VEC_HPP
2 +#define CPU_MICRO_GEMM_VEC_HPP
3 +#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
4 +
5 +namespace cpu_micro_gemm {
6 +namespace {
7 +// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
8 +template <typename scalar_t>
9 +class TileGemm82 {
10 + public:
11 + FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
12 + switch (m) {
13 + case 1:
14 + gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
15 + break;
16 + case 2:
17 + gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
18 + break;
19 + case 3:
20 + gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
21 + break;
22 + case 4:
23 + gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
24 + break;
25 + case 5:
26 + gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
27 + break;
28 + case 6:
29 + gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
30 + break;
31 + case 7:
32 + gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
33 + break;
34 + case 8:
35 + gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
36 + break;
37 + }
38 + }
39 +
40 + template <int32_t M>
41 + static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
42 + static_assert(0 < M <= 8);
43 + using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
44 +
45 + scalar_t* __restrict__ curr_b_0 = b_ptr;
46 + scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
47 + float* __restrict__ curr_c_0 = c_ptr;
48 + float* __restrict__ curr_c_1 = c_ptr + 16;
49 +
50 + vec_op::FP32Vec16 c_regs[M * 2];
51 + if (accum_c) {
52 + float* __restrict__ curr_m_c_0 = curr_c_0;
53 + float* __restrict__ curr_m_c_1 = curr_c_1;
54 + vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
55 + c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
56 + c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
57 +
58 + // update
59 + curr_m_c_0 += ldc;
60 + curr_m_c_1 += ldc;
61 + });
62 + }
63 +
64 + scalar_t* __restrict__ curr_a = a_ptr;
65 + for (int32_t k_idx = 0; k_idx < k; ++k_idx) {
66 + load_vec_t b_0_reg(curr_b_0);
67 + vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
68 + load_vec_t b_1_reg(curr_b_1);
69 + vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
70 +
71 + scalar_t* __restrict__ curr_m_a = curr_a;
72 + vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
73 + scalar_t v = *curr_m_a;
74 + load_vec_t a_reg_original(v);
75 + vec_op::FP32Vec16 a_reg(a_reg_original);
76 + c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
77 + c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
78 +
79 + // update
80 + curr_m_a += lda;
81 + });
82 +
83 + // update
84 + curr_a += 1;
85 + curr_b_0 += 16;
86 + curr_b_1 += 16;
87 + }
88 +
89 + vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
90 + c_regs[i * 2].save(curr_c_0);
91 + c_regs[i * 2 + 1].save(curr_c_1);
92 +
93 + // update
94 + curr_c_0 += ldc;
95 + curr_c_1 += ldc;
96 + });
97 + }
98 +};
99 +} // namespace
100 +
101 +// Gemm kernel uses vector instructions, requires B matrix to be packed
102 +template <typename scalar_t>
103 +class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
104 + public:
105 + static constexpr int32_t MaxMSize = 8;
106 + static constexpr int32_t NSize = 32;
107 +
108 + public:
109 + void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
110 + TileGemm82<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
111 + }
112 +};
113 +} // namespace cpu_micro_gemm
114 +
115 +#endif
csrc/cpu/torch_bindings.cpp
+16
@@ -103,6 +103,13 @@ void cpu_attention_with_kv_cache(
103 // Note: just for avoiding importing errors
104 void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
105
106 +void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
107 + torch::Tensor& output, const torch::Tensor& scales,
108 + const std::optional<torch::Tensor>& zeros,
109 + const std::optional<torch::Tensor>& g_idx,
110 + const std::optional<torch::Tensor>& bias,
111 + const int64_t pack_factor, const std::string& isa_hint);
112 +
113 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
114 // vLLM custom ops
115
@@ -283,6 +290,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
290 ops.def("static_scaled_fp8_quant() -> ()", placeholder_op);
291 ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op);
292 ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
293 +
294 + // WNA16
295 +#if defined(__AVX512F__)
296 + ops.def(
297 + "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
298 + "Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
299 + "pack_factor, str isa_hint) -> ()");
300 + ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
301 +#endif
302 }
303
304 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _utils), utils) {
csrc/cpu/utils.hpp new
+55
@@ -0,0 +1,55 @@
1 +#ifndef UTILS_HPP
2 +#define UTILS_HPP
3 +
4 +#include <atomic>
5 +#include <cassert>
6 +#include <cstdint>
7 +#include <unistd.h>
8 +
9 +#include "cpu_types.hpp"
10 +
11 +namespace cpu_utils {
12 +enum class ISA { AMX, VEC };
13 +
14 +template <typename T>
15 +struct VecTypeTrait {
16 + using vec_t = void;
17 +};
18 +
19 +template <>
20 +struct VecTypeTrait<float> {
21 + using vec_t = vec_op::FP32Vec16;
22 +};
23 +
24 +template <>
25 +struct VecTypeTrait<c10::BFloat16> {
26 + using vec_t = vec_op::BF16Vec16;
27 +};
28 +
29 +template <>
30 +struct VecTypeTrait<c10::Half> {
31 + using vec_t = vec_op::FP16Vec16;
32 +};
33 +
34 +struct Counter {
35 + std::atomic<int64_t> counter;
36 + char _padding[56];
37 +
38 + Counter() : counter(0) {}
39 +
40 + void reset_counter() { counter.store(0); }
41 +
42 + int64_t acquire_counter() { return counter++; }
43 +};
44 +
45 +inline int64_t get_l2_size() {
46 + static int64_t size = []() {
47 + long l2_cache_size = sysconf(_SC_LEVEL2_CACHE_SIZE);
48 + assert(l2_cache_size != -1);
49 + return l2_cache_size >> 1; // use 50% of L2 cache
50 + }();
51 + return size;
52 +}
53 +} // namespace cpu_utils
54 +
55 +#endif
docs/getting_started/installation/cpu.md
+1 -3
@@ -97,7 +97,6 @@ Currently, there are no pre-built CPU wheels.
97 - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads, can be set as CPU id lists, `auto` (by default), or `nobind` (to disable binding to individual CPU cores and to inherit user-defined OpenMP variables). For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node respectively. If set to `nobind`, the number of OpenMP threads is determined by the standard `OMP_NUM_THREADS` environment variable.
98 - `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `None`. If the value is not set and use `auto` thread binding, no CPU will be reserved for `world_size == 1`, 1 CPU per rank will be reserved for `world_size > 1`.
99 - `CPU_VISIBLE_MEMORY_NODES`: specify visible NUMA memory nodes for vLLM CPU workers, similar to ```CUDA_VISIBLE_DEVICES```. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. The variable provides more control for the auto thread-binding feature, such as masking nodes and changing nodes binding sequence.
100 -- `VLLM_CPU_MOE_PREPACK` (x86 only): whether to use prepack for MoE layer. This will be passed to `ipex.llm.modules.GatedMLPMOE`. Default is `1` (True). On unsupported CPUs, you might need to set this to `0` (False).
100 - `VLLM_CPU_SGL_KERNEL` (x86 only, Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False).
101
102 ## FAQ
@@ -191,10 +190,9 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
190 - GPTQ (x86 only)
191 - compressed-tensor INT8 W8A8 (x86, s390x)
192
194 -### (x86 only) What is the purpose of `VLLM_CPU_MOE_PREPACK` and `VLLM_CPU_SGL_KERNEL`?
193 +### (x86 only) What is the purpose of `VLLM_CPU_SGL_KERNEL`?
194
195 - Both of them require `amx` CPU flag.
197 - - `VLLM_CPU_MOE_PREPACK` can provide better performance for MoE models
196 - `VLLM_CPU_SGL_KERNEL` can provide better performance for MoE models and small-batch scenarios.
197
198 ### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
requirements/cpu.txt
-1
@@ -22,7 +22,6 @@ datasets # for benchmark scripts
22
23 # Intel Extension for PyTorch, only for x86_64 CPUs
24 intel-openmp==2024.2.1; platform_machine == "x86_64"
25 -intel_extension_for_pytorch==2.8.0; platform_machine == "x86_64"
25 triton==3.2.0; platform_machine == "x86_64" # Triton is required for torch 2.6+cpu, as it is imported in torch.compile.
26
27 # Use this to gather CPU info and optimize based on ARM Neoverse cores
tests/quantization/test_cpu_wna16.py new
+23
@@ -0,0 +1,23 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +
5 +from vllm.platforms import current_platform
6 +
7 +if not current_platform.is_cpu():
8 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
9 +
10 +MODELS = [
11 + "TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
12 + "TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
13 +]
14 +DTYPE = ["bfloat16"]
15 +
16 +
17 +@pytest.mark.parametrize("model", MODELS)
18 +@pytest.mark.parametrize("dtype", DTYPE)
19 +def test_ipex_quant(vllm_runner, model, dtype):
20 + with vllm_runner(model, dtype=dtype) as llm:
21 + output = llm.generate_greedy(["The capital of France is"], max_tokens=32)
22 + assert output
23 + print(output)
vllm/_custom_ops.py
+25
@@ -2702,6 +2702,31 @@ def cpu_attention_with_kv_cache(
2702 )
2703
2704
2705 +def cpu_gemm_wna16(
2706 + input: torch.Tensor,
2707 + q_weight: torch.Tensor,
2708 + scales: torch.Tensor,
2709 + zeros: torch.Tensor | None,
2710 + g_idx: torch.Tensor | None,
2711 + bias: torch.Tensor | None,
2712 + pack_factor: int,
2713 + isa_hint: str,
2714 +) -> torch.Tensor:
2715 + output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype)
2716 + torch.ops._C.cpu_gemm_wna16(
2717 + input,
2718 + q_weight,
2719 + output,
2720 + scales,
2721 + zeros,
2722 + g_idx,
2723 + bias,
2724 + pack_factor,
2725 + isa_hint,
2726 + )
2727 + return output
2728 +
2729 +
2730 if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
2731
2732 @register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
vllm/config/model.py
+2
@@ -1020,6 +1020,8 @@ class ModelConfig:
1020 # Ensure heavy backends are probed last to avoid unnecessary
1021 # imports during override detection (e.g., MXFP4 imports Triton)
1022 "mxfp4",
1023 + "cpu_gptq",
1024 + "cpu_awq",
1025 ]
1026 quantization_methods = [
1027 q for q in supported_quantization if q not in overrides
vllm/envs.py
-5
@@ -50,7 +50,6 @@ if TYPE_CHECKING:
50 VLLM_CPU_KVCACHE_SPACE: int | None = 0
51 VLLM_CPU_OMP_THREADS_BIND: str = ""
52 VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
53 - VLLM_CPU_MOE_PREPACK: bool = True
53 VLLM_CPU_SGL_KERNEL: bool = False
54 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
55 VLLM_XLA_CHECK_RECOMPILATION: bool = False
@@ -665,10 +664,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
664 )
665 if "VLLM_CPU_NUM_OF_RESERVED_CPU" in os.environ
666 else None,
668 - # (CPU backend only) whether to use prepack for MoE layer. This will be
669 - # passed to ipex.llm.modules.GatedMLPMOE. On unsupported CPUs, you might
670 - # need to set this to "0" (False).
671 - "VLLM_CPU_MOE_PREPACK": lambda: bool(int(os.getenv("VLLM_CPU_MOE_PREPACK", "1"))),
667 # (CPU backend only) whether to use SGL kernels, optimized for small batch.
668 "VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
669 # If the env var is set, Ray Compiled Graph uses the specified
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
-49
@@ -6,7 +6,6 @@ import torch
6 from torch.nn import functional as F
7
8 from vllm import _custom_ops as ops
9 -from vllm import envs
9
10
11 def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
@@ -130,54 +129,6 @@ def select_experts(
129 )
130
131
133 -class IPEXFusedMOE:
134 - def __init__(self, layer: torch.nn.Module) -> None:
135 - import intel_extension_for_pytorch as ipex
136 -
137 - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
138 - layer.w13_weight,
139 - layer.w2_weight,
140 - use_prepack=envs.VLLM_CPU_MOE_PREPACK,
141 - )
142 -
143 - def __call__(
144 - self,
145 - layer: torch.nn.Module,
146 - x: torch.Tensor,
147 - use_grouped_topk: bool,
148 - top_k: int,
149 - router_logits: torch.Tensor,
150 - renormalize: bool,
151 - topk_group: int | None = None,
152 - num_expert_group: int | None = None,
153 - global_num_experts: int = -1,
154 - expert_map: torch.Tensor | None = None,
155 - custom_routing_function: Callable | None = None,
156 - scoring_func: str = "softmax",
157 - routed_scaling_factor: float = 1.0,
158 - e_score_correction_bias: torch.Tensor | None = None,
159 - apply_router_weight_on_input: bool = False,
160 - activation: str = "silu",
161 - ) -> torch.Tensor:
162 - assert activation == "silu", f"{activation} is not supported."
163 - assert not apply_router_weight_on_input
164 - assert routed_scaling_factor == 1.0, (
165 - f"routed_scaling_factor {routed_scaling_factor} is not supported."
166 - )
167 - return layer.ipex_fusion(
168 - x,
169 - use_grouped_topk,
170 - top_k,
171 - router_logits,
172 - renormalize,
173 - topk_group,
174 - num_expert_group,
175 - custom_routing_function,
176 - scoring_func,
177 - e_score_correction_bias,
178 - )
179 -
180 -
132 class SGLFusedMOE:
133 def __init__(self, layer: torch.nn.Module) -> None:
134 pass
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+1 -1
@@ -260,7 +260,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
260 layer.w2_weight.copy_(packed_w2_weight)
261 layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
262 else:
263 - layer.cpu_fused_moe = cpu_fused_moe.IPEXFusedMOE(layer)
263 + layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
264 else:
265 layer.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
266
vllm/model_executor/layers/quantization/__init__.py
+5
@@ -38,6 +38,8 @@ QuantizationMethods = Literal[
38 "inc",
39 "mxfp4",
40 "petit_nvfp4",
41 + "cpu_gptq",
42 + "cpu_awq",
43 ]
44 QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
45
@@ -107,6 +109,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
109 from .compressed_tensors.compressed_tensors import (
110 CompressedTensorsConfig,
111 )
112 + from .cpu_wna16 import CPUAWQConfig, CPUGPTQConfig
113 from .deepspeedfp import DeepSpeedFPConfig
114 from .experts_int8 import ExpertsInt8Config
115 from .fbgemm_fp8 import FBGEMMFp8Config
@@ -159,6 +162,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
162 "inc": INCConfig,
163 "mxfp4": Mxfp4Config,
164 "petit_nvfp4": PetitNvFp4Config,
165 + "cpu_gptq": CPUGPTQConfig,
166 + "cpu_awq": CPUAWQConfig,
167 }
168 # Update the `method_to_config` with customized quantization methods.
169 method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
vllm/model_executor/layers/quantization/cpu_wna16.py new
+625
@@ -0,0 +1,625 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any, Optional
5 +
6 +import torch
7 +from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
8 +
9 +from vllm._custom_ops import (
10 + cpu_gemm_wna16,
11 +)
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.layers.linear import (
14 + LinearBase,
15 + LinearMethodBase,
16 + UnquantizedLinearMethod,
17 +)
18 +from vllm.model_executor.layers.quantization import QuantizationMethods
19 +from vllm.model_executor.layers.quantization.base_config import (
20 + QuantizationConfig,
21 + QuantizeMethodBase,
22 +)
23 +from vllm.model_executor.layers.quantization.utils.gptq_utils import (
24 + get_linear_quant_method,
25 +)
26 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
27 + marlin_repeat_scales_on_all_ranks,
28 +)
29 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 + is_layer_skipped,
31 + pack_cols,
32 + unpack_cols,
33 +)
34 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
35 +from vllm.model_executor.models.utils import WeightsMapper
36 +from vllm.model_executor.parameter import (
37 + ChannelQuantScaleParameter,
38 + GroupQuantScaleParameter,
39 + PackedColumnParameter,
40 + PackedvLLMParameter,
41 + RowvLLMParameter,
42 +)
43 +from vllm.model_executor.utils import set_weight_attrs
44 +from vllm.platforms import current_platform
45 +from vllm.transformers_utils.config import get_safetensors_params_metadata
46 +from vllm.utils.collection_utils import is_list_of
47 +
48 +logger = init_logger(__name__)
49 +
50 +
51 +class CPUGPTQConfig(QuantizationConfig):
52 + """Config class for CPU GPTQ quant"""
53 +
54 + def __init__(
55 + self,
56 + weight_bits: int,
57 + group_size: int,
58 + desc_act: bool,
59 + is_sym: bool,
60 + lm_head_quantized: bool,
61 + dynamic: dict[str, dict[str, int | bool]],
62 + full_config: dict[str, Any],
63 + modules_in_block_to_quantize: list[str] | None = None,
64 + ) -> None:
65 + super().__init__()
66 + if desc_act and group_size == -1:
67 + # In this case, act_order == True is the same as act_order == False
68 + # (since we have only one group per output channel)
69 + desc_act = False
70 +
71 + # GPTQModel use `dynamic` config property to allow per module
72 + # quantization config so each module can be individually optimized.
73 + # Format is dict[str, dict] where key is a regex string that can
74 + # perform both positive ("+:" prefixed) or negative ("-:" prefixed)
75 + # matching of a module.
76 + # Default to positive match, override base quant config mode, if no
77 + # prefix is used. Value is in dict format of field key and override
78 + # value.
79 + # Negative matching will skip quantization init for this module
80 + # entirely:
81 + # non-quantized inference. More details and quantization examples can be
82 + # found at: https://github.com/ModelCloud/GPTQModel
83 + # Example:
84 + # # last 1/2 of the layers 10-21 has 8bit vs 4bit for 0-9
85 + # # last 1/4 of the layers 16-21 has 8bit and group_size 64
86 + # dynamic = {
87 + # #`.*\.` matches the layers_node prefix
88 + # # positive match layer 10-15
89 + # r"+:.*\.(?:1[0-5])\..*": {"bits": 8,},
90 + # # positive match layer 16-21
91 + # r"+:.*\.(?:1[6-9]|20|21)\..*": {"bits": 8, "group_size": 64,},
92 + # r"-:.*\.moe\..*": {}, # negative match (skip) all `moe` layers
93 + # }
94 + assert weight_bits == 4
95 + self.dynamic = dynamic
96 + self.weight_bits = weight_bits
97 + self.is_sym = is_sym
98 + self.pack_factor = 32 // weight_bits # packed into int32
99 + self.group_size = group_size
100 + self.desc_act = desc_act
101 + self.lm_head_quantized = lm_head_quantized
102 + self.full_config = full_config
103 + self.modules_in_block_to_quantize = modules_in_block_to_quantize or []
104 +
105 + def __repr__(self) -> str:
106 + return (
107 + f"CPUWNA16Config("
108 + f"group_size={self.group_size}, "
109 + f"desc_act={self.desc_act}, "
110 + f"lm_head_quantized={self.lm_head_quantized}, "
111 + f"dynamic={self.dynamic}, "
112 + f"modules_in_block_to_quantize={self.modules_in_block_to_quantize})"
113 + )
114 +
115 + @classmethod
116 + def get_name(cls) -> QuantizationMethods:
117 + return "cpu_gptq"
118 +
119 + @classmethod
120 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:
121 + return [torch.half, torch.bfloat16]
122 +
123 + @classmethod
124 + def get_min_capability(cls) -> int:
125 + return -1
126 +
127 + @classmethod
128 + def get_config_filenames(cls) -> list[str]:
129 + return ["quantize_config.json"]
130 +
131 + @classmethod
132 + def from_config(cls, config: dict[str, Any]) -> "CPUGPTQConfig":
133 + weight_bits = cls.get_from_keys(config, ["bits"])
134 + desc_act = cls.get_from_keys_or(config, ["desc_act"], default=False)
135 + dynamic = cls.get_from_keys_or(config, ["dynamic"], default={})
136 + group_size = cls.get_from_keys(config, ["group_size"])
137 + is_sym = cls.get_from_keys(config, ["sym"])
138 + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
139 + modules_in_block_to_quantize = cls.get_from_keys_or(
140 + config, ["modules_in_block_to_quantize"], default=None
141 + )
142 + return cls(
143 + weight_bits,
144 + group_size,
145 + desc_act,
146 + is_sym,
147 + lm_head_quantized,
148 + dynamic,
149 + config,
150 + modules_in_block_to_quantize,
151 + )
152 +
153 + @classmethod
154 + def override_quantization_method(
155 + cls, hf_quant_cfg, user_quant
156 + ) -> QuantizationMethods | None:
157 + quant_method = hf_quant_cfg.get("quant_method", "").lower()
158 + if current_platform.is_cpu() and (quant_method == "gptq"):
159 + return cls.get_name()
160 + return None
161 +
162 + def get_quant_method(
163 + self, layer: torch.nn.Module, prefix: str
164 + ) -> Optional["QuantizeMethodBase"]:
165 + return get_linear_quant_method(self, layer, prefix, CPUGPTQLinearMethod) # type: ignore
166 +
167 + def apply_vllm_mapper(self, hf_to_vllm_mapper):
168 + if self.modules_in_block_to_quantize is not None:
169 + self.modules_in_block_to_quantize = hf_to_vllm_mapper.apply_list(
170 + self.modules_in_block_to_quantize
171 + )
172 +
173 + def maybe_update_config(self, model_name: str, revision: str | None = None):
174 + if self.modules_in_block_to_quantize:
175 + if is_list_of(self.modules_in_block_to_quantize, list):
176 + # original modules_in_block_to_quantize: list[list[str]]
177 + # flatten original modules_in_block_to_quantize
178 + self.modules_in_block_to_quantize = [
179 + item
180 + for sublist in self.modules_in_block_to_quantize
181 + for item in sublist
182 + ]
183 + return
184 +
185 + unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
186 + metadata = get_safetensors_params_metadata(model_name, revision=revision)
187 + quant_layers: set[str] = {
188 + param_name.rsplit(".", 1)[0]
189 + for param_name, info in metadata.items()
190 + if (dtype := info.get("dtype", None))
191 + and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
192 + }
193 + self.modules_in_block_to_quantize = list(quant_layers)
194 +
195 +
196 +class CPUGPTQLinearMethod(LinearMethodBase):
197 + """Linear method for GPTQ on CPU.
198 +
199 + Args:
200 + quant_config: The CPUWNA16 quantization config.
201 + """
202 +
203 + def __init__(self, quant_config: CPUGPTQConfig) -> None:
204 + self.quant_config = quant_config
205 + assert self.quant_config.is_sym, "GPTQ asym quant is not supported on CPU"
206 +
207 + def create_weights(
208 + self,
209 + layer: torch.nn.Module,
210 + input_size_per_partition: int,
211 + output_partition_sizes: list[int],
212 + input_size: int,
213 + output_size: int,
214 + params_dtype: torch.dtype,
215 + **extra_weight_attrs,
216 + ) -> None:
217 + output_size_per_partition = sum(output_partition_sizes)
218 + assert output_size_per_partition * self.quant_config.weight_bits % 32 == 0
219 + assert output_size_per_partition % 32 == 0
220 + assert input_size_per_partition % 32 == 0
221 +
222 + is_row_parallel = input_size != input_size_per_partition
223 + weight_loader = extra_weight_attrs.get("weight_loader")
224 +
225 + # Normalize group_size
226 + if self.quant_config.group_size != -1:
227 + group_size = self.quant_config.group_size
228 + else:
229 + group_size = input_size
230 +
231 + # Determine sharding
232 + if marlin_repeat_scales_on_all_ranks(
233 + self.quant_config.desc_act, self.quant_config.group_size, is_row_parallel
234 + ):
235 + # By setting scale_dim == None, weight_loader will
236 + # repeat the scales on each rank in TP>1 case.
237 + scales_and_zp_input_dim = None
238 + scales_and_zp_size = input_size // group_size
239 + else:
240 + # By setting scale_dim == 0, weight_loader will
241 + # shard the scales in TP>1 case.
242 + scales_and_zp_input_dim = 0
243 + scales_and_zp_size = input_size_per_partition // group_size
244 +
245 + # Quantized weights
246 + qweight = PackedvLLMParameter(
247 + data=torch.empty(
248 + input_size_per_partition // self.quant_config.pack_factor,
249 + output_size_per_partition,
250 + dtype=torch.int32,
251 + ),
252 + input_dim=0,
253 + output_dim=1,
254 + packed_dim=0,
255 + packed_factor=self.quant_config.pack_factor,
256 + weight_loader=weight_loader,
257 + )
258 +
259 + # Activation order
260 + g_idx = RowvLLMParameter(
261 + data=torch.empty(
262 + input_size_per_partition,
263 + dtype=torch.int32,
264 + ),
265 + input_dim=0,
266 + weight_loader=weight_loader,
267 + )
268 + set_weight_attrs(
269 + g_idx,
270 + {"ignore_warning": True},
271 + )
272 +
273 + qzeros_args = {
274 + "data": torch.empty(
275 + scales_and_zp_size,
276 + output_size_per_partition // self.quant_config.pack_factor,
277 + dtype=torch.int32,
278 + ),
279 + "weight_loader": weight_loader,
280 + }
281 + weight_scale_args = {
282 + "data": torch.empty(
283 + scales_and_zp_size,
284 + output_size_per_partition,
285 + dtype=params_dtype,
286 + ),
287 + "weight_loader": weight_loader,
288 + }
289 +
290 + if scales_and_zp_input_dim is None:
291 + scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
292 + qzeros = PackedColumnParameter(
293 + output_dim=1,
294 + packed_dim=1,
295 + packed_factor=self.quant_config.pack_factor,
296 + **qzeros_args,
297 + )
298 +
299 + else:
300 + scales = GroupQuantScaleParameter(
301 + output_dim=1, input_dim=0, **weight_scale_args
302 + )
303 + qzeros = PackedvLLMParameter(
304 + input_dim=0,
305 + output_dim=1,
306 + packed_dim=1,
307 + packed_factor=self.quant_config.pack_factor,
308 + **qzeros_args,
309 + )
310 +
311 + layer.register_parameter("qweight", qweight)
312 + layer.register_parameter("g_idx", g_idx)
313 + layer.register_parameter("scales", scales)
314 + layer.register_parameter("qzeros", qzeros)
315 +
316 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
317 + torch.set_printoptions(profile="full", linewidth=5000, sci_mode=False)
318 + packed_weight = layer.qweight.data
319 + bits = self.quant_config.weight_bits
320 + pack_factor = int(self.quant_config.pack_factor)
321 + p_w_k, p_w_n = packed_weight.size()
322 + input_size = p_w_k * pack_factor
323 + output_size = p_w_n
324 + isa_hint = _get_isa_hint(layer.scales.dtype)
325 + layer.isa_hint = isa_hint
326 +
327 + layer.qzeros = None
328 + if not self.quant_config.desc_act:
329 + layer.g_idx = None
330 +
331 + # convert input dim packed to output dim packed
332 + weight = unpack_cols(packed_weight, bits, p_w_k, p_w_n * pack_factor).view(
333 + p_w_k, p_w_n, pack_factor
334 + )
335 + weight = weight.permute(0, 2, 1).reshape(input_size, output_size).contiguous()
336 + weight = pack_cols(weight, bits, input_size, output_size)
337 + # make 16 output channel as a block and transpose to the make
338 + # the block contigous
339 + weight = (
340 + weight.view(input_size, -1, 16 // pack_factor)
341 + .permute(1, 0, 2)
342 + .reshape(-1, input_size * 16 // pack_factor)
343 + .contiguous()
344 + )
345 + layer.qweight.data = weight
346 +
347 + def apply(
348 + self,
349 + layer: torch.nn.Module,
350 + x: torch.Tensor,
351 + bias: torch.Tensor | None = None,
352 + ) -> torch.Tensor:
353 + x = cpu_gemm_wna16(
354 + input=x,
355 + q_weight=layer.qweight,
356 + scales=layer.scales,
357 + zeros=layer.qzeros,
358 + g_idx=layer.g_idx,
359 + bias=bias,
360 + pack_factor=8,
361 + isa_hint=layer.isa_hint,
362 + )
363 + return x
364 +
365 +
366 +class CPUAWQConfig(QuantizationConfig):
367 + """Config class for CPU AWQ"""
368 +
369 + def __init__(
370 + self,
371 + weight_bits: int,
372 + group_size: int,
373 + zero_point: bool,
374 + lm_head_quantized: bool,
375 + modules_to_not_convert: list[str] | None,
376 + full_config: dict[str, Any],
377 + ) -> None:
378 + super().__init__()
379 + assert weight_bits == 4
380 + self.pack_factor = 32 // weight_bits # packed into int32
381 + self.group_size = group_size
382 + self.zero_point = zero_point
383 + self.lm_head_quantized = lm_head_quantized
384 + self.weight_bits = weight_bits
385 + self.modules_to_not_convert = modules_to_not_convert or []
386 + self.full_config = full_config
387 +
388 + def __repr__(self) -> str:
389 + return (
390 + f"AWQMarlinConfig("
391 + f"group_size={self.group_size}, "
392 + f"zero_point={self.zero_point}, "
393 + f"lm_head_quantized={self.lm_head_quantized}, "
394 + f"modules_to_not_convert={self.modules_to_not_convert})"
395 + )
396 +
397 + @classmethod
398 + def get_name(cls) -> "QuantizationMethods":
399 + return "cpu_awq"
400 +
401 + @classmethod
402 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:
403 + return [torch.half, torch.bfloat16]
404 +
405 + @classmethod
406 + def get_min_capability(cls) -> int:
407 + return -1
408 +
409 + @classmethod
410 + def get_config_filenames(cls) -> list[str]:
411 + return ["quantize_config.json"]
412 +
413 + @classmethod
414 + def from_config(cls, config: dict[str, Any]) -> "CPUAWQConfig":
415 + weight_bits = cls.get_from_keys(config, ["bits"])
416 + group_size = cls.get_from_keys(config, ["group_size"])
417 + zero_point = cls.get_from_keys(config, ["zero_point"])
418 + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
419 + modules_to_not_convert = cls.get_from_keys_or(
420 + config, ["modules_to_not_convert"], None
421 + )
422 + return cls(
423 + weight_bits,
424 + group_size,
425 + zero_point,
426 + lm_head_quantized,
427 + modules_to_not_convert,
428 + config,
429 + )
430 +
431 + @classmethod
432 + def override_quantization_method(
433 + cls, hf_quant_cfg, user_quant
434 + ) -> Optional["QuantizationMethods"]:
435 + quant_method = hf_quant_cfg.get("quant_method", "").lower()
436 + if current_platform.is_cpu() and (quant_method == "awq"):
437 + return cls.get_name()
438 + return None
439 +
440 + def get_quant_method(
441 + self, layer: torch.nn.Module, prefix: str
442 + ) -> Optional["QuantizeMethodBase"]:
443 + if isinstance(layer, LinearBase) or (
444 + isinstance(layer, ParallelLMHead) and self.lm_head_quantized
445 + ):
446 + if is_layer_skipped(
447 + prefix,
448 + self.modules_to_not_convert,
449 + self.packed_modules_mapping,
450 + skip_with_substr=True,
451 + ):
452 + return UnquantizedLinearMethod()
453 + return CPUAWQLinearMethod(self)
454 + return None
455 +
456 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
457 + if self.modules_to_not_convert:
458 + self.modules_to_not_convert = hf_to_vllm_mapper.apply_list(
459 + self.modules_to_not_convert
460 + )
461 +
462 + def maybe_update_config(self, model_name: str, revision: str | None = None):
463 + if self.modules_to_not_convert:
464 + return
465 +
466 + unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
467 + metadata = get_safetensors_params_metadata(model_name, revision=revision)
468 + layers = {param_name.rsplit(".", 1)[0] for param_name in metadata}
469 + quant_layers: set[str] = {
470 + param_name.rsplit(".", 1)[0]
471 + for param_name, info in metadata.items()
472 + if (dtype := info.get("dtype", None))
473 + and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
474 + }
475 + self.modules_to_not_convert = list(layers - quant_layers)
476 +
477 +
478 +class CPUAWQLinearMethod(LinearMethodBase):
479 + """Linear method for CPU AWQ.
480 +
481 + Args:
482 + quant_config: The CPU AWQ quantization config.
483 + """
484 +
485 + def __init__(self, quant_config: CPUAWQConfig) -> None:
486 + self.quant_config = quant_config
487 + assert self.quant_config.zero_point
488 +
489 + def create_weights(
490 + self,
491 + layer: torch.nn.Module,
492 + input_size_per_partition: int,
493 + output_partition_sizes: list[int],
494 + input_size: int,
495 + output_size: int,
496 + params_dtype: torch.dtype,
497 + **extra_weight_attrs,
498 + ) -> None:
499 + del output_size
500 + output_size_per_partition = sum(output_partition_sizes)
501 + weight_loader = extra_weight_attrs.get("weight_loader")
502 +
503 + # Normalize group_size
504 + if self.quant_config.group_size != -1:
505 + group_size = self.quant_config.group_size
506 + else:
507 + group_size = input_size
508 +
509 + qweight = PackedvLLMParameter(
510 + data=torch.empty(
511 + input_size_per_partition,
512 + output_size_per_partition // self.quant_config.pack_factor,
513 + dtype=torch.int32,
514 + ),
515 + input_dim=0,
516 + output_dim=1,
517 + packed_dim=1,
518 + packed_factor=self.quant_config.pack_factor,
519 + weight_loader=weight_loader,
520 + )
521 +
522 + num_groups = input_size_per_partition // group_size
523 +
524 + qzeros = PackedvLLMParameter(
525 + data=torch.empty(
526 + num_groups,
527 + output_size_per_partition // self.quant_config.pack_factor,
528 + dtype=torch.int32,
529 + ),
530 + input_dim=0,
531 + output_dim=1,
532 + packed_dim=1,
533 + packed_factor=self.quant_config.pack_factor,
534 + weight_loader=weight_loader,
535 + )
536 +
537 + scales = GroupQuantScaleParameter(
538 + data=torch.empty(
539 + num_groups,
540 + output_size_per_partition,
541 + dtype=params_dtype,
542 + ),
543 + input_dim=0,
544 + output_dim=1,
545 + weight_loader=weight_loader,
546 + )
547 +
548 + layer.register_parameter("qweight", qweight)
549 + layer.register_parameter("qzeros", qzeros)
550 + layer.register_parameter("scales", scales)
551 +
552 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
553 + torch.set_printoptions(profile="full", linewidth=5000, sci_mode=False)
554 + packed_weight = layer.qweight.data
555 + packed_zeros = layer.qzeros.data
556 + group_num = packed_zeros.size(0)
557 + bits = self.quant_config.weight_bits
558 + pack_factor = int(self.quant_config.pack_factor)
559 + input_size, packed_output_size = packed_weight.size()
560 + output_size = packed_output_size * pack_factor
561 + isa_hint = _get_isa_hint(layer.scales.dtype)
562 + layer.isa_hint = isa_hint
563 +
564 + interleave_map = (0, 4, 1, 5, 2, 6, 3, 7)
565 + weight = unpack_cols(
566 + packed_weight,
567 + bits,
568 + input_size,
569 + output_size,
570 + )
571 + zeros = unpack_cols(
572 + packed_zeros,
573 + bits,
574 + group_num,
575 + output_size,
576 + )
577 + weight = (
578 + weight.view(input_size, -1, pack_factor)[:, :, interleave_map]
579 + .reshape(input_size, output_size)
580 + .contiguous()
581 + )
582 + zeros = (
583 + zeros.view(group_num, -1, pack_factor)[:, :, interleave_map]
584 + .reshape(group_num, output_size)
585 + .contiguous()
586 + )
587 +
588 + zeros = pack_cols(zeros, bits, group_num, output_size).contiguous()
589 + # make 16 output channel as a block and transpose to
590 + # the make the block contigous
591 + weight = pack_cols(weight, bits, input_size, output_size)
592 + weight = (
593 + weight.view(input_size, -1, 16 // pack_factor)
594 + .permute(1, 0, 2)
595 + .reshape(-1, input_size * 16 // pack_factor)
596 + .contiguous()
597 + )
598 + layer.qweight.data = weight
599 + layer.qzeros.data = zeros
600 +
601 + def apply(
602 + self,
603 + layer: torch.nn.Module,
604 + x: torch.Tensor,
605 + bias: torch.Tensor | None = None,
606 + ) -> torch.Tensor:
607 + x = cpu_gemm_wna16(
608 + input=x,
609 + q_weight=layer.qweight,
610 + scales=layer.scales,
611 + zeros=layer.qzeros,
612 + g_idx=None,
613 + bias=bias,
614 + pack_factor=8,
615 + isa_hint=layer.isa_hint,
616 + )
617 + return x
618 +
619 +
620 +def _get_isa_hint(dtype: torch.dtype) -> str:
621 + supports_amx = torch._C._cpu._is_amx_tile_supported()
622 + if supports_amx and dtype in (torch.bfloat16,):
623 + return "amx"
624 + else:
625 + return "vec"
vllm/model_executor/layers/quantization/ipex_quant.py
+1 -1
@@ -134,7 +134,7 @@ class IPEXConfig(QuantizationConfig):
134 def override_quantization_method(
135 cls, hf_quant_cfg, user_quant
136 ) -> QuantizationMethods | None:
137 - if not current_platform.is_cpu() and not current_platform.is_xpu():
137 + if not current_platform.is_xpu():
138 return None
139
140 quant_method = hf_quant_cfg.get("quant_method", "").lower()