@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f09daea261
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 5, 18:30 UTC

[CPU] Support int8 compute mode in CPU AWQ (#35697)

Signed-off-by: Yintong Lu <yintong.lu@intel.com>

Yintong Lu committed Mar 31, 2026 at 15:27 UTC f09daea261ee98340512e7c7a5fce09db6f8ab72
10 files changed +1197 -11
.buildkite/hardware_tests/cpu.yaml
+3 -1
@@ -13,12 +13,14 @@ steps:
13 - tests/kernels/attention/test_cpu_attn.py
14 - tests/kernels/moe/test_cpu_fused_moe.py
15 - tests/kernels/test_onednn.py
16 + - tests/kernels/test_awq_int4_to_int8.py
17 commands:
18 - |
19 bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
20 pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
21 pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
21 - pytest -x -v -s tests/kernels/test_onednn.py"
22 + pytest -x -v -s tests/kernels/test_onednn.py
23 + pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py"
24
25 - label: CPU-Compatibility Tests
26 depends_on: []
cmake/cpu_extension.cmake
+1
@@ -373,6 +373,7 @@ if (ENABLE_X86_ISA)
373 "csrc/cpu/sgl-kernels/gemm.cpp"
374 "csrc/cpu/sgl-kernels/gemm_int8.cpp"
375 "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
376 + "csrc/cpu/sgl-kernels/gemm_int4.cpp"
377 "csrc/cpu/sgl-kernels/moe.cpp"
378 "csrc/cpu/sgl-kernels/moe_int8.cpp"
379 "csrc/cpu/sgl-kernels/moe_fp8.cpp")
csrc/cpu/sgl-kernels/common.h
+8
@@ -117,6 +117,14 @@ inline void parallel_for(int n, const func_t& f) {
117 #endif
118 }
119
120 +inline int get_thread_num() {
121 +#if defined(_OPENMP)
122 + return omp_get_thread_num();
123 +#else
124 + return 0;
125 +#endif
126 +}
127 +
128 // for 1d parallel, use `actual_nth`
129 // for 2d parallel, use even nths, e.g. 43->42
130 int inline adjust_num_threads(int m) {
csrc/cpu/sgl-kernels/gemm.h
+36 -3
@@ -17,8 +17,8 @@ constexpr int block_size_n() { return 2 * TILE_N; }
17 template <typename T> inline bool can_use_brgemm(int M);
18 template <> inline bool can_use_brgemm<at::BFloat16>(int M) { return M > 4; }
19 template <> inline bool can_use_brgemm<at::Half>(int M) { return true; }
20 -// TODO: add u8s8 brgemm, this requires PyTorch 2.7
21 -template <> inline bool can_use_brgemm<int8_t>(int M) { return false; }
20 +template <> inline bool can_use_brgemm<int8_t>(int M) { return M > 4; }
21 +template <> inline bool can_use_brgemm<uint8_t>(int M) { return M > 4; }
22 template <> inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) { return M > 4; }
23 template <> inline bool can_use_brgemm<at::quint4x2>(int M) { return M > 4; }
24
@@ -40,9 +40,17 @@ inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
40 return use_int8_w8a8 ? K + sizeof(int32_t) : K;
41 }
42
43 -// pack weight to vnni format
43 +inline int64_t get_4bit_block_k_size(int64_t group_size) {
44 + return group_size > 128 ? 128 : group_size;
45 +}
46 +
47 +// pack weight into vnni format
48 at::Tensor convert_weight_packed(at::Tensor& weight);
49
50 +// pack weight to vnni format for int4 (adapted from sglang)
51 +std::tuple<at::Tensor, at::Tensor, at::Tensor>
52 +convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
53 +
54 // moe implementations for int8 w8a8
55 template <typename scalar_t>
56 void fused_experts_int8_kernel_impl(
@@ -233,6 +241,31 @@ void tinygemm_kernel(
241 int64_t strideBs,
242 bool brg);
243
244 +// int4 scaled GEMM (adapted from sglang)
245 +at::Tensor int4_scaled_mm_cpu(
246 + at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
247 +
248 +// int4 tinygemm kernel interface(adapted from sglang)
249 +template <typename scalar_t>
250 +void tinygemm_kernel(
251 + scalar_t* C,
252 + float* C_temp,
253 + const uint8_t* A,
254 + const float* scales_a,
255 + const int32_t* qzeros_a,
256 + const uint8_t* B,
257 + const float* scales_b,
258 + const int8_t* qzeros_b,
259 + const int32_t* compensation,
260 + int8_t* dqB_tmp,
261 + int64_t M,
262 + int64_t K,
263 + int64_t lda,
264 + int64_t ldc_f,
265 + int64_t ldc_s,
266 + bool store_out,
267 + bool use_brgemm);
268 +
269 // TODO: debug print, remove me later
270 inline void print_16x32i(const __m512i x) {
271 int32_t a[16];
csrc/cpu/sgl-kernels/gemm_int4.cpp new
+755
@@ -0,0 +1,755 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// Adapted from sgl-project/sglang
3 +// https://github.com/sgl-project/sglang/pull/8226
4 +
5 +#include <ATen/ATen.h>
6 +
7 +#include "common.h"
8 +#include "gemm.h"
9 +#include "vec.h"
10 +
11 +namespace {
12 +
13 +#define BLOCK_N block_size_n()
14 +#define BLOCK_M 128
15 +
16 +template <bool sym_quant_act>
17 +struct ActDtype;
18 +template <>
19 +struct ActDtype<true> {
20 + using type = int8_t;
21 +};
22 +template <>
23 +struct ActDtype<false> {
24 + using type = uint8_t;
25 +};
26 +
27 +struct alignas(32) m256i_wrapper {
28 + __m256i data;
29 +};
30 +
31 +#if defined(CPU_CAPABILITY_AVX512)
32 +inline std::array<m256i_wrapper, 2> load_zps_4vnni(
33 + const int8_t* __restrict__ zps) {
34 + __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps));
35 + __m256i vzps_high =
36 + _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps + 8));
37 + __m256i shuffle_mask =
38 + _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3,
39 + 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
40 + vzps_low = _mm256_shuffle_epi8(vzps_low, shuffle_mask);
41 + vzps_high = _mm256_shuffle_epi8(vzps_high, shuffle_mask);
42 + m256i_wrapper vzps_low_wp, vzps_high_wp;
43 + vzps_low_wp.data = vzps_low;
44 + vzps_high_wp.data = vzps_high;
45 + return {vzps_low_wp, vzps_high_wp};
46 +}
47 +
48 +inline std::array<m256i_wrapper, 2> load_uint4_as_int8(
49 + const uint8_t* __restrict__ qB) {
50 + __m256i packed = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(qB));
51 + const __m256i low_mask = _mm256_set1_epi8(0x0f);
52 + __m256i high = _mm256_srli_epi16(packed, 4);
53 + high = _mm256_and_si256(high, low_mask);
54 + __m256i low = _mm256_and_si256(packed, low_mask);
55 + m256i_wrapper low_wp, high_wp;
56 + low_wp.data = low;
57 + high_wp.data = high;
58 + return {low_wp, high_wp};
59 +}
60 +
61 +template <int N, int ldb>
62 +void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB,
63 + const int8_t* __restrict__ qzeros, int64_t K) {
64 + #pragma GCC unroll 2
65 + for (int n = 0; n < N; n += 16) {
66 + auto [zps_low_wp, zps_high_wp] = load_zps_4vnni(&qzeros[n]);
67 + auto zps_low = zps_low_wp.data;
68 + auto zps_high = zps_high_wp.data;
69 + for (int k = 0; k < K; k += 4) {
70 + auto [vb_low_wp, vb_high_wp] =
71 + load_uint4_as_int8(B + ldb * k + n / 2 * 4);
72 + auto vb_low = vb_low_wp.data;
73 + auto vb_high = vb_high_wp.data;
74 + vb_high = _mm256_sub_epi8(vb_high, zps_high);
75 + vb_low = _mm256_sub_epi8(vb_low, zps_low);
76 + _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4),
77 + vb_low);
78 + _mm256_storeu_si256(
79 + reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
80 + }
81 + }
82 +}
83 +
84 +template <bool sym_quant_act, int N, bool accum>
85 +void _dequant_and_store(float* __restrict__ output,
86 + const int32_t* __restrict__ input,
87 + const float* __restrict__ scale_a,
88 + const int32_t* __restrict__ zp_a,
89 + const float* __restrict__ scale_b,
90 + const int32_t* __restrict__ comp_b, int M, int ldi,
91 + int ldo, int ldsa = 1) {
92 + for (int m = 0; m < M; ++m) {
93 + float a_scale = *(scale_a + m * ldsa);
94 + __m512 va_scale = _mm512_set1_ps(a_scale);
95 + int32_t a_zp;
96 + __m512i va_zp;
97 + if constexpr (!sym_quant_act) {
98 + a_zp = *(zp_a + m * ldsa);
99 + va_zp = _mm512_set1_epi32(a_zp);
100 + }
101 + int n = 0;
102 + #pragma GCC unroll 2
103 + for (; n < N; n += 16) {
104 + __m512i vc = _mm512_loadu_si512(input + m * ldi + n);
105 + if constexpr (!sym_quant_act) {
106 + __m512i vb_comp = _mm512_loadu_si512(comp_b + n);
107 + vc = _mm512_sub_epi32(vc, _mm512_mullo_epi32(vb_comp, va_zp));
108 + }
109 + __m512 vc_f = _mm512_cvtepi32_ps(vc);
110 + __m512 vc_f_mul = _mm512_mul_ps(vc_f, va_scale);
111 + __m512 vb_s = _mm512_loadu_ps(scale_b + n);
112 + vc_f_mul = _mm512_mul_ps(vc_f_mul, vb_s);
113 + if constexpr (accum) {
114 + __m512 vo = _mm512_loadu_ps(output + m * ldo + n);
115 + _mm512_storeu_ps(output + m * ldo + n, _mm512_add_ps(vo, vc_f_mul));
116 + } else {
117 + _mm512_storeu_ps(output + m * ldo + n, vc_f_mul);
118 + }
119 + }
120 + for (; n < N; ++n) {
121 + float dq_val;
122 + if constexpr (sym_quant_act) {
123 + dq_val = (float)input[m * ldi + n] * a_scale * scale_b[n];
124 + } else {
125 + dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale *
126 + scale_b[n];
127 + }
128 + if constexpr (accum) {
129 + output[m * ldo + n] += dq_val;
130 + } else {
131 + output[m * ldo + n] = dq_val;
132 + }
133 + }
134 + }
135 +}
136 +
137 +#else
138 +template <int N, int ldb>
139 +void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB,
140 + const int8_t* qzeros, int64_t K) {
141 + for (int k = 0; k < K; ++k) {
142 + for (int n = 0; n < N / 2; ++n) {
143 + int32_t b = (int32_t)B[k * ldb + n];
144 + dqB[k * N + n * 2] = (b & 0xf) - qzeros[n];
145 + dqB[k * N + n * 2 + 1] = (b >> 4) - qzeros[n];
146 + }
147 + }
148 +}
149 +#endif
150 +
151 +#if defined(CPU_CAPABILITY_AVX512)
152 +inline __m512i combine_m256i(__m256i a, __m256i b) {
153 + __m512i c = _mm512_castsi256_si512(a);
154 + return _mm512_inserti64x4(c, b, 1);
155 +}
156 +
157 +inline __m512i combine_m256i(std::array<m256i_wrapper, 2> two_256) {
158 + return combine_m256i(two_256[0].data, two_256[1].data);
159 +}
160 +
161 +static inline __m512i _mm512_sign_epi8(__m512i a, __m512i b) {
162 + __m512i zero = _mm512_setzero_si512();
163 + __mmask64 blt0 = _mm512_movepi8_mask(b);
164 + return _mm512_mask_sub_epi8(a, blt0, zero, a);
165 +}
166 +
167 +template <bool sym_quant_act, int M, int N, int ldb>
168 +void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
169 + const float* scales_a, const int32_t* qzeros_a,
170 + const uint8_t* B, const float* scales_b,
171 + const int8_t* qzeros_b, int64_t K, int64_t lda,
172 + int64_t ldc) {
173 + constexpr int COLS = N / 16;
174 + __m512i ones = _mm512_set1_epi8(1);
175 + __m512i va;
176 + __m512i vb[COLS];
177 + __m512i vc[M * COLS];
178 + __m512 vscales[COLS];
179 + __m512i vzps[COLS];
180 + __m512i vcompensate[COLS];
181 +
182 + Unroll<COLS>{}([&](auto i) {
183 + vscales[i] = _mm512_loadu_ps(scales_b + i * 16);
184 + vzps[i] = combine_m256i(load_zps_4vnni(qzeros_b + i * 16));
185 + if constexpr (!sym_quant_act) {
186 + vcompensate[i] = _mm512_setzero_epi32();
187 + }
188 + });
189 + Unroll<M * COLS>{}([&](auto i) { vc[i] = _mm512_setzero_epi32(); });
190 +
191 + auto compute = [&](auto i, int k) {
192 + constexpr const int row = i / COLS;
193 + constexpr const int col = i % COLS;
194 +
195 + if constexpr (col == 0) {
196 + va = _mm512_set1_epi32(*(int32_t*)(A + row * lda + k));
197 + }
198 +
199 + if constexpr (row == 0) {
200 + int B_offset = k * ldb + col * 16 * 2;
201 + vb[col] = combine_m256i(load_uint4_as_int8(B + B_offset));
202 + vb[col] = _mm512_sub_epi8(vb[col], vzps[col]);
203 + if constexpr (!sym_quant_act) {
204 + vcompensate[col] = _mm512_dpbusd_epi32(vcompensate[col], ones, vb[col]);
205 + }
206 + _mm_prefetch(B + B_offset + 128 * ldb, _MM_HINT_T0);
207 + }
208 + if constexpr (sym_quant_act) {
209 + auto vsb = _mm512_sign_epi8(vb[col], va);
210 + auto vabsa = _mm512_sign_epi8(va, va);
211 + vc[i] = _mm512_dpbusds_epi32(vc[i], vabsa, vsb);
212 + } else {
213 + vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
214 + }
215 + };
216 +
217 + constexpr const int unroll = 4;
218 + int k = 0;
219 + for (; k < K / 4 / unroll; k++) {
220 + Unroll<unroll>{}(
221 + [&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
222 + }
223 + k *= 4 * unroll;
224 + for (; k < K; k += 4) {
225 + Unroll<M * COLS>{}(compute, k);
226 + }
227 +
228 + auto store = [&](auto i) {
229 + constexpr const int row = i / COLS;
230 + constexpr const int col = i % COLS;
231 + __m512 vc_float;
232 + if constexpr (!sym_quant_act) {
233 + vc[i] = _mm512_sub_epi32(
234 + vc[i], _mm512_mullo_epi32(vcompensate[col],
235 + _mm512_set1_epi32(*(qzeros_a + row))));
236 + }
237 + vc_float = _mm512_cvtepi32_ps(vc[i]);
238 + vc_float = _mm512_mul_ps(vc_float, _mm512_set1_ps(*(scales_a + row)));
239 +
240 + vc_float = _mm512_mul_ps(vc_float, vscales[col]);
241 + auto vc_old = _mm512_loadu_ps(C + row * ldc + col * 16);
242 + vc_float = _mm512_add_ps(vc_float, vc_old);
243 + _mm512_storeu_ps(C + row * ldc + col * 16, vc_float);
244 + };
245 + Unroll<M * COLS>{}(store);
246 +}
247 +
248 + #define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \
249 + _dequant_gemm_accum_small_M<sym_quant_act, M, N, ldb>( \
250 + C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
251 +#endif
252 +
253 +template <bool sym_quant_act, int N, int ldb>
254 +void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
255 + const int32_t* qzeros_a, const uint8_t* B,
256 + const float* scales_b, const int8_t* qzeros_b,
257 + const int32_t* compensation, int8_t* dqB, int64_t M,
258 + int64_t K, int64_t lda, int64_t ldc, bool use_brgemm) {
259 +#if defined(CPU_CAPABILITY_AVX512)
260 + if (!use_brgemm) {
261 + switch (M) {
262 + case 1:
263 + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(1);
264 + break;
265 + case 2:
266 + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(2);
267 + break;
268 + case 3:
269 + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(3);
270 + break;
271 + case 4:
272 + CALL_DEQUANT_GEMM_ACCUM_SMALL_M(4);
273 + break;
274 + default:
275 + TORCH_CHECK(false, "tinygemm_kernel: unexpected M for AVX path!");
276 + }
277 + return;
278 + }
279 +
280 + _dequant_weight_zp_only<N, ldb>(B, dqB, qzeros_b, K);
281 + using Tin = typename ActDtype<sym_quant_act>::type;
282 + Tin* A_ptr = (Tin*)A;
283 + if (use_brgemm) {
284 + int32_t C_i32[M * N];
285 + at::native::cpublas::brgemm(M, N, K, lda, N /*ldb*/, N /*ldc*/,
286 + false /* add_C */, A_ptr, dqB, C_i32,
287 + true /* is_vnni */);
288 + _mm_prefetch(B + N * K / 2, _MM_HINT_T0);
289 + _mm_prefetch(A + K, _MM_HINT_T0);
290 + _dequant_and_store<sym_quant_act, N, true>(C, C_i32, scales_a, qzeros_a,
291 + scales_b, compensation, M,
292 + N /*ldi*/, ldc, 1 /*ldsa*/);
293 + } else
294 +#endif
295 + {
296 + TORCH_CHECK(false, "tinygemm_kernel: scalar path not implemented!");
297 + }
298 +}
299 +
300 +template <int N>
301 +inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
302 + if (bias_ptr) {
303 + for (int i = 0; i < m; ++i) {
304 + int j = 0;
305 +#if defined(CPU_CAPABILITY_AVX512)
306 + #pragma GCC unroll 2
307 + for (; j < N; j += 16) {
308 + __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j);
309 + _mm512_storeu_ps(y_buf + i * N + j, bias_vec);
310 + }
311 +#endif
312 + for (; j < N; ++j) {
313 + y_buf[i * N + j] = bias_ptr[j];
314 + }
315 + }
316 + } else {
317 + for (int i = 0; i < m; ++i) {
318 + int j = 0;
319 +#if defined(CPU_CAPABILITY_AVX512)
320 + #pragma GCC unroll 2
321 + for (; j < N; j += 16) {
322 + __m512 zero_vec = _mm512_setzero_ps();
323 + _mm512_storeu_ps(y_buf + i * N + j, zero_vec);
324 + }
325 +#endif
326 + for (; j < N; ++j) {
327 + y_buf[i * N + j] = 0;
328 + }
329 + }
330 + }
331 +}
332 +
333 +template <int N, typename out_dtype>
334 +inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
335 + int64_t lda) {
336 + for (int i = 0; i < m; ++i) {
337 + int j = 0;
338 + if constexpr (std::is_same<out_dtype, float>::value) {
339 +#if defined(CPU_CAPABILITY_AVX512)
340 + #pragma GCC unroll 2
341 + for (; j < N; j += 16) {
342 + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
343 + _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
344 + }
345 +#endif
346 + for (; j < N; ++j) {
347 + c_ptr[i * lda + j] = y_buf[i * N + j];
348 + }
349 + } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
350 +#if defined(CPU_CAPABILITY_AVX512)
351 + #pragma GCC unroll 2
352 + for (; j < N; j += 16) {
353 + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
354 + __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
355 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
356 + y_bf16_vec);
357 + }
358 +#endif
359 + for (; j < N; ++j) {
360 + c_ptr[i * lda + j] = at::BFloat16(y_buf[i * N + j]);
361 + }
362 + } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
363 +#if defined(CPU_CAPABILITY_AVX512)
364 + #pragma GCC unroll 2
365 + for (; j < N; j += 16) {
366 + __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
367 + __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
368 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
369 + y_fp16_vec);
370 + }
371 +#endif
372 + for (; j < N; ++j) {
373 + c_ptr[i * lda + j] = at::Half(y_buf[i * N + j]);
374 + }
375 + } else {
376 + TORCH_CHECK(false, "Unsupported output dtype");
377 + }
378 + }
379 +}
380 +
381 +void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) {
382 + using iVec = at::vec::Vectorized<int32_t>;
383 + constexpr int VecSize = iVec::size();
384 + const iVec fill_val_vec = iVec(value);
385 + int64_t d;
386 +#pragma GCC unroll 4
387 + for (d = 0; d <= size - VecSize; d += VecSize) {
388 + fill_val_vec.store(output + d);
389 + }
390 + for (; d < size; ++d) {
391 + output[d] = value;
392 + }
393 +}
394 +
395 +template <bool sym_quant_act, typename act_dtype, typename out_dtype>
396 +void _da8w4_linear_impl(
397 + act_dtype* __restrict__ input, const float* __restrict__ input_scales,
398 + const int32_t* __restrict__ input_qzeros,
399 + const uint8_t* __restrict__ weight, const float* __restrict__ weight_scales,
400 + const int8_t* __restrict__ weight_qzeros, const float* __restrict__ bias,
401 + out_dtype* __restrict__ output, float* __restrict__ output_temp,
402 + int8_t* __restrict__ dequant_weight_temp, int64_t M, int64_t N, int64_t K,
403 + int64_t num_groups) {
404 + const bool use_brgemm = can_use_brgemm<act_dtype>(M);
405 + int64_t block_m = [&]() -> long {
406 + if (M <= 48) {
407 + return M;
408 + } else if (M < 64) {
409 + return 32;
410 + } else if (M < 96) {
411 + return 64;
412 + } else {
413 + return 128;
414 + }
415 + }();
416 + int64_t Mc = div_up(M, block_m);
417 + bool parallel_on_M = M > 128;
418 + int64_t Nc = N / BLOCK_N;
419 + int64_t num_blocks = parallel_on_M ? Mc * Nc : Nc;
420 + int64_t group_size = div_up(K, num_groups);
421 + int64_t _block_k = get_4bit_block_k_size(group_size);
422 + int64_t Kc = K / _block_k;
423 + int64_t block_per_group = group_size / _block_k;
424 +
425 + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
426 + int tid = get_thread_num();
427 + float* C_tmp = output_temp + tid * block_m * BLOCK_N;
428 + int8_t* dqB_tmp = dequant_weight_temp + tid * _block_k * BLOCK_N;
429 + for (const auto i : c10::irange(begin, end)) {
430 + int64_t mc = parallel_on_M ? i / Nc : 0;
431 + int64_t nc = parallel_on_M ? i % Nc : i;
432 + int64_t mc_end = parallel_on_M ? mc + 1 : Mc;
433 +
434 + for (int mci = mc; mci < mc_end; ++mci) {
435 + int64_t m_size =
436 + mci * block_m + block_m > M ? M - mci * block_m : block_m;
437 + auto bias_data = bias ? bias + nc * BLOCK_N : nullptr;
438 + copy_bias<BLOCK_N>(bias_data, C_tmp, m_size);
439 + for (int kci = 0; kci < Kc; ++kci) {
440 + int32_t* compensation_ptr =
441 + sym_quant_act
442 + ? nullptr
443 + : (int32_t*)(void*)(weight +
444 + (nc * Kc + kci) *
445 + (BLOCK_N *
446 + (_block_k / 2 + sizeof(int32_t))) +
447 + _block_k * BLOCK_N / 2);
448 + _dequant_gemm_accum<sym_quant_act, BLOCK_N, BLOCK_N / 2>(
449 + /*C*/ C_tmp,
450 + /*A*/ (uint8_t*)input + mci * block_m * K + kci * _block_k,
451 + /*scales_a*/ input_scales + mci * block_m,
452 + /*qzeros_a*/ input_qzeros + mci * block_m,
453 + /*B*/ weight + (nc * Kc + kci) *
454 + (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
455 + /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups +
456 + kci / block_per_group * BLOCK_N,
457 + /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups +
458 + kci / block_per_group * BLOCK_N,
459 + /*Bcomp*/ compensation_ptr,
460 + /*dqB_tmp*/ dqB_tmp,
461 + /*M*/ m_size,
462 + /*K*/ _block_k,
463 + /*lda*/ K,
464 + /*ldc*/ BLOCK_N,
465 + /*use_brgemm*/ use_brgemm);
466 + }
467 + store_out<BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N,
468 + m_size, N /*lda*/);
469 + }
470 + }
471 + if (use_brgemm) {
472 + at::native::cpublas::brgemm_release();
473 + }
474 + });
475 +}
476 +
477 +} // anonymous namespace
478 +
479 +std::tuple<at::Tensor, at::Tensor, at::Tensor>
480 +convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
481 + const at::Tensor& scales,
482 + const at::Tensor& qzeros) {
483 + TORCH_CHECK(weight.dim() == 2,
484 + "DA8W4 CPU: Weight should be a 2D tensor for packing");
485 + TORCH_CHECK(
486 + weight.size(1) % 2 == 0,
487 + "DA8W4 CPU: Weight should have even number of columns for packing");
488 +
489 + auto new_scales = scales;
490 + auto new_qzeros = qzeros;
491 + if (new_scales.dim() == 1) {
492 + new_scales.unsqueeze_(1);
493 + }
494 + new_scales = new_scales.to(at::kFloat);
495 + if (new_qzeros.dim() == 1) {
496 + new_qzeros.unsqueeze_(1);
497 + }
498 + new_qzeros = new_qzeros.to(at::kChar);
499 + int64_t N = weight.size(0);
500 + int64_t K = weight.size(1);
501 + int64_t G = scales.size(1);
502 + int64_t group_size = K / G;
503 + int64_t _block_k = get_4bit_block_k_size(group_size);
504 + constexpr int block_n = block_size_n();
505 + int64_t Nc = N / block_n;
506 + int64_t Kc = K / _block_k;
507 +
508 + auto weight_view = weight.view({Nc, block_n, Kc, _block_k});
509 + at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
510 + at::Tensor blocked_weight;
511 + at::Tensor blocked_scales =
512 + new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
513 + at::Tensor blocked_qzeros =
514 + new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
515 + auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) -
516 + new_qzeros.view({Nc, block_n, G, -1});
517 + weight_sub_qzero = weight_sub_qzero.view({Nc, block_n, Kc, _block_k});
518 + at::Tensor compensation = weight_sub_qzero.sum(-1);
519 + compensation = compensation.permute({0, 2, 1}).contiguous().to(at::kInt);
520 + int64_t buffer_size_nbytes =
521 + _block_k * block_n / 2 + block_n * sizeof(int32_t);
522 + blocked_weight = at::empty({Nc, Kc, buffer_size_nbytes}, weight.options());
523 +
524 + auto weight_ptr = weight_reordered.data_ptr<uint8_t>();
525 + auto compensation_ptr = compensation.data_ptr<int32_t>();
526 + auto blocked_weight_ptr = blocked_weight.data_ptr<uint8_t>();
527 + int64_t num_blocks = Nc * Kc;
528 + at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
529 + for (const auto i : c10::irange(begin, end)) {
530 + auto in_ptr = weight_ptr + i * _block_k * block_n;
531 + auto out_ptr =
532 + blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
533 + int32_t* comp_in_prt = compensation_ptr + i * block_n;
534 + int32_t* comp_out_prt =
535 + (int32_t*)(void*)(blocked_weight_ptr +
536 + i * block_n * (_block_k / 2 + sizeof(int32_t)) +
537 + _block_k * block_n / 2);
538 + constexpr int n_group_size = 8;
539 + constexpr int vnni_size = 4;
540 + constexpr int n_group = block_n / n_group_size;
541 + for (int nb = 0; nb < n_group; nb += 2) {
542 + for (int k = 0; k < _block_k; k += vnni_size) {
543 + for (int ni = 0; ni < n_group_size; ++ni) {
544 + for (int ki = 0; ki < vnni_size; ++ki) {
545 + int src_idx_1 = nb * n_group_size + ni + (k + ki) * block_n;
546 + int src_idx_2 = (nb + 1) * n_group_size + ni + (k + ki) * block_n;
547 + int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size +
548 + k * block_n / 2 + ki;
549 + uint8_t src_1 = *(in_ptr + src_idx_1);
550 + uint8_t src_2 = *(in_ptr + src_idx_2);
551 + uint8_t dst = (src_1 & 0x0f) | ((src_2 & 0x0f) << 4);
552 + *(out_ptr + dst_idx) = dst;
553 + }
554 + }
555 + }
556 + }
557 + for (int nb = 0; nb < block_n; nb++) {
558 + *(comp_out_prt + nb) = *(comp_in_prt + nb);
559 + }
560 + }
561 + });
562 +
563 + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales),
564 + std::move(blocked_qzeros));
565 +}
566 +
567 +std::tuple<at::Tensor, at::Tensor> autoawq_to_int4pack(at::Tensor qweight,
568 + at::Tensor qzeros) {
569 + auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
570 + auto qweight_unsq = qweight.unsqueeze(-1);
571 + auto unpacked = at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF;
572 + auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte);
573 +
574 + auto qzeros_unsq = qzeros.unsqueeze(-1);
575 + auto qzeros_unpacked = at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF;
576 + auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte);
577 +
578 + return std::make_tuple(qweight_final, qzeros_final);
579 +}
580 +
581 +std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
582 + at::Tensor qweight, at::Tensor qzeros, at::Tensor scales) {
583 + auto res = autoawq_to_int4pack(qweight, qzeros);
584 + auto _qweight = std::get<0>(res);
585 + auto _qzeros = std::get<1>(res);
586 + auto _scales = scales;
587 + _qzeros = _qzeros.transpose(-2, -1).contiguous();
588 + _scales = _scales.transpose(-2, -1).contiguous();
589 + if (_qweight.dim() == 3) {
590 + int64_t E = _qweight.size(0);
591 + int64_t K = _qweight.size(2);
592 + int64_t G = _scales.size(2);
593 + int64_t group_size = K / G;
594 + int64_t _block_k = get_4bit_block_k_size(group_size);
595 + int64_t block_n = block_size_n();
596 + int64_t Nc = _qweight.size(1) / block_n;
597 + int64_t Kc = K / _block_k;
598 + int64_t buffer_size_nbytes =
599 + _block_k * block_n / 2 + block_n * sizeof(int32_t);
600 + auto blocked_weight =
601 + at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
602 + auto blocked_scales =
603 + at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
604 + auto blocked_qzeros =
605 + at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
606 + for (int i = 0; i < _qweight.size(0); i++) {
607 + auto res_ = convert_int4_weight_packed_with_compensation(
608 + _qweight[i], _scales[i], _qzeros[i]);
609 + blocked_weight[i] = std::get<0>(res_);
610 + blocked_scales[i] = std::get<1>(res_);
611 + blocked_qzeros[i] = std::get<2>(res_);
612 + }
613 + _qweight = blocked_weight;
614 + _scales = blocked_scales;
615 + _qzeros = blocked_qzeros;
616 + } else {
617 + auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales,
618 + _qzeros);
619 + _qweight = std::get<0>(res_);
620 + _scales = std::get<1>(res_);
621 + _qzeros = std::get<2>(res_);
622 + }
623 +
624 + return std::make_tuple(_qweight, _qzeros, _scales);
625 +}
626 +
627 +at::Tensor int4_scaled_mm_cpu_with_quant(const at::Tensor& input,
628 + const at::Tensor& weight,
629 + const at::Tensor& weight_scales,
630 + const at::Tensor& weight_qzeros,
631 + const std::optional<at::Tensor>& bias,
632 + at::ScalarType output_dtype) {
633 + RECORD_FUNCTION("vllm::int4_scaled_mm_cpu_with_quant",
634 + std::vector<c10::IValue>({input, weight}));
635 +
636 + int64_t M_a = input.size(0);
637 + int64_t K_a = input.size(1);
638 + int64_t lda = input.stride(0);
639 +
640 + const auto st = input.scalar_type();
641 + TORCH_CHECK(
642 + st == at::kBFloat16 || st == at::kHalf,
643 + "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");
644 +
645 + constexpr bool sym_quant_act = false;
646 + using Tin = typename ActDtype<sym_quant_act>::type;
647 + int64_t act_buffer_size =
648 + M_a * K_a + M_a * sizeof(float) + M_a * sizeof(int32_t);
649 + auto act_buffer =
650 + at::empty({act_buffer_size}, input.options().dtype(at::kByte));
651 + auto Aq_data = act_buffer.data_ptr<uint8_t>();
652 + auto As_data = reinterpret_cast<float*>(Aq_data + M_a * K_a);
653 + auto Azp_data = reinterpret_cast<int32_t*>(As_data + M_a);
654 + fill_val_stub(Azp_data, 128, M_a);
655 +
656 + auto out_sizes = input.sizes().vec();
657 + int64_t N = weight_scales.size(0) * weight_scales.size(-1);
658 + out_sizes.back() = N;
659 + auto output = at::empty(out_sizes, input.options());
660 + int64_t Nc = weight.size(0);
661 + int64_t Kc = weight.size(1);
662 + int64_t _block_k = K_a / Kc;
663 + TORCH_CHECK(N == Nc * BLOCK_N, "DA8W4: weight and input shapes mismatch");
664 + int64_t num_groups = weight_scales.size(1);
665 +
666 + const uint8_t* b_ptr = weight.data_ptr<uint8_t>();
667 + const float* b_scales_ptr = weight_scales.data_ptr<float>();
668 + const int8_t* b_qzeros_ptr = weight_qzeros.data_ptr<int8_t>();
669 + const float* bias_ptr =
670 + bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
671 + int num_threads = at::get_num_threads();
672 + int64_t temp_buffer_size = num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
673 + num_threads * _block_k * BLOCK_N;
674 + auto c_temp_buffer =
675 + at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
676 + float* c_temp_ptr = (float*)((void*)(c_temp_buffer.data_ptr<int8_t>()));
677 + int8_t* dqB_temp_ptr =
678 + (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));
679 +
680 +#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act) \
681 + AT_DISPATCH_FLOATING_TYPES_AND2( \
682 + at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, \
683 + "int4_scaled_mm_cpu", [&] { \
684 + const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>(); \
685 + scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>(); \
686 + at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) { \
687 + for (int64_t m = begin; m < end; ++m) { \
688 + quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m], \
689 + A_data + m * lda, K_a); \
690 + } \
691 + }); \
692 + _da8w4_linear_impl<sym_quant_act, Tin, scalar_t>( \
693 + Aq_data, As_data, Azp_data, b_ptr, b_scales_ptr, b_qzeros_ptr, \
694 + bias_ptr, c_ptr, c_temp_ptr, dqB_temp_ptr, M_a, N, K_a, \
695 + num_groups); \
696 + });
697 +
698 + LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act);
699 +
700 + return output;
701 +}
702 +
703 +namespace {
704 +
705 +template <typename scalar_t>
706 +inline void copy_stub(scalar_t* __restrict__ out,
707 + const float* __restrict__ input, int64_t size) {
708 + using Vec = at::vec::Vectorized<scalar_t>;
709 + using fVec = at::vec::Vectorized<float>;
710 +#pragma GCC unroll 4
711 + for (int64_t d = 0; d < size; d += Vec::size()) {
712 + fVec x0 = fVec::loadu(input + d);
713 + fVec x1 = fVec::loadu(input + d + fVec::size());
714 + Vec res = convert_from_float_ext<scalar_t>(x0, x1);
715 + res.store(out + d);
716 + }
717 +}
718 +
719 +} // anonymous namespace
720 +
721 +template <typename scalar_t>
722 +void tinygemm_kernel(scalar_t* C, float* C_temp, const uint8_t* A,
723 + const float* scales_a, const int32_t* qzeros_a,
724 + const uint8_t* B, const float* scales_b,
725 + const int8_t* qzeros_b, const int32_t* compensation,
726 + int8_t* dqB_tmp, int64_t M, int64_t K, int64_t lda,
727 + int64_t ldc_f, int64_t ldc_s, bool store_out,
728 + bool use_brgemm) {
729 + _dequant_gemm_accum<false, BLOCK_N, BLOCK_N / 2>(
730 + C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation,
731 + dqB_tmp, M, K, lda, ldc_f, use_brgemm);
732 + if (store_out) {
733 + for (int64_t m = 0; m < M; ++m) {
734 + copy_stub<scalar_t>(C + m * ldc_s, C_temp + m * ldc_f, BLOCK_N);
735 + }
736 + }
737 +}
738 +
739 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
740 + template void tinygemm_kernel<TYPE>( \
741 + TYPE * C, float* C_temp, const uint8_t* A, const float* scales_a, \
742 + const int32_t* qzeros_a, const uint8_t* B, const float* scales_b, \
743 + const int8_t* qzeros_b, const int32_t* compensation, int8_t* dqB_tmp, \
744 + int64_t M, int64_t K, int64_t lda, int64_t ldc_f, int64_t ldc_s, \
745 + bool store_out, bool use_brgemm)
746 +
747 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
748 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
749 +
750 +at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
751 + at::Tensor& w_scales,
752 + std::optional<at::Tensor> bias) {
753 + return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias,
754 + x.scalar_type());
755 +}
csrc/cpu/torch_bindings.cpp
+20
@@ -79,6 +79,14 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
79 const std::optional<at::Tensor>& bias,
80 at::ScalarType out_dtype, bool is_vnni);
81
82 +// Adapted from sglang: INT4 W4A8 kernels
83 +std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
84 + at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
85 +
86 +at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
87 + at::Tensor& w_scales,
88 + std::optional<at::Tensor> bias);
89 +
90 torch::Tensor get_scheduler_metadata(
91 const int64_t num_req, const int64_t num_heads_q,
92 const int64_t num_heads_kv, const int64_t head_dim,
@@ -285,6 +293,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
293 "Tensor? bias, ScalarType out_dtype, bool is_vnni) -> Tensor");
294 ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
295 &int8_scaled_mm_with_quant);
296 +
297 + // Adapted from sglang: INT4 W4A8 kernels
298 + ops.def(
299 + "convert_weight_packed_scale_zp(Tensor qweight, Tensor qzeros, "
300 + "Tensor scales) -> (Tensor, Tensor, Tensor)");
301 + ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
302 + &convert_weight_packed_scale_zp);
303 +
304 + ops.def(
305 + "int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
306 + "Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
307 + ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
308 #endif
309
310 // CPU attention kernels
tests/kernels/test_awq_int4_to_int8.py new
+281
@@ -0,0 +1,281 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Unit tests for AWQ INT4 W4A8 GEMM pipeline (SGLang kernel migration).
5 +
6 +Part 1: Weight packing tests
7 + - convert_weight_packed_scale_zp correctness
8 +
9 +Part 2: INT4 W4A8 GEMM tests
10 + - int4_scaled_mm_cpu correctness w.r.t. float reference
11 + - Bias, 3D input, various shapes
12 +
13 +Part 3: create_weights shapes
14 +
15 +cmd:
16 + VLLM_CPU_INT4_W4A8=1 python -m pytest tests/kernels/test_awq_int4_to_int8.py -v -s
17 +"""
18 +
19 +import numpy as np
20 +import pytest
21 +import torch
22 +
23 +from vllm._custom_ops import _supports_cpu_w4a8_int8
24 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
25 + pack_cols,
26 +)
27 +from vllm.platforms import current_platform
28 +
29 +if not current_platform.is_cpu():
30 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
31 +
32 +requires_cpu_w4a8_int8 = pytest.mark.skipif(
33 + not _supports_cpu_w4a8_int8,
34 + reason="Requires vLLM CPU build with SGLang INT4 W4A8 kernels",
35 +)
36 +
37 +
38 +def make_awq_checkpoint_data(K, N, group_size, seed=42):
39 + """Create synthetic AWQ checkpoint data in packed int32 format.
40 +
41 + Returns:
42 + packed_qweight: [K, N//8] int32 (AWQ interleaved + packed)
43 + packed_qzeros: [num_groups, N//8] int32 (AWQ interleaved + packed)
44 + scales: [num_groups, N] float32
45 + float_ref: [K, N] float32, reference dequantized weights
46 + weight_int4_orig: [K, N] int32, original int4 values (0-15)
47 + zeros_int4_orig: [num_groups, N] int32, original zero points (0-15)
48 + """
49 + rng = np.random.RandomState(seed)
50 + num_groups = K // group_size
51 +
52 + weight_int4_orig = torch.from_numpy(
53 + rng.randint(0, 16, size=(K, N)).astype(np.int32)
54 + )
55 + zeros_int4_orig = torch.from_numpy(
56 + rng.randint(0, 16, size=(num_groups, N)).astype(np.int32)
57 + )
58 + scales = torch.from_numpy((rng.randn(num_groups, N) * 0.05).astype(np.float32))
59 +
60 + scales_exp = scales.repeat_interleave(group_size, dim=0)
61 + zeros_exp = zeros_int4_orig.repeat_interleave(group_size, dim=0)
62 + float_ref = (weight_int4_orig.float() - zeros_exp.float()) * scales_exp
63 +
64 + awq_interleave = [0, 2, 4, 6, 1, 3, 5, 7]
65 + weight_interleaved = (
66 + weight_int4_orig.reshape(-1, 8)[:, awq_interleave].reshape(K, N).contiguous()
67 + )
68 + packed_qweight = pack_cols(weight_interleaved, 4, K, N)
69 +
70 + zeros_interleaved = (
71 + zeros_int4_orig.reshape(-1, 8)[:, awq_interleave]
72 + .reshape(num_groups, N)
73 + .contiguous()
74 + )
75 + packed_qzeros = pack_cols(zeros_interleaved, 4, num_groups, N)
76 +
77 + return (
78 + packed_qweight,
79 + packed_qzeros,
80 + scales,
81 + float_ref,
82 + weight_int4_orig,
83 + zeros_int4_orig,
84 + )
85 +
86 +
87 +class TestConvertWeightPackedScaleZp:
88 + """Tests for convert_weight_packed_scale_zp weightpacking."""
89 +
90 + @requires_cpu_w4a8_int8
91 + @pytest.mark.parametrize(
92 + "K,N,group_size",
93 + [
94 + (128, 128, 128),
95 + (256, 256, 128),
96 + (512, 256, 64),
97 + ],
98 + )
99 + def test_packing_output_shapes(self, K, N, group_size):
100 + """Packed outputs should have expected shapes."""
101 + (packed_qweight, packed_qzeros, scales, _, _, _) = make_awq_checkpoint_data(
102 + K, N, group_size
103 + )
104 +
105 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
106 + packed_qweight, packed_qzeros, scales
107 + )
108 +
109 + block_n = 32
110 + Nc = N // block_n
111 +
112 + assert blocked_w.dim() >= 2, (
113 + f"blocked_w should have >= 2 dims, got {blocked_w.dim()}"
114 + )
115 + assert blocked_s.size(0) == Nc, (
116 + f"Expected Nc={Nc} scale blocks, got {blocked_s.size(0)}"
117 + )
118 + assert blocked_zp.size(0) == Nc, (
119 + f"Expected Nc={Nc} qzeros blocks, got {blocked_zp.size(0)}"
120 + )
121 +
122 + print(
123 + f" [PASS] packing shapes K={K}, N={N}, gs={group_size}: "
124 + f"blocked_w={list(blocked_w.shape)}, "
125 + f"blocked_s={list(blocked_s.shape)}, blocked_zp={list(blocked_zp.shape)}"
126 + )
127 +
128 +
129 +class TestInt4ScaledMmCpu:
130 + """Tests for int4_scaled_mm_cpu GEMM kernel."""
131 +
132 + @requires_cpu_w4a8_int8
133 + @pytest.mark.parametrize(
134 + "M,K,N,group_size",
135 + [
136 + (1, 128, 128, 128),
137 + (4, 256, 256, 128),
138 + (16, 512, 256, 64),
139 + (32, 256, 512, 128),
140 + (64, 512, 512, 128),
141 + ],
142 + )
143 + def test_gemm_vs_float_reference(self, M, K, N, group_size):
144 + """INT4 W4A8 GEMM should approximate float matmul."""
145 + (packed_qweight, packed_qzeros, scales, float_ref, _, _) = (
146 + make_awq_checkpoint_data(K, N, group_size)
147 + )
148 +
149 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
150 + packed_qweight, packed_qzeros, scales
151 + )
152 +
153 + x = torch.randn(M, K, dtype=torch.bfloat16)
154 + out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
155 +
156 + ref_out = torch.mm(x.float(), float_ref)
157 +
158 + abs_diff = (out.float() - ref_out).abs()
159 + mean_abs = abs_diff.mean().item()
160 + pct95 = torch.quantile(abs_diff, 0.95).item()
161 + ref_mag = ref_out.abs().mean().item() + 1e-6
162 + mean_rel = mean_abs / ref_mag
163 +
164 + assert mean_rel < 0.05, (
165 + f"Mean relative error {mean_rel:.4f} exceeds 5% threshold"
166 + )
167 + assert pct95 < ref_mag * 0.15, (
168 + f"95th-pctile abs_diff {pct95:.4f} exceeds 15% of ref magnitude"
169 + )
170 + print(f" [PASS] INT4 GEMM correct: M={M}, K={K}, N={N}")
171 +
172 + @requires_cpu_w4a8_int8
173 + @pytest.mark.parametrize("M", [1, 8, 32])
174 + def test_gemm_with_bias(self, M):
175 + """INT4 W4A8 GEMM with bias should match reference."""
176 + K, N, group_size = 256, 128, 128
177 + (packed_qweight, packed_qzeros, scales, float_ref, _, _) = (
178 + make_awq_checkpoint_data(K, N, group_size)
179 + )
180 +
181 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
182 + packed_qweight, packed_qzeros, scales
183 + )
184 +
185 + bias = torch.randn(N, dtype=torch.float32)
186 + x = torch.randn(M, K, dtype=torch.bfloat16)
187 +
188 + out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, bias)
189 +
190 + ref_out = torch.mm(x.float(), float_ref) + bias
191 + abs_diff = (out.float() - ref_out).abs()
192 + mean_abs = abs_diff.mean().item()
193 + ref_mag = ref_out.abs().mean().item() + 1e-6
194 + mean_rel = mean_abs / ref_mag
195 + assert mean_rel < 0.05, (
196 + f"Mean relative error {mean_rel:.4f} with bias exceeds 5%"
197 + )
198 + print(f" [PASS] INT4 GEMM with bias: M={M}")
199 +
200 + @requires_cpu_w4a8_int8
201 + def test_gemm_3d_input(self):
202 + """apply() reshapes 3D input [B, S, K] -> [B*S, K] -> back to 3D."""
203 + K, N, group_size = 256, 128, 128
204 + (packed_qweight, packed_qzeros, scales, float_ref, _, _) = (
205 + make_awq_checkpoint_data(K, N, group_size)
206 + )
207 +
208 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
209 + packed_qweight, packed_qzeros, scales
210 + )
211 +
212 + B, S = 2, 8
213 + x_3d = torch.randn(B, S, K, dtype=torch.bfloat16)
214 + x_2d = x_3d.reshape(-1, K)
215 +
216 + out_2d = torch.ops._C.int4_scaled_mm_cpu(
217 + x_2d, blocked_w, blocked_zp, blocked_s, None
218 + )
219 + out_3d = out_2d.reshape(B, S, N)
220 +
221 + ref_out = torch.mm(x_2d.float(), float_ref).reshape(B, S, N)
222 +
223 + assert out_3d.shape == (B, S, N)
224 + abs_diff = (out_3d.float() - ref_out).abs()
225 + mean_abs = abs_diff.mean().item()
226 + ref_mag = ref_out.abs().mean().item() + 1e-6
227 + mean_rel = mean_abs / ref_mag
228 +
229 + assert mean_rel < 0.05, f"Mean relative error {mean_rel:.4f} for 3D exceeds 5%"
230 + print(f" [PASS] 3D input [{B},{S},{K}] -> output [{B},{S},{N}]")
231 +
232 + @requires_cpu_w4a8_int8
233 + def test_gemm_fp16_input(self):
234 + """INT4 GEMM should also work with fp16 input."""
235 + K, N, group_size, M = 256, 256, 128, 8
236 + (packed_qweight, packed_qzeros, scales, float_ref, _, _) = (
237 + make_awq_checkpoint_data(K, N, group_size)
238 + )
239 +
240 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
241 + packed_qweight, packed_qzeros, scales
242 + )
243 +
244 + x = torch.randn(M, K, dtype=torch.float16)
245 + out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
246 +
247 + ref_out = torch.mm(x.float(), float_ref)
248 + abs_diff = (out.float() - ref_out).abs()
249 + ref_mag = ref_out.abs().mean().item() + 1e-6
250 + mean_rel = abs_diff.mean().item() / ref_mag
251 +
252 + assert mean_rel < 0.05, (
253 + f"Mean relative error {mean_rel:.4f} for fp16 exceeds 5%"
254 + )
255 + print(f" [PASS] fp16 input M={M}, K={K}, N={N}")
256 +
257 +
258 +class TestCreateWeightsUnchanged:
259 + """Create_weights should still produce correct int4 placeholder shapes."""
260 +
261 + @pytest.mark.parametrize(
262 + "K,N,group_size",
263 + [
264 + (128, 128, 128),
265 + (256, 256, 128),
266 + (512, 256, 64),
267 + ],
268 + )
269 + def test_int4_placeholder_shapes(self, K, N, group_size):
270 + """Verify qweight, qzeros, scales shapes."""
271 + pack_factor = 8
272 + num_groups = K // group_size
273 +
274 + qweight = torch.empty(K, N // pack_factor, dtype=torch.int32)
275 + qzeros = torch.empty(num_groups, N // pack_factor, dtype=torch.int32)
276 + scales = torch.empty(num_groups, N, dtype=torch.bfloat16)
277 +
278 + assert qweight.shape == (K, N // pack_factor)
279 + assert qzeros.shape == (num_groups, N // pack_factor)
280 + assert scales.shape == (num_groups, N)
281 + print(f" [PASS] create_weights shapes: K={K}, N={N}, gs={group_size}")
vllm/_custom_ops.py
+32
@@ -2967,6 +2967,38 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
2967 return torch.empty((M, N), dtype=out_dtype)
2968
2969
2970 +if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
2971 +
2972 + @register_fake("_C::convert_weight_packed_scale_zp")
2973 + def convert_weight_packed_scale_zp_fake(
2974 + qweight: torch.Tensor,
2975 + qzeros: torch.Tensor,
2976 + scales: torch.Tensor,
2977 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
2978 + return (
2979 + torch.empty_like(qweight),
2980 + torch.empty_like(qzeros),
2981 + torch.empty_like(scales),
2982 + )
2983 +
2984 +
2985 +if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
2986 +
2987 + @register_fake("_C::int4_scaled_mm_cpu")
2988 + def int4_scaled_mm_cpu_fake(
2989 + x: torch.Tensor,
2990 + w: torch.Tensor,
2991 + w_zeros: torch.Tensor,
2992 + w_scales: torch.Tensor,
2993 + bias: torch.Tensor | None,
2994 + ) -> torch.Tensor:
2995 + N = w_scales.size(0) * w_scales.size(-1)
2996 + return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device)
2997 +
2998 +
2999 +_supports_cpu_w4a8_int8 = bool(hasattr(torch.ops._C, "convert_weight_packed_scale_zp"))
3000 +
3001 +
3002 class CPUDNNLGEMMHandler:
3003 def __init__(self) -> None:
3004 self.handler_tensor: torch.Tensor | None = None
vllm/envs.py
+3
@@ -52,6 +52,7 @@ if TYPE_CHECKING:
52 VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
53 VLLM_CPU_SGL_KERNEL: bool = False
54 VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
55 + VLLM_CPU_INT4_W4A8: bool = True
56 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
57 VLLM_XLA_CHECK_RECOMPILATION: bool = False
58 VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
@@ -728,6 +729,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
729 "VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool(
730 int(os.getenv("VLLM_ZENTORCH_WEIGHT_PREPACK", "1"))
731 ),
732 + # (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ.
733 + "VLLM_CPU_INT4_W4A8": lambda: bool(int(os.getenv("VLLM_CPU_INT4_W4A8", "1"))),
734 # If the env var is set, Ray Compiled Graph uses the specified
735 # channel type to communicate between workers belonging to
736 # different pipeline-parallel stages.
vllm/model_executor/layers/quantization/cpu_wna16.py
+58 -7
@@ -7,9 +7,8 @@ import torch
7 from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
8 from transformers import PretrainedConfig
9
10 -from vllm._custom_ops import (
11 - cpu_gemm_wna16,
12 -)
10 +import vllm.envs as envs
11 +from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.linear import (
14 LinearBase,
@@ -230,7 +229,14 @@ class CPUAWQLinearMethod(LinearMethodBase):
229 layer.register_parameter("scales", scales)
230
231 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
233 - torch.set_printoptions(profile="full", linewidth=5000, sci_mode=False)
232 + layer.use_w4a8 = envs.VLLM_CPU_INT4_W4A8 and torch.cpu._is_amx_tile_supported()
233 + if layer.use_w4a8:
234 + self._process_weights_sglang_int4(layer)
235 + else:
236 + self._process_weights_woq(layer)
237 +
238 + def _process_weights_woq(self, layer: torch.nn.Module) -> None:
239 + """Original WOQ int4 repack path."""
240 packed_weight = layer.qweight.data
241 packed_zeros = layer.qzeros.data
242 group_num = packed_zeros.size(0)
@@ -266,8 +272,6 @@ class CPUAWQLinearMethod(LinearMethodBase):
272 )
273
274 zeros = pack_cols(zeros, bits, group_num, output_size).contiguous()
269 - # make 16 output channel as a block and transpose to
270 - # the make the block contiguous
275 weight = pack_cols(weight, bits, input_size, output_size)
276 weight = (
277 weight.view(input_size, -1, 16 // pack_factor)
@@ -278,13 +282,40 @@ class CPUAWQLinearMethod(LinearMethodBase):
282 layer.qweight.data = weight
283 layer.qzeros.data = zeros
284
285 + def _process_weights_sglang_int4(self, layer: torch.nn.Module) -> None:
286 + """SGLang INT4 W4A8 path: pack int4 weights with VNNI reordering."""
287 + packed_weight = layer.qweight.data
288 + packed_zeros = layer.qzeros.data
289 + scales = layer.scales.data
290 + blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
291 + packed_weight, packed_zeros, scales
292 + )
293 +
294 + layer.packed_weight = blocked_w
295 + layer.packed_qzeros = blocked_zp
296 + layer.packed_scales = blocked_s
297 + layer.qweight = None
298 + layer.qzeros = None
299 + layer.scales = None
300 +
301 def apply(
302 self,
303 layer: torch.nn.Module,
304 x: torch.Tensor,
305 bias: torch.Tensor | None = None,
306 ) -> torch.Tensor:
287 - x = cpu_gemm_wna16(
307 + if layer.use_w4a8:
308 + return self._apply_sglang_int4(layer, x, bias)
309 + return self._apply_woq(layer, x, bias)
310 +
311 + def _apply_woq(
312 + self,
313 + layer: torch.nn.Module,
314 + x: torch.Tensor,
315 + bias: torch.Tensor | None = None,
316 + ) -> torch.Tensor:
317 + """Original WOQ int4 GEMM path."""
318 + x = ops.cpu_gemm_wna16(
319 input=x,
320 q_weight=layer.qweight,
321 scales=layer.scales,
@@ -296,6 +327,26 @@ class CPUAWQLinearMethod(LinearMethodBase):
327 )
328 return x
329
330 + def _apply_sglang_int4(
331 + self,
332 + layer: torch.nn.Module,
333 + x: torch.Tensor,
334 + bias: torch.Tensor | None = None,
335 + ) -> torch.Tensor:
336 + """SGLang INT4 W4A8 GEMM path."""
337 + x_shape = x.shape
338 + x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
339 +
340 + out = torch.ops._C.int4_scaled_mm_cpu(
341 + x_2d,
342 + layer.packed_weight,
343 + layer.packed_qzeros,
344 + layer.packed_scales,
345 + bias,
346 + )
347 + out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
348 + return out
349 +
350
351 def _get_isa_hint(dtype: torch.dtype) -> str:
352 supports_amx = torch.cpu._is_amx_tile_supported()