From 72e7874fa669617fd20c716e08cb486032f542c1 Mon Sep 17 00:00:00 2001 From: Rukhaiya2004 <162583766+Rukhaiya2004@users.noreply.github.com> Date: Wed, 30 Sep 2026 17:14:48 +0530 Subject: [PATCH] [Hardware][Power] Enable W4A16 (AWQ & GPTQ) quantization on POWER10 using VSX (#59149) Signed-off-by: Rukhaiya Co-authored-by: Li, Jiang --- cmake/cpu_extension.cmake | 3 +- csrc/cpu/cpu_types_vsx.hpp | 46 ++++++++++++++++++- csrc/cpu/cpu_wna16.cpp | 33 ++++++++++++- csrc/cpu/torch_bindings.cpp | 3 +- docs/getting_started/installation/cpu.md | 4 +- tests/quantization/test_cpu_wna16.py | 14 +++++- .../kernels/linear/mixed_precision/cpu.py | 5 ++ 7 files changed, 99 insertions(+), 9 deletions(-) diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 94dd188822..664e309e94 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -530,8 +530,9 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) endif() endif() -if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) +if (POWER9_FOUND OR POWER10_FOUND OR POWER11_FOUND) set(VLLM_EXT_SRC + "csrc/cpu/cpu_wna16.cpp" "csrc/cpu/shm.cpp" ${VLLM_EXT_SRC}) endif() diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index eaea921d98..31065508d7 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -417,7 +417,7 @@ struct FP32Vec8 : public Vec { tmp.val[1] = two_x; FP32Vec8 temp_vec(tmp); FP32Vec8 exp_vec = temp_vec.exp(); - vector float e = temp_vec.exp().reg.val[0]; + vector float e = exp_vec.reg.val[0]; vector float num = vec_sub(e, one); vector float den = vec_add(e, one); @@ -624,6 +624,24 @@ struct FP32Vec16 : public Vec { explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {} + // De-pack 16 x 4-bit nibbles from a 64-bit value and convert with offset. + // Used by WNA16 (AWQ/GPTQ) dequantization. + // Fully vectorized using VSX / AltiVec intrinsics with zero stack stores. + FORCE_INLINE explicit FP32Vec16(int64_t value, const FP32Vec16& lut) { + uint64_t u = static_cast(value); + float f[16]; +#pragma GCC unroll 16 + for (int i = 0; i < 16; i++) { + f[i] = static_cast((u >> (i * 4)) & 0xF); + } + // Extract base offset from LUT[0] (0.0f for AWQ, -8.0f for GPTQ) + __vector float offset = vec_splats(vec_extract(lut.reg.val[0], 0)); + reg.val[0] = vec_add(vec_xl(0, f), offset); + reg.val[1] = vec_add(vec_xl(16, f), offset); + reg.val[2] = vec_add(vec_xl(32, f), offset); + reg.val[3] = vec_add(vec_xl(48, f), offset); + } + // FP8 stub: dead code on PowerPC (fp8 KV cache is x86-only), needed for // load_b_pair_vec template to compile on all platforms. explicit FP32Vec16(const BF16Vec32&, int) : reg{} {} @@ -1102,6 +1120,32 @@ inline void prefetch(const void* addr) { __asm__ __volatile__("dcbt 0, %0" : : "r"(addr) : "memory"); } +static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1, + void* ptr) { + __vector signed short h0 = vec_mergeh(vec0.reg.val[0], vec1.reg.val[0]); + __vector signed short l0 = vec_mergel(vec0.reg.val[0], vec1.reg.val[0]); + __vector signed short h1 = vec_mergeh(vec0.reg.val[1], vec1.reg.val[1]); + __vector signed short l1 = vec_mergel(vec0.reg.val[1], vec1.reg.val[1]); + + vec_xst(h0, 0, reinterpret_cast(ptr)); + vec_xst(l0, 16, reinterpret_cast(ptr)); + vec_xst(h1, 32, reinterpret_cast(ptr)); + vec_xst(l1, 48, reinterpret_cast(ptr)); +} + +static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1, + void* ptr) { + __vector signed short h0 = vec_mergeh(vec0.reg.val[0], vec1.reg.val[0]); + __vector signed short l0 = vec_mergel(vec0.reg.val[0], vec1.reg.val[0]); + __vector signed short h1 = vec_mergeh(vec0.reg.val[1], vec1.reg.val[1]); + __vector signed short l1 = vec_mergel(vec0.reg.val[1], vec1.reg.val[1]); + + vec_xst(h0, 0, reinterpret_cast(ptr)); + vec_xst(l0, 16, reinterpret_cast(ptr)); + vec_xst(h1, 32, reinterpret_cast(ptr)); + vec_xst(l1, 48, reinterpret_cast(ptr)); +} + struct INT8Vec64 { __vector signed char data[4]; INT8Vec64() = default; diff --git a/csrc/cpu/cpu_wna16.cpp b/csrc/cpu/cpu_wna16.cpp index d3e42660f2..209c89454e 100644 --- a/csrc/cpu/cpu_wna16.cpp +++ b/csrc/cpu/cpu_wna16.cpp @@ -7,6 +7,9 @@ #if defined(__riscv_v) #include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp" #endif +#if defined(__powerpc__) + #include "cpu/micro_gemm/cpu_micro_gemm_vsx.hpp" +#endif #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp" #define VLLM_DISPATCH_CASE_16B_TYPES(...) \ @@ -104,8 +107,8 @@ class Dequantizer4b { scalar_vec_t output_vec_0(wb_0); scalar_vec_t output_vec_1(wb_1); - // AMX needs to interleave K elements to pack as 32 bits - if constexpr (isa == ISA::AMX) { + // AMX and VSX need to interleave K elements to pack as 32 bits + if constexpr (isa == ISA::AMX || isa == ISA::VSX) { vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight); } else { output_vec_0.save(curr_weight); @@ -303,6 +306,8 @@ void cpu_gemm_wna16( return ISA::VEC; } else if (isa_hint == "rvv") { return ISA::RVV; + } else if (isa_hint == "vsx") { + return ISA::VSX; } else { TORCH_CHECK(false, "unsupported isa hint: " + isa_hint); } @@ -378,6 +383,30 @@ void cpu_gemm_wna16( scales_group_stride, zeros_group_stride, group_size, pack_factor); return; } + } else if (isa == ISA::VSX) { + TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, + "VSX WNA16 GEMM only supports BFloat16"); + using gemm_t = cpu_micro_gemm::MicroGemm; + if (has_zp) { + using dequantizer_t = Dequantizer4b; + cpu_gemm_wna16_impl( + input.data_ptr(), q_weight.data_ptr(), + output.data_ptr(), scales.data_ptr(), zeros_ptr, + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); + return; + } + { + using dequantizer_t = Dequantizer4b; + cpu_gemm_wna16_impl( + input.data_ptr(), q_weight.data_ptr(), + output.data_ptr(), scales.data_ptr(), zeros_ptr, + bias.has_value() ? bias->data_ptr() : nullptr, a_m_size, + b_n_size, a_k_size, a_m_stride, output_m_stride, + scales_group_stride, zeros_group_stride, group_size, pack_factor); + return; + } } }); } diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 918daf5a8b..6bbb06efa6 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -902,7 +902,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op); // WNA16 -#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) +#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__) || \ + defined(__powerpc__) ops.def( "cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, " "Tensor scales, Tensor? zeros, Tensor? bias, SymInt " diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index 53b8eb596d..20834529b1 100644 --- a/docs/getting_started/installation/cpu.md +++ b/docs/getting_started/installation/cpu.md @@ -312,8 +312,8 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel ### Which quantization configs does vLLM CPU support? - vLLM CPU supports quantizations: - - AWQ (x86, s390x) - - GPTQ (x86, s390x) + - AWQ (x86, s390x, Power) + - GPTQ (x86, s390x, Power) - compressed-tensor INT8 W8A8 (x86, s390x only) ### Why do I see `get_mempolicy: Operation not permitted` when running in Docker? diff --git a/tests/quantization/test_cpu_wna16.py b/tests/quantization/test_cpu_wna16.py index 2af633be6b..d15d4b8147 100644 --- a/tests/quantization/test_cpu_wna16.py +++ b/tests/quantization/test_cpu_wna16.py @@ -3,7 +3,7 @@ import pytest import torch -from vllm.platforms import current_platform +from vllm.platforms import CpuArchEnum, current_platform if not current_platform.is_cpu(): pytest.skip("skipping CPU-only tests", allow_module_level=True) @@ -54,7 +54,7 @@ def test_cpu_gemm_wna16_3d_input(): MODELS = [ "TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ", - "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", + "Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx "RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp "OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", "Qwen/Qwen3-0.6B-FP8", # FP8 W8A16 block-quantized linear @@ -67,6 +67,16 @@ MODELS = [ DTYPE = ["bfloat16"] +@pytest.mark.skipif( + current_platform.get_cpu_architecture() != CpuArchEnum.POWERPC, + reason="PowerPC-only test", +) +def test_cpu_wna16_power_bf16_uses_vsx(): + """Use the POWER10 BF16 MMA microkernel instead of generic vectors.""" + assert _get_isa_hint(torch.bfloat16) == "vsx" + assert _get_isa_hint(torch.float16) == "vec" + + @pytest.mark.parametrize("model", MODELS) @pytest.mark.parametrize("dtype", DTYPE) def test_cpu_quant(vllm_runner, model, dtype): diff --git a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py index 19f5283acb..a9c13051e2 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/cpu.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/cpu.py @@ -213,5 +213,10 @@ def _get_isa_hint(dtype: torch.dtype) -> str: return "amx" elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV: return "rvv" + elif ( + dtype == torch.bfloat16 + and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC + ): + return "vsx" else: return "vec"