@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e515668edf
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 15, 11:00 UTC

[Hardware][Power] Enable compressed tensor W8A8 INT8 quantization for POWER (#17153)

Signed-off-by: Akash Kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash Kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: mgoin <mgoin64@gmail.com>

Akash kaothalkar committed May 8, 2025 at 11:05 UTC e515668edf510d86a0543ac5d7981dd91b2026d7
4 files changed +669 -5
cmake/cpu_extension.cmake
+32 -1
@@ -167,6 +167,33 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
167
168 FetchContent_MakeAvailable(oneDNN)
169
170 + list(APPEND LIBS dnnl)
171 +elseif(POWER10_FOUND)
172 + FetchContent_Declare(
173 + oneDNN
174 + GIT_REPOSITORY https://github.com/oneapi-src/oneDNN.git
175 + GIT_TAG v3.7.2
176 + GIT_PROGRESS TRUE
177 + GIT_SHALLOW TRUE
178 + )
179 +
180 + set(ONEDNN_LIBRARY_TYPE "STATIC")
181 + set(ONEDNN_BUILD_DOC "OFF")
182 + set(ONEDNN_BUILD_EXAMPLES "OFF")
183 + set(ONEDNN_BUILD_TESTS "OFF")
184 + set(ONEDNN_ENABLE_WORKLOAD "INFERENCE")
185 + set(ONEDNN_ENABLE_PRIMITIVE "MATMUL;REORDER")
186 + set(ONEDNN_BUILD_GRAPH "OFF")
187 + set(ONEDNN_ENABLE_JIT_PROFILING "OFF")
188 + set(ONEDNN_ENABLE_ITT_TASKS "OFF")
189 + set(ONEDNN_ENABLE_MAX_CPU_ISA "OFF")
190 + set(ONEDNN_ENABLE_CPU_ISA_HINTS "OFF")
191 + set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
192 +
193 + set(DNNL_CPU_RUNTIME "OMP")
194 +
195 + FetchContent_MakeAvailable(oneDNN)
196 +
197 list(APPEND LIBS dnnl)
198 endif()
199
@@ -197,6 +224,10 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
224 "csrc/cpu/quant.cpp"
225 "csrc/cpu/shm.cpp"
226 ${VLLM_EXT_SRC})
227 +elseif(POWER10_FOUND)
228 + set(VLLM_EXT_SRC
229 + "csrc/cpu/quant.cpp"
230 + ${VLLM_EXT_SRC})
231 endif()
232
233 #
@@ -214,4 +245,4 @@ define_gpu_extension_target(
245 WITH_SOABI
246 )
247
217 -message(STATUS "Enabling C extension.")
\ No newline at end of file
248 +message(STATUS "Enabling C extension.")
csrc/cpu/cpu_types_vsx.hpp
+265
@@ -4,6 +4,7 @@
4
5 #include <altivec.h>
6 #include <cmath>
7 +#include <algorithm>
8 #include <torch/all.h>
9
10 namespace vec_op {
@@ -62,6 +63,10 @@ typedef struct f32x4x4_t {
63 __vector float val[4];
64 } f32x4x4_t;
65
66 +typedef struct i32x4x4_t {
67 + __vector int32_t val[4];
68 +} i32x4x4_t;
69 +
70 struct FP32Vec8;
71 struct FP32Vec16;
72
@@ -98,6 +103,28 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
103 vec_xst(reg.val[0], 0, (signed short*)ptr);
104 vec_xst(reg.val[1], 16, (signed short*)ptr);
105 }
106 +
107 + void save(void* ptr, const int elem_num) const {
108 + const int clamped_elem = std::max(0, std::min(elem_num, 16));
109 +
110 + // Calculate elements to store in each 128-bit part (8 elements each)
111 + const int elements_val0 = std::min(clamped_elem, 8);
112 + const int elements_val1 = std::max(clamped_elem - 8, 0);
113 +
114 + // Convert elements to bytes (2 bytes per element)
115 + const size_t bytes_val0 = elements_val0 * sizeof(signed short);
116 + const size_t bytes_val1 = elements_val1 * sizeof(signed short);
117 +
118 + signed short* dest = static_cast<signed short*>(ptr);
119 + // Store the first part using vec_xst_len
120 + if (bytes_val0 > 0) {
121 + vec_xst_len(reg.val[0], dest, bytes_val0);
122 + }
123 + // Store the second part if needed
124 + if (bytes_val1 > 0) {
125 + vec_xst_len(reg.val[1], dest + elements_val0, bytes_val1);
126 + }
127 + }
128 };
129
130 const static __vector signed short zero = vec_splats((signed short)0);
@@ -257,6 +284,64 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
284 }
285 };
286
287 +struct INT32Vec16 : public Vec<INT32Vec16> {
288 + constexpr static int VEC_ELEM_NUM = 16;
289 + union AliasReg {
290 + i32x4x4_t reg;
291 + int32_t values[VEC_ELEM_NUM];
292 + };
293 +
294 + i32x4x4_t reg;
295 +
296 + explicit INT32Vec16(const void* data_ptr) {
297 + reg.val[0] = vec_xl(0, reinterpret_cast<const __vector int32_t*>(data_ptr));
298 + reg.val[1] =
299 + vec_xl(16, reinterpret_cast<const __vector int32_t*>(data_ptr));
300 + reg.val[2] =
301 + vec_xl(32, reinterpret_cast<const __vector int32_t*>(data_ptr));
302 + reg.val[3] =
303 + vec_xl(48, reinterpret_cast<const __vector int32_t*>(data_ptr));
304 + }
305 +
306 + void save(int32_t* ptr) const {
307 + vec_xst(reg.val[0], 0, reinterpret_cast<__vector int32_t*>(ptr));
308 + vec_xst(reg.val[1], 16, reinterpret_cast<__vector int32_t*>(ptr));
309 + vec_xst(reg.val[2], 32, reinterpret_cast<__vector int32_t*>(ptr));
310 + vec_xst(reg.val[3], 48, reinterpret_cast<__vector int32_t*>(ptr));
311 + }
312 +
313 + void save(int32_t* ptr, const int elem_num) const {
314 + const int elements_in_chunk1 =
315 + (elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0;
316 + const int elements_in_chunk2 =
317 + (elem_num > 4) ? ((elem_num >= 8) ? 4 : elem_num - 4) : 0;
318 + const int elements_in_chunk3 =
319 + (elem_num > 8) ? ((elem_num >= 12) ? 4 : elem_num - 8) : 0;
320 + const int elements_in_chunk4 =
321 + (elem_num > 12) ? ((elem_num >= 16) ? 4 : elem_num - 12) : 0;
322 +
323 + const size_t bytes_chunk1 =
324 + static_cast<size_t>(elements_in_chunk1 * sizeof(int32_t));
325 + const size_t bytes_chunk2 =
326 + static_cast<size_t>(elements_in_chunk2 * sizeof(int32_t));
327 + const size_t bytes_chunk3 =
328 + static_cast<size_t>(elements_in_chunk3 * sizeof(int32_t));
329 + const size_t bytes_chunk4 =
330 + static_cast<size_t>(elements_in_chunk4 * sizeof(int32_t));
331 +
332 + vec_xst_len(reg.val[0], reinterpret_cast<int32_t*>(ptr), bytes_chunk1);
333 + vec_xst_len(reg.val[1],
334 + reinterpret_cast<int32_t*>(reinterpret_cast<char*>(ptr) + 16),
335 + bytes_chunk2);
336 + vec_xst_len(reg.val[2],
337 + reinterpret_cast<int32_t*>(reinterpret_cast<char*>(ptr) + 32),
338 + bytes_chunk3);
339 + vec_xst_len(reg.val[3],
340 + reinterpret_cast<int32_t*>(reinterpret_cast<char*>(ptr) + 48),
341 + bytes_chunk4);
342 + }
343 +};
344 +
345 struct FP32Vec16 : public Vec<FP32Vec16> {
346 constexpr static int VEC_ELEM_NUM = 16;
347 union AliasReg {
@@ -319,6 +404,13 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
404
405 explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
406
407 + explicit FP32Vec16(const INT32Vec16& v) {
408 + reg.val[0] = vec_ctf(v.reg.val[0], 0);
409 + reg.val[1] = vec_ctf(v.reg.val[1], 0);
410 + reg.val[2] = vec_ctf(v.reg.val[2], 0);
411 + reg.val[3] = vec_ctf(v.reg.val[3], 0);
412 + }
413 +
414 FP32Vec16 operator*(const FP32Vec16& b) const {
415 return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
416 vec_mul(reg.val[1], b.reg.val[1]),
@@ -347,6 +439,117 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
439 vec_div(reg.val[3], b.reg.val[3])}));
440 }
441
442 + FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
443 + return FP32Vec16(f32x4x4_t(
444 + {vec_min(max.reg.val[0], vec_max(min.reg.val[0], reg.val[0])),
445 + vec_min(max.reg.val[1], vec_max(min.reg.val[1], reg.val[1])),
446 + vec_min(max.reg.val[2], vec_max(min.reg.val[2], reg.val[2])),
447 + vec_min(max.reg.val[3], vec_max(min.reg.val[3], reg.val[3]))}));
448 + }
449 +
450 + FP32Vec16 max(const FP32Vec16& b) const {
451 + return FP32Vec16(f32x4x4_t({vec_max(reg.val[0], b.reg.val[0]),
452 + vec_max(reg.val[1], b.reg.val[1]),
453 + vec_max(reg.val[2], b.reg.val[2]),
454 + vec_max(reg.val[3], b.reg.val[3])}));
455 + }
456 +
457 + FP32Vec16 max(const FP32Vec16& b, int elem_num) const {
458 + FP32Vec16 result;
459 +
460 + // Create a vector of element indices for each chunk
461 + __vector unsigned int indices = {0, 1, 2, 3};
462 + __vector unsigned int elem_num_vec =
463 + vec_splats(static_cast<unsigned int>(elem_num));
464 +
465 + // Compute masks for each chunk
466 + __vector unsigned int chunk_offset0 = {0, 0, 0,
467 + 0}; // Chunk 0: Elements 0-3
468 + __vector unsigned int chunk_offset1 = {4, 4, 4,
469 + 4}; // Chunk 1: Elements 4-7
470 + __vector unsigned int chunk_offset2 = {8, 8, 8,
471 + 8}; // Chunk 2: Elements 8-11
472 + __vector unsigned int chunk_offset3 = {12, 12, 12,
473 + 12}; // Chunk 3: Elements 12-15
474 +
475 + // Compute masks for each chunk
476 + __vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec);
477 + __vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec);
478 + __vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec);
479 + __vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec);
480 +
481 + // Apply masks to compute the result for each chunk
482 + result.reg.val[0] = vec_sel(this->reg.val[0],
483 + vec_max(this->reg.val[0], b.reg.val[0]), mask0);
484 + result.reg.val[1] = vec_sel(this->reg.val[1],
485 + vec_max(this->reg.val[1], b.reg.val[1]), mask1);
486 + result.reg.val[2] = vec_sel(this->reg.val[2],
487 + vec_max(this->reg.val[2], b.reg.val[2]), mask2);
488 + result.reg.val[3] = vec_sel(this->reg.val[3],
489 + vec_max(this->reg.val[3], b.reg.val[3]), mask3);
490 +
491 + return FP32Vec16(result.reg);
492 + }
493 +
494 + FP32Vec16 min(const FP32Vec16& b) const {
495 + return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
496 + vec_min(reg.val[1], b.reg.val[1]),
497 + vec_min(reg.val[2], b.reg.val[2]),
498 + vec_min(reg.val[3], b.reg.val[3])}));
499 + }
500 +
501 + FP32Vec16 min(const FP32Vec16& b, int elem_num) const {
502 + FP32Vec16 result;
503 +
504 + vector unsigned int indices = {0, 1, 2, 3};
505 + vector unsigned int elem_num_vec =
506 + vec_splats(static_cast<unsigned int>(elem_num));
507 +
508 + vector unsigned int chunk_offset0 = {0, 0, 0, 0};
509 + vector unsigned int chunk_offset1 = {4, 4, 4, 4};
510 + vector unsigned int chunk_offset2 = {8, 8, 8, 8};
511 + vector unsigned int chunk_offset3 = {12, 12, 12, 12};
512 +
513 + vector bool int mask0 = vec_cmplt(indices + chunk_offset0, elem_num_vec);
514 + vector bool int mask1 = vec_cmplt(indices + chunk_offset1, elem_num_vec);
515 + vector bool int mask2 = vec_cmplt(indices + chunk_offset2, elem_num_vec);
516 + vector bool int mask3 = vec_cmplt(indices + chunk_offset3, elem_num_vec);
517 +
518 + result.reg.val[0] = vec_sel(this->reg.val[0],
519 + vec_min(this->reg.val[0], b.reg.val[0]), mask0);
520 + result.reg.val[1] = vec_sel(this->reg.val[1],
521 + vec_min(this->reg.val[1], b.reg.val[1]), mask1);
522 + result.reg.val[2] = vec_sel(this->reg.val[2],
523 + vec_min(this->reg.val[2], b.reg.val[2]), mask2);
524 + result.reg.val[3] = vec_sel(this->reg.val[3],
525 + vec_min(this->reg.val[3], b.reg.val[3]), mask3);
526 +
527 + return FP32Vec16(result.reg);
528 + }
529 +
530 + FP32Vec16 abs() const {
531 + return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]),
532 + vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
533 + }
534 +
535 + float reduce_max() {
536 + __vector float max01 = vec_max(reg.val[0], reg.val[1]);
537 + __vector float max23 = vec_max(reg.val[2], reg.val[3]);
538 + __vector float max_all = vec_max(max01, max23);
539 + __vector float temp = vec_max(max_all, vec_sld(max_all, max_all, 8));
540 + temp = vec_max(temp, vec_sld(temp, temp, 4));
541 + return vec_extract(temp, 0);
542 + }
543 +
544 + float reduce_min() {
545 + __vector float min01 = vec_min(reg.val[0], reg.val[1]);
546 + __vector float min23 = vec_min(reg.val[2], reg.val[3]);
547 + __vector float min_all = vec_min(min01, min23);
548 + __vector float temp = vec_min(min_all, vec_sld(min_all, min_all, 8));
549 + temp = vec_min(temp, vec_sld(temp, temp, 4));
550 + return vec_extract(temp, 0);
551 + }
552 +
553 float reduce_sum() const {
554 AliasReg ar;
555 ar.reg = reg;
@@ -377,6 +580,68 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
580 vec_xst(reg.val[2], 32, ptr);
581 vec_xst(reg.val[3], 48, ptr);
582 }
583 +
584 + void save(float* ptr, const int elem_num) const {
585 + const int elements_in_chunk1 =
586 + (elem_num >= 0) ? ((elem_num >= 4) ? 4 : elem_num) : 0;
587 + const int elements_in_chunk2 =
588 + (elem_num > 4) ? ((elem_num >= 8) ? 4 : elem_num - 4) : 0;
589 + const int elements_in_chunk3 =
590 + (elem_num > 8) ? ((elem_num >= 12) ? 4 : elem_num - 8) : 0;
591 + const int elements_in_chunk4 =
592 + (elem_num > 12) ? ((elem_num >= 16) ? 4 : elem_num - 12) : 0;
593 +
594 + const size_t bytes_chunk1 =
595 + static_cast<size_t>(elements_in_chunk1 * sizeof(float));
596 + const size_t bytes_chunk2 =
597 + static_cast<size_t>(elements_in_chunk2 * sizeof(float));
598 + const size_t bytes_chunk3 =
599 + static_cast<size_t>(elements_in_chunk3 * sizeof(float));
600 + const size_t bytes_chunk4 =
601 + static_cast<size_t>(elements_in_chunk4 * sizeof(float));
602 +
603 + vec_xst_len(reg.val[0], ptr, bytes_chunk1);
604 + vec_xst_len(reg.val[1],
605 + reinterpret_cast<float*>(reinterpret_cast<char*>(ptr) + 16),
606 + bytes_chunk2);
607 + vec_xst_len(reg.val[2],
608 + reinterpret_cast<float*>(reinterpret_cast<char*>(ptr) + 32),
609 + bytes_chunk3);
610 + vec_xst_len(reg.val[3],
611 + reinterpret_cast<float*>(reinterpret_cast<char*>(ptr) + 48),
612 + bytes_chunk4);
613 + }
614 +};
615 +
616 +struct INT8Vec16 : public Vec<INT8Vec16> {
617 + constexpr static int VEC_NUM_ELEM = 16; // 128 bits / 8 bits = 16
618 +
619 + union AliasReg {
620 + __vector signed char reg;
621 + int8_t values[VEC_NUM_ELEM];
622 + };
623 +
624 + __vector signed char reg;
625 +
626 + explicit INT8Vec16(const FP32Vec16& vec) {
627 + __vector signed int ret[4];
628 + ret[0] = vec_cts(vec.reg.val[0], 0);
629 + ret[1] = vec_cts(vec.reg.val[1], 0);
630 + ret[2] = vec_cts(vec.reg.val[2], 0);
631 + ret[3] = vec_cts(vec.reg.val[3], 0);
632 +
633 + __vector signed short packed1 = vec_packs(ret[0], ret[1]);
634 + __vector signed short packed2 = vec_packs(ret[2], ret[3]);
635 +
636 + reg = vec_packs(packed1, packed2);
637 + }
638 +
639 + void save(void* ptr) const {
640 + *reinterpret_cast<__vector signed char*>(ptr) = reg;
641 + }
642 + void save(signed char* ptr, const int elem_num) {
643 + vec_xst_len(reg, ptr, static_cast<size_t>(elem_num));
644 + }
645 };
646
647 template <typename T>
csrc/cpu/quant.cpp
+337 -4
@@ -239,6 +239,280 @@ void static_quant_epilogue(const float* input, scalar_t* output,
239 }
240 }
241
242 +template <bool AZP, bool PerChannel, bool Bias, typename scalar_t>
243 +void dynamic_quant_epilogue(const float* input, scalar_t* output,
244 + const float* a_scale, const float* b_scale,
245 + const int32_t* azp, const int32_t* azp_adj,
246 + const scalar_t* bias, const int num_tokens,
247 + const int hidden_size) {
248 + CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue)
249 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
250 + using azp_adj_load_vec_t =
251 + typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
252 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
253 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
254 +
255 + #pragma omp parallel for
256 + for (int i = 0; i < num_tokens; ++i) {
257 + int j = 0;
258 + cvt_vec_t token_scale_vec(a_scale[i]);
259 + cvt_vec_t token_zp_scale_vec;
260 + if constexpr (AZP) {
261 + float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
262 + if constexpr (!PerChannel) {
263 + zp_scale_val *= *b_scale;
264 + }
265 + token_zp_scale_vec = cvt_vec_t(zp_scale_val);
266 + }
267 +
268 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
269 + cvt_vec_t elems_fp32(input + i * hidden_size + j);
270 + elems_fp32 = elems_fp32 * token_scale_vec;
271 +
272 + if constexpr (AZP) {
273 + azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
274 + cvt_vec_t azp_adj_fp32(azp_adj_vec);
275 + azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
276 +
277 + if constexpr (PerChannel) {
278 + cvt_vec_t b_scale_vec(b_scale + j);
279 + azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
280 + }
281 +
282 + elems_fp32 = elems_fp32 - azp_adj_fp32;
283 + }
284 +
285 + if constexpr (Bias) {
286 + load_vec_t bias_vec(bias + j);
287 + cvt_vec_t bias_vec_fp32(bias_vec);
288 + elems_fp32 = elems_fp32 + bias_vec_fp32;
289 + }
290 +
291 + load_vec_t elems_out(elems_fp32);
292 + elems_out.save(output + i * hidden_size + j);
293 + }
294 +
295 + cvt_vec_t elems_fp32(input + i * hidden_size + j);
296 + elems_fp32 = elems_fp32 * token_scale_vec;
297 +
298 + if constexpr (AZP) {
299 + azp_adj_load_vec_t azp_adj_vec(azp_adj + j);
300 + cvt_vec_t azp_adj_fp32(azp_adj_vec);
301 + azp_adj_fp32 = azp_adj_fp32 * token_zp_scale_vec;
302 +
303 + if constexpr (PerChannel) {
304 + cvt_vec_t b_scale_vec(b_scale + j);
305 + azp_adj_fp32 = azp_adj_fp32 * b_scale_vec;
306 + }
307 +
308 + elems_fp32 = elems_fp32 - azp_adj_fp32;
309 + }
310 +
311 + if constexpr (Bias) {
312 + load_vec_t bias_vec(bias + j);
313 + cvt_vec_t bias_vec_fp32(bias_vec);
314 + elems_fp32 = elems_fp32 + bias_vec_fp32;
315 + }
316 +
317 + load_vec_t elems_out(elems_fp32);
318 + elems_out.save(output + i * hidden_size + j, hidden_size - j);
319 + }
320 +}
321 +#elif defined(__powerpc64__)
322 +template <bool AZP, typename scalar_t>
323 +void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
324 + const float* scale, const int32_t* azp,
325 + const int num_tokens,
326 + const int hidden_size) {
327 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
328 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
329 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
330 +
331 + constexpr float i8_min =
332 + static_cast<float>(std::numeric_limits<int8_t>::min());
333 + constexpr float i8_max =
334 + static_cast<float>(std::numeric_limits<int8_t>::max());
335 +
336 + const cvt_vec_t inv_scale(1.0 / *scale);
337 + const cvt_vec_t i8_min_vec(i8_min);
338 + const cvt_vec_t i8_max_vec(i8_max);
339 +
340 + cvt_vec_t zp_vec;
341 + if constexpr (AZP) {
342 + zp_vec = cvt_vec_t(static_cast<float>(*azp));
343 + }
344 + #pragma omp parallel for
345 + for (int i = 0; i < num_tokens; ++i) {
346 + int j = 0;
347 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
348 + load_vec_t elems(input + i * hidden_size + j);
349 + cvt_vec_t elems_fp32(elems);
350 + elems_fp32 = elems_fp32 * inv_scale;
351 + if constexpr (AZP) {
352 + elems_fp32 = elems_fp32 + zp_vec;
353 + }
354 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
355 + vec_op::INT8Vec16 elems_int8(elems_fp32);
356 + elems_int8.save(output + i * hidden_size + j);
357 + }
358 + load_vec_t elems(input + i * hidden_size + j);
359 + cvt_vec_t elems_fp32(elems);
360 + elems_fp32 = elems_fp32 * inv_scale;
361 +
362 + if constexpr (AZP) {
363 + elems_fp32 = elems_fp32 + zp_vec;
364 + }
365 +
366 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
367 + vec_op::INT8Vec16 elems_int8(elems_fp32);
368 + elems_int8.save(output + i * hidden_size + j, hidden_size - j);
369 + }
370 +}
371 +template <bool AZP, typename scalar_t>
372 +void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
373 + float* scale, int32_t* azp,
374 + const int num_tokens,
375 + const int hidden_size) {
376 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
377 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
378 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
379 +
380 + constexpr float i8_min =
381 + static_cast<float>(std::numeric_limits<int8_t>::min());
382 + constexpr float i8_max =
383 + static_cast<float>(std::numeric_limits<int8_t>::max());
384 + const cvt_vec_t i8_min_vec(i8_min);
385 + const cvt_vec_t i8_max_vec(i8_max);
386 +
387 + #pragma omp parallel for
388 + for (int i = 0; i < num_tokens; ++i) {
389 + cvt_vec_t max_value(std::numeric_limits<float>::lowest());
390 + cvt_vec_t min_value(std::numeric_limits<float>::max());
391 + {
392 + int j = 0;
393 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
394 + load_vec_t elems(input + i * hidden_size + j);
395 + cvt_vec_t elems_fp32(elems);
396 + if constexpr (AZP) {
397 + max_value = max_value.max(elems_fp32);
398 + min_value = min_value.min(elems_fp32);
399 + } else {
400 + max_value = max_value.max(elems_fp32.abs());
401 + }
402 + }
403 +
404 + load_vec_t elems(input + i * hidden_size + j);
405 + cvt_vec_t elems_fp32(elems);
406 +
407 + if (j + vec_elem_num == hidden_size) {
408 + if constexpr (AZP) {
409 + max_value = max_value.max(elems_fp32);
410 + min_value = min_value.min(elems_fp32);
411 + } else {
412 + max_value = max_value.max(elems_fp32.abs());
413 + }
414 + } else {
415 + if constexpr (AZP) {
416 + max_value = max_value.max(elems_fp32, hidden_size - j);
417 + min_value = min_value.min(elems_fp32, hidden_size - j);
418 + } else {
419 + max_value = max_value.max(elems_fp32.abs(), hidden_size - j);
420 + }
421 + }
422 + }
423 +
424 + float scale_val, azp_val;
425 + if constexpr (AZP) {
426 + float max_scalar = max_value.reduce_max();
427 + float min_scalar = min_value.reduce_min();
428 + scale_val = (max_scalar - min_scalar) / 255.0f;
429 + azp_val = std::nearbyint(-128.0f - min_scalar / scale_val);
430 + azp[i] = static_cast<int32_t>(azp_val);
431 + scale[i] = scale_val;
432 + } else {
433 + scale_val = max_value.reduce_max() / 127.0f;
434 + scale[i] = scale_val;
435 + }
436 +
437 + const cvt_vec_t inv_scale(1.0 / scale_val);
438 + const cvt_vec_t azp_vec(azp_val);
439 +
440 + {
441 + int j = 0;
442 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
443 + load_vec_t elems(input + i * hidden_size + j);
444 + cvt_vec_t elems_fp32(elems);
445 + elems_fp32 = (elems_fp32 * inv_scale);
446 +
447 + if constexpr (AZP) {
448 + elems_fp32 = elems_fp32 + azp_vec;
449 + }
450 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
451 + vec_op::INT8Vec16 elems_int8(elems_fp32);
452 + elems_int8.save(output + i * hidden_size + j);
453 + }
454 +
455 + load_vec_t elems(input + i * hidden_size + j);
456 + cvt_vec_t elems_fp32(elems);
457 + elems_fp32 = (elems_fp32 * inv_scale);
458 +
459 + if constexpr (AZP) {
460 + elems_fp32 = elems_fp32 + azp_vec;
461 + }
462 + elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
463 + vec_op::INT8Vec16 elems_int8(elems_fp32);
464 + elems_int8.save(output + i * hidden_size + j, hidden_size - j);
465 + }
466 + }
467 +}
468 +template <bool PerChannel, typename scalar_t>
469 +void static_quant_epilogue(const float* input, scalar_t* output,
470 + const float a_scale, const float* b_scale,
471 + const int32_t* azp_with_adj, const int num_tokens,
472 + const int hidden_size) {
473 + CPU_KERNEL_GUARD_IN(dynamic_output_scale_impl)
474 + using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
475 + using azp_adj_load_vec_t =
476 + typename KernelVecType<scalar_t>::azp_adj_load_vec_type;
477 + using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
478 + constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
479 +
480 + #pragma omp parallel for
481 + for (int i = 0; i < num_tokens; ++i) {
482 + cvt_vec_t a_scale_vec(a_scale);
483 + cvt_vec_t b_scale_vec(*b_scale);
484 + cvt_vec_t scale_vec = a_scale_vec * b_scale_vec;
485 +
486 + int j = 0;
487 + for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
488 + cvt_vec_t elems_fp32(input + i * hidden_size + j);
489 + azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
490 + cvt_vec_t azp_adj_fp32(azp_adj_vec);
491 +
492 + if constexpr (PerChannel) {
493 + b_scale_vec = cvt_vec_t(b_scale + j);
494 + scale_vec = b_scale_vec * a_scale_vec;
495 + }
496 + elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
497 + load_vec_t elems_out(elems_fp32);
498 + elems_out.save(output + i * hidden_size + j);
499 + }
500 +
501 + cvt_vec_t elems_fp32(input + i * hidden_size + j);
502 + azp_adj_load_vec_t azp_adj_vec(azp_with_adj + j);
503 + cvt_vec_t azp_adj_fp32(azp_adj_vec);
504 +
505 + if constexpr (PerChannel) {
506 + b_scale_vec = cvt_vec_t(b_scale + j);
507 + scale_vec = b_scale_vec * a_scale_vec;
508 + }
509 +
510 + elems_fp32 = elems_fp32 - scale_vec * azp_adj_fp32;
511 +
512 + load_vec_t elems_out(elems_fp32);
513 + elems_out.save(output + i * hidden_size + j, hidden_size - j);
514 + }
515 +}
516 template <bool AZP, bool PerChannel, bool Bias, typename scalar_t>
517 void dynamic_quant_epilogue(const float* input, scalar_t* output,
518 const float* a_scale, const float* b_scale,
@@ -324,7 +598,8 @@ void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
598 const float* scale, const int32_t* azp,
599 const int num_tokens,
600 const int hidden_size) {
327 - TORCH_CHECK(false, "static_scaled_int8_quant_impl requires AVX512 support.")
601 + TORCH_CHECK(
602 + false, "static_scaled_int8_quant_impl requires AVX512/powerpc64 support.")
603 }
604
605 template <typename scalar_t>
@@ -332,7 +607,9 @@ void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
607 float* scale, int32_t* azp,
608 const int num_tokens,
609 const int hidden_size) {
335 - TORCH_CHECK(false, "dynamic_scaled_int8_quant_impl requires AVX512 support.")
610 + TORCH_CHECK(
611 + false,
612 + "dynamic_scaled_int8_quant_impl requires AVX512/powerpc64 support.")
613 }
614
615 template <bool PerChannel, typename scalar_t>
@@ -340,7 +617,7 @@ void static_quant_epilogue(const float* input, scalar_t* output,
617 const float a_scale, const float* b_scale,
618 const int32_t* azp_with_adj, const int num_tokens,
619 const int hidden_size) {
343 - TORCH_CHECK(false, "static_quant_epilogue requires AVX512 support.")
620 + TORCH_CHECK(false, "static_quant_epilogue requires AVX512/powerpc64 support.")
621 }
622
623 template <typename scalar_t>
@@ -349,7 +626,8 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
626 const int32_t* azp, const int32_t* azp_with_adj,
627 const scalar_t* bias, const int num_tokens,
628 const int hidden_size) {
352 - TORCH_CHECK(false, "dynamic_quant_epilogue requires AVX512 support.")
629 + TORCH_CHECK(false,
630 + "dynamic_quant_epilogue requires AVX512/powerpc64 support.")
631 }
632 #endif
633 } // namespace
@@ -611,3 +889,58 @@ void dynamic_scaled_int8_quant(
889 }
890 });
891 }
892 +
893 +#if defined(__powerpc64__)
894 +void int8_scaled_mm_ppc64le(torch::Tensor& c, // [M, OC], row-major
895 + const torch::Tensor& a, // [M, IC], row-major
896 + const torch::Tensor& b, // [IC, OC], column-major
897 + const torch::Tensor& a_scales,
898 + const torch::Tensor& b_scales,
899 + const std::optional<torch::Tensor>& bias // [OC]
900 +) {
901 + CPU_KERNEL_GUARD_IN(cutlass_scaled_mm)
902 + // Checks for conformality
903 + TORCH_CHECK(a.dtype() == torch::kInt8 && b.dtype() == torch::kInt8,
904 + "int8_scaled_mm_ppc64le only supports INT8 inputs.");
905 + TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
906 + TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(0) &&
907 + b.size(1) == c.size(1));
908 + // We dont need this
909 + TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
910 + TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == b.size(1));
911 +
912 + // Check for strides and alignment
913 + TORCH_CHECK(a.stride(1) == 1 && c.stride(1) == 1); // Row-major
914 + TORCH_CHECK(b.stride(0) == 1); // Column-major
915 + TORCH_CHECK(c.stride(0) % 16 == 0 &&
916 + b.stride(1) % 16 == 0); // 16 Byte Alignment
917 + TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
918 +
919 + if (bias) {
920 + TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
921 + bias->dim() == 1);
922 + }
923 + VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "int8_scaled_mm_ppc64le", [&] {
924 + torch::Tensor tmp_fp32_out = torch::empty_like(c, ::at::ScalarType::Float);
925 + // Compute C_inter=s_b * (A@B)
926 + DNNLPrimitiveHelper<true>::gemm_s8s8_jit<float, void>(
927 + a.data_ptr<int8_t>(), b.data_ptr<int8_t>(),
928 + tmp_fp32_out.data_ptr<float>(), nullptr, a.size(0), b.size(1),
929 + a.size(1), nullptr, b_scales.data_ptr<float>(), 0, b_scales.numel());
930 + if (bias.has_value()) {
931 + // Compute C=s_a * C_inter + bias
932 + dynamic_quant_epilogue<false, true, true>(
933 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
934 + a_scales.data_ptr<float>(), nullptr, nullptr, nullptr,
935 + bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
936 + } else {
937 + // Compute C=s_a * C_inter
938 + dynamic_quant_epilogue<false, true, false, scalar_t>(
939 + tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
940 + a_scales.data_ptr<float>(), nullptr, nullptr, nullptr, nullptr,
941 + c.size(0), c.size(1));
942 + }
943 + });
944 +}
945 +
946 +#endif
csrc/cpu/torch_bindings.cpp
+35
@@ -18,6 +18,14 @@ void int8_scaled_mm_azp(torch::Tensor& c, const torch::Tensor& a,
18 const std::optional<torch::Tensor>& azp,
19 const std::optional<torch::Tensor>& bias);
20
21 +#if defined(__powerpc64__)
22 +void int8_scaled_mm_ppc64le(torch::Tensor& c, const torch::Tensor& a,
23 + const torch::Tensor& b,
24 + const torch::Tensor& a_scales,
25 + const torch::Tensor& b_scales,
26 + const std::optional<torch::Tensor>& bias);
27 +#endif
28 +
29 void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
30 torch::Tensor& kv_cache, double scale,
31 torch::Tensor& block_tables, torch::Tensor& seq_lens);
@@ -150,6 +158,33 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
158 " Tensor b_scales, Tensor azp_adj,"
159 " Tensor? azp, Tensor? bias) -> ()");
160 ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
161 +#elif defined(__powerpc64__)
162 + // Compute int8 quantized tensor for given scaling factor.
163 + ops.def(
164 + "static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
165 + "Tensor? azp) -> ()");
166 + ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
167 +
168 + // Compute int8 quantized tensor and scaling factor
169 + ops.def(
170 + "dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
171 + "Tensor!? azp) -> ()");
172 + ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
173 + &dynamic_scaled_int8_quant);
174 + // W8A8 GEMM, supporting symmetric quantization.
175 + ops.def(
176 + "cutlass_scaled_mm(Tensor! out, Tensor a,"
177 + " Tensor b, Tensor a_scales,"
178 + " Tensor b_scales, Tensor? bias) -> ()");
179 + ops.impl("cutlass_scaled_mm", torch::kCPU, &int8_scaled_mm_ppc64le);
180 + // w8a8 GEMM, supporting asymmetric per-tensor or per-row/column
181 + // quantization.
182 + ops.def(
183 + "cutlass_scaled_mm_azp(Tensor! out, Tensor a,"
184 + " Tensor b, Tensor a_scales,"
185 + " Tensor b_scales, Tensor azp_adj,"
186 + " Tensor? azp, Tensor? bias) -> ()");
187 + ops.impl("cutlass_scaled_mm_azp", torch::kCPU, &int8_scaled_mm_azp);
188 #endif
189
190 // SHM CCL