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
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 "