iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Hardware][Power] Enable W4A16 (AWQ & GPTQ) quantization on POWER10 using VSX (#59149)

Browse Source 
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Rukhaiya2004
2026-09-30 19:44:48 +08:00
committed by GitHub
co-authored by Li, Jiang <jiang1.li@intel.com>
parent 863475eb9c
commit 72e7874fa6
7 changed files with 99 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/cpu_extension.cmake
+2 -1
View File
@@ -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()
csrc/cpu/cpu_types_vsx.hpp
+45 -1
View File
@@ -417,7 +417,7 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
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<FP32Vec16> {
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<uint64_t>(value);
float f[16];
#pragma GCC unroll 16
for (int i = 0; i < 16; i++) {
f[i] = static_cast<float>((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<signed short*>(ptr));
vec_xst(l0, 16, reinterpret_cast<signed short*>(ptr));
vec_xst(h1, 32, reinterpret_cast<signed short*>(ptr));
vec_xst(l1, 48, reinterpret_cast<signed short*>(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<signed short*>(ptr));
vec_xst(l0, 16, reinterpret_cast<signed short*>(ptr));
vec_xst(h1, 32, reinterpret_cast<signed short*>(ptr));
vec_xst(l1, 48, reinterpret_cast<signed short*>(ptr));
}
struct INT8Vec64 {
__vector signed char data[4];
INT8Vec64() = default;
csrc/cpu/cpu_wna16.cpp
+31 -2
View File
@@ -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<ISA::VSX, scalar_t>;
if (has_zp) {
using dequantizer_t = Dequantizer4b<scalar_t, ISA::VSX, true>;
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
bias.has_value() ? bias->data_ptr<scalar_t>() : 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<scalar_t, ISA::VSX, false>;
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
bias.has_value() ? bias->data_ptr<scalar_t>() : 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;
}
}
});
}
csrc/cpu/torch_bindings.cpp
+2 -1
View File
@@ -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 "
docs/getting_started/installation/cpu.md
+2 -2
View File
@@ -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?
tests/quantization/test_cpu_wna16.py
+12 -2
View File
@@ -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):
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+5
View File
@@ -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"