[Hardware][Power] Enable W4A16 (AWQ & GPTQ) quantization on POWER10 using VSX (#59149)
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rukhaiya2004 committed
Sep 30, 2026 at 17:14 UTC
72e7874fa669617fd20c716e08cb486032f542c1
7 files changed
+99
-9
cmake/cpu_extension.cmake
+2
-1
@@ -530,8 +530,9 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
530
endif()
531
endif()
532
533
-if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
533
+if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND)
534
set(VLLM_EXT_SRC
535
+ "csrc/cpu/cpu_wna16.cpp"
536
"csrc/cpu/shm.cpp"
537
${VLLM_EXT_SRC})
538
endif()
csrc/cpu/cpu_types_vsx.hpp
+45
-1
@@ -417,7 +417,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
417
tmp.val[1] = two_x;
418
FP32Vec8 temp_vec(tmp);
419
FP32Vec8 exp_vec = temp_vec.exp();
420
- vector float e = temp_vec.exp().reg.val[0];
420
+ vector float e = exp_vec.reg.val[0];
421
422
vector float num = vec_sub(e, one);
423
vector float den = vec_add(e, one);
@@ -624,6 +624,24 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
624
625
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
626
627
+ // De-pack 16 x 4-bit nibbles from a 64-bit value and convert with offset.
628
+ // Used by WNA16 (AWQ/GPTQ) dequantization.
629
+ // Fully vectorized using VSX / AltiVec intrinsics with zero stack stores.
630
+ FORCE_INLINE explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
631
+ uint64_t u = static_cast<uint64_t>(value);
632
+ float f[16];
633
+#pragma GCC unroll 16
634
+ for (int i = 0; i < 16; i++) {
635
+ f[i] = static_cast<float>((u >> (i * 4)) & 0xF);
636
+ }
637
+ // Extract base offset from LUT[0] (0.0f for AWQ, -8.0f for GPTQ)
638
+ __vector float offset = vec_splats(vec_extract(lut.reg.val[0], 0));
639
+ reg.val[0] = vec_add(vec_xl(0, f), offset);
640
+ reg.val[1] = vec_add(vec_xl(16, f), offset);
641
+ reg.val[2] = vec_add(vec_xl(32, f), offset);
642
+ reg.val[3] = vec_add(vec_xl(48, f), offset);
643
+ }
644
+
645
// FP8 stub: dead code on PowerPC (fp8 KV cache is x86-only), needed for
646
// load_b_pair_vec template to compile on all platforms.
647
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
@@ -1102,6 +1120,32 @@ inline void prefetch(const void* addr) {
1120
__asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory");
1121
}
1122
1123
+static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
1124
+ void* ptr) {
1125
+ __vector signed short h0 = vec_mergeh(vec0.reg.val[0], vec1.reg.val[0]);
1126
+ __vector signed short l0 = vec_mergel(vec0.reg.val[0], vec1.reg.val[0]);
1127
+ __vector signed short h1 = vec_mergeh(vec0.reg.val[1], vec1.reg.val[1]);
1128
+ __vector signed short l1 = vec_mergel(vec0.reg.val[1], vec1.reg.val[1]);
1129
+
1130
+ vec_xst(h0, 0, reinterpret_cast<signed short*>(ptr));
1131
+ vec_xst(l0, 16, reinterpret_cast<signed short*>(ptr));
1132
+ vec_xst(h1, 32, reinterpret_cast<signed short*>(ptr));
1133
+ vec_xst(l1, 48, reinterpret_cast<signed short*>(ptr));
1134
+}
1135
+
1136
+static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
1137
+ void* ptr) {
1138
+ __vector signed short h0 = vec_mergeh(vec0.reg.val[0], vec1.reg.val[0]);
1139
+ __vector signed short l0 = vec_mergel(vec0.reg.val[0], vec1.reg.val[0]);
1140
+ __vector signed short h1 = vec_mergeh(vec0.reg.val[1], vec1.reg.val[1]);
1141
+ __vector signed short l1 = vec_mergel(vec0.reg.val[1], vec1.reg.val[1]);
1142
+
1143
+ vec_xst(h0, 0, reinterpret_cast<signed short*>(ptr));
1144
+ vec_xst(l0, 16, reinterpret_cast<signed short*>(ptr));
1145
+ vec_xst(h1, 32, reinterpret_cast<signed short*>(ptr));
1146
+ vec_xst(l1, 48, reinterpret_cast<signed short*>(ptr));
1147
+}
1148
+
1149
struct INT8Vec64 {
1150
__vector signed char data[4];
1151
INT8Vec64() = default;
csrc/cpu/cpu_wna16.cpp
+31
-2
@@ -7,6 +7,9 @@
7
#if defined(__riscv_v)
8
#include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
9
#endif
10
+#if defined(__powerpc__)
11
+ #include "cpu/micro_gemm/cpu_micro_gemm_vsx.hpp"
12
+#endif
13
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
14
15
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
@@ -104,8 +107,8 @@ class Dequantizer4b {
107
scalar_vec_t output_vec_0(wb_0);
108
scalar_vec_t output_vec_1(wb_1);
109
107
- // AMX needs to interleave K elements to pack as 32 bits
108
- if constexpr (isa == ISA::AMX) {
110
+ // AMX and VSX need to interleave K elements to pack as 32 bits
111
+ if constexpr (isa == ISA::AMX || isa == ISA::VSX) {
112
vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
113
} else {
114
output_vec_0.save(curr_weight);
@@ -303,6 +306,8 @@ void cpu_gemm_wna16(
306
return ISA::VEC;
307
} else if (isa_hint == "rvv") {
308
return ISA::RVV;
309
+ } else if (isa_hint == "vsx") {
310
+ return ISA::VSX;
311
} else {
312
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
313
}
@@ -378,6 +383,30 @@ void cpu_gemm_wna16(
383
scales_group_stride, zeros_group_stride, group_size, pack_factor);
384
return;
385
}
386
+ } else if (isa == ISA::VSX) {
387
+ TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16,
388
+ "VSX WNA16 GEMM only supports BFloat16");
389
+ using gemm_t = cpu_micro_gemm::MicroGemm<ISA::VSX, scalar_t>;
390
+ if (has_zp) {
391
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::VSX, true>;
392
+ cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
393
+ input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
394
+ output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
395
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
396
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
397
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
398
+ return;
399
+ }
400
+ {
401
+ using dequantizer_t = Dequantizer4b<scalar_t, ISA::VSX, false>;
402
+ cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
403
+ input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
404
+ output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
405
+ bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr, a_m_size,
406
+ b_n_size, a_k_size, a_m_stride, output_m_stride,
407
+ scales_group_stride, zeros_group_stride, group_size, pack_factor);
408
+ return;
409
+ }
410
}
411
});
412
}
csrc/cpu/torch_bindings.cpp
+2
-1
@@ -902,7 +902,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
902
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
903
904
// WNA16
905
-#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__)
905
+#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) || \
906
+ defined(__powerpc__)
907
ops.def(
908
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
909
"Tensor scales, Tensor? zeros, Tensor? bias, SymInt "
docs/getting_started/installation/cpu.md
+2
-2
@@ -312,8 +312,8 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
312
### Which quantization configs does vLLM CPU support?
313
314
- vLLM CPU supports quantizations:
315
- - AWQ (x86, s390x)
316
- - GPTQ (x86, s390x)
315
+ - AWQ (x86, s390x, Power)
316
+ - GPTQ (x86, s390x, Power)
317
- compressed-tensor INT8 W8A8 (x86, s390x only)
318
319
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
tests/quantization/test_cpu_wna16.py
+12
-2
@@ -3,7 +3,7 @@
3
import pytest
4
import torch
5
6
-from vllm.platforms import current_platform
6
+from vllm.platforms import CpuArchEnum, current_platform
7
8
if not current_platform.is_cpu():
9
pytest.skip("skipping CPU-only tests", allow_module_level=True)
@@ -54,7 +54,7 @@ def test_cpu_gemm_wna16_3d_input():
54
55
MODELS = [
56
"TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
57
- "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4",
57
+ "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
58
"RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
59
"OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
60
"Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear
@@ -67,6 +67,16 @@ MODELS = [
67
DTYPE = ["bfloat16"]
68
69
70
+@pytest.mark.skipif(
71
+ current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC,
72
+ reason="PowerPC-only test",
73
+)
74
+def test_cpu_wna16_power_bf16_uses_vsx():
75
+ """Use the POWER10 BF16 MMA microkernel instead of generic vectors."""
76
+ assert _get_isa_hint(torch.bfloat16) == "vsx"
77
+ assert _get_isa_hint(torch.float16) == "vec"
78
+
79
+
80
@pytest.mark.parametrize("model", MODELS)
81
@pytest.mark.parametrize("dtype", DTYPE)
82
def test_cpu_quant(vllm_runner, model, dtype):
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+5
@@ -213,5 +213,10 @@ def _get_isa_hint(dtype: torch.dtype) -> str:
213
return "amx"
214
elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
215
return "rvv"
216
+ elif (
217
+ dtype == torch.bfloat16
218
+ and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
219
+ ):
220
+ return "vsx"
221
else:
222
return "vec"