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

[CPU] Enable GPTQ and AWQ quantization for s390x (#51148)

Browse Source 
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
This commit is contained in:
Rehan Khan
2026-08-10 11:41:43 +08:00
committed by GitHub
co-authored by Li, Jiang <jiang1.li@intel.com>
parent 7ce84b99ce
commit b22afe45ac
5 changed files with 88 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmake/cpu_extension.cmake
+6
View File
@@ -472,6 +472,12 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
${VLLM_EXT_SRC})
endif()
if (S390_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/cpu_wna16.cpp"
${VLLM_EXT_SRC})
endif()
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
set(VLLM_EXT_SRC
"csrc/cpu/shm.cpp"
csrc/cpu/cpu_types_vxe.hpp
+78 -20
View File
@@ -704,6 +704,58 @@ 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 look each up in a
// 16-element float LUT. Used by WNA16 (AWQ/GPTQ) dequantization.
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
uint64_t uval = static_cast<uint64_t>(value);
uval = (uval >> 32) | (uval << 32);
// Process 4 floats per output vector register
for (int v = 0; v < 4; ++v) {
// Extract 4 nibble indices for this output vector
uint8_t n0 = (uval >> ((v * 4 + 0) * 4)) & 0xF;
uint8_t n1 = (uval >> ((v * 4 + 1) * 4)) & 0xF;
uint8_t n2 = (uval >> ((v * 4 + 2) * 4)) & 0xF;
uint8_t n3 = (uval >> ((v * 4 + 3) * 4)) & 0xF;
// Build permute control: index % 8 * 4 gives byte offset within a
// 32-byte window (two 16-byte LUT vectors concatenated).
alignas(16) uint8_t ctrl[16] = {
(uint8_t)((n0 % 8) * 4 + 0), (uint8_t)((n0 % 8) * 4 + 1),
(uint8_t)((n0 % 8) * 4 + 2), (uint8_t)((n0 % 8) * 4 + 3),
(uint8_t)((n1 % 8) * 4 + 0), (uint8_t)((n1 % 8) * 4 + 1),
(uint8_t)((n1 % 8) * 4 + 2), (uint8_t)((n1 % 8) * 4 + 3),
(uint8_t)((n2 % 8) * 4 + 0), (uint8_t)((n2 % 8) * 4 + 1),
(uint8_t)((n2 % 8) * 4 + 2), (uint8_t)((n2 % 8) * 4 + 3),
(uint8_t)((n3 % 8) * 4 + 0), (uint8_t)((n3 % 8) * 4 + 1),
(uint8_t)((n3 % 8) * 4 + 2), (uint8_t)((n3 % 8) * 4 + 3),
};
__vector unsigned char perm =
(__vector unsigned char)vec_xl(0, (const signed char*)ctrl);
// Gather from both LUT halves via vec_perm (VPERM)
__vector unsigned char from_lo =
vec_perm((__vector unsigned char)lut.reg.val[0],
(__vector unsigned char)lut.reg.val[1], perm);
__vector unsigned char from_hi =
vec_perm((__vector unsigned char)lut.reg.val[2],
(__vector unsigned char)lut.reg.val[3], perm);
// Build selection mask: 0xFF bytes for indices >= 8, 0x00 otherwise
uint8_t m0 = (n0 >= 8) ? 0xFF : 0x00;
uint8_t m1 = (n1 >= 8) ? 0xFF : 0x00;
uint8_t m2 = (n2 >= 8) ? 0xFF : 0x00;
uint8_t m3 = (n3 >= 8) ? 0xFF : 0x00;
alignas(16) uint8_t sel[16] = {
m0, m0, m0, m0, m1, m1, m1, m1, m2, m2, m2, m2, m3, m3, m3, m3,
};
__vector __bool char mask =
(__vector __bool char)vec_xl(0, (const signed char*)sel);
reg.val[v] = (__vector float)vec_sel(from_lo, from_hi, mask);
}
}
// FP8 stub: dead code on s390x (fp8 KV cache is x86-only), needed for
// load_b_pair_vec template to compile on all platforms.
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
@@ -820,31 +872,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
return vec_extract(m, 0);
}
FP32Vec16 exp() const {
FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
auto lo_e = lo.exp();
auto hi_e = hi.exp();
return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
hi_e.reg.val[0], hi_e.reg.val[1]});
}
FP32Vec16 abs() const {
return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]),
vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
}
FP32Vec16 min(const FP32Vec16& b) const {
return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
vec_min(reg.val[1], b.reg.val[1]),
vec_min(reg.val[2], b.reg.val[2]),
vec_min(reg.val[3], b.reg.val[3])}));
}
FP32Vec16 clamp(const FP32Vec16& lo, const FP32Vec16& hi) const {
return this->max(lo).min(hi);
}
float reduce_min() const {
__vector float m = vec_min(vec_min(reg.val[0], reg.val[1]),
vec_min(reg.val[2], reg.val[3]));
@@ -1387,6 +1419,32 @@ FORCE_INLINE void prefetch(const void* addr) {
__builtin_prefetch(addr, 0, 3); // 0=read, 3=high temporal locality
}
static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
void* ptr) {
alignas(16) uint16_t v0[BF16Vec16::VEC_ELEM_NUM];
alignas(16) uint16_t v1[BF16Vec16::VEC_ELEM_NUM];
vec0.save(v0);
vec1.save(v1);
auto* packed = reinterpret_cast<uint32_t*>(ptr);
for (int i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
packed[i] =
static_cast<uint32_t>(v0[i]) | (static_cast<uint32_t>(v1[i]) << 16);
}
}
static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
void* ptr) {
alignas(16) uint16_t v0[FP16Vec16::VEC_ELEM_NUM];
alignas(16) uint16_t v1[FP16Vec16::VEC_ELEM_NUM];
vec0.save(v0);
vec1.save(v1);
auto* packed = reinterpret_cast<uint32_t*>(ptr);
for (int i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
packed[i] =
static_cast<uint32_t>(v0[i]) | (static_cast<uint32_t>(v1[i]) << 16);
}
}
}; // namespace vec_op
#endif
csrc/cpu/torch_bindings.cpp
+1 -1
View File
@@ -605,7 +605,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
// WNA16
#if defined(__AVX512F__) || defined(__riscv_v)
#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__)
ops.def(
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
"Tensor scales, Tensor? zeros, Tensor? g_idx, 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 only)
- GPTQ (x86 only)
- AWQ (x86, s390x)
- GPTQ (x86, s390x)
- compressed-tensor INT8 W8A8 (x86, s390x only)
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
docs/getting_started/installation/cpu.s390x.inc.md
+1 -1
View File
@@ -3,7 +3,7 @@
vLLM has experimental support for s390x architecture on IBM Z platform. For now, users must build from source to natively run on IBM Z platform.
Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16, as well as AWQ and GPTQ 4-bit quantization and compressed-tensors INT8 W8A8.
--8<-- [end:installation]
--8<-- [start:requirements]