[CPU] Enable GPTQ and AWQ quantization for s390x (#51148)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rehan Khan committed
Aug 10, 2026 at 09:11 UTC
b22afe45ac797ae58e67a7a3ad79ee5714024420
5 files changed
+88
-24
cmake/cpu_extension.cmake
+6
@@ -472,6 +472,12 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64" AND VLLM_RVV_VLEN AND
472
${VLLM_EXT_SRC})
473
endif()
474
475
+if (S390_FOUND)
476
+ set(VLLM_EXT_SRC
477
+ "csrc/cpu/cpu_wna16.cpp"
478
+ ${VLLM_EXT_SRC})
479
+endif()
480
+
481
if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
482
set(VLLM_EXT_SRC
483
"csrc/cpu/shm.cpp"
csrc/cpu/cpu_types_vxe.hpp
+78
-20
@@ -704,6 +704,58 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
704
705
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
706
707
+ // De-pack 16 x 4-bit nibbles from a 64-bit value and look each up in a
708
+ // 16-element float LUT. Used by WNA16 (AWQ/GPTQ) dequantization.
709
+ explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
710
+ uint64_t uval = static_cast<uint64_t>(value);
711
+ uval = (uval >> 32) | (uval << 32);
712
+
713
+ // Process 4 floats per output vector register
714
+ for (int v = 0; v < 4; ++v) {
715
+ // Extract 4 nibble indices for this output vector
716
+ uint8_t n0 = (uval >> ((v * 4 + 0) * 4)) & 0xF;
717
+ uint8_t n1 = (uval >> ((v * 4 + 1) * 4)) & 0xF;
718
+ uint8_t n2 = (uval >> ((v * 4 + 2) * 4)) & 0xF;
719
+ uint8_t n3 = (uval >> ((v * 4 + 3) * 4)) & 0xF;
720
+
721
+ // Build permute control: index % 8 * 4 gives byte offset within a
722
+ // 32-byte window (two 16-byte LUT vectors concatenated).
723
+ alignas(16) uint8_t ctrl[16] = {
724
+ (uint8_t)((n0 % 8) * 4 + 0), (uint8_t)((n0 % 8) * 4 + 1),
725
+ (uint8_t)((n0 % 8) * 4 + 2), (uint8_t)((n0 % 8) * 4 + 3),
726
+ (uint8_t)((n1 % 8) * 4 + 0), (uint8_t)((n1 % 8) * 4 + 1),
727
+ (uint8_t)((n1 % 8) * 4 + 2), (uint8_t)((n1 % 8) * 4 + 3),
728
+ (uint8_t)((n2 % 8) * 4 + 0), (uint8_t)((n2 % 8) * 4 + 1),
729
+ (uint8_t)((n2 % 8) * 4 + 2), (uint8_t)((n2 % 8) * 4 + 3),
730
+ (uint8_t)((n3 % 8) * 4 + 0), (uint8_t)((n3 % 8) * 4 + 1),
731
+ (uint8_t)((n3 % 8) * 4 + 2), (uint8_t)((n3 % 8) * 4 + 3),
732
+ };
733
+ __vector unsigned char perm =
734
+ (__vector unsigned char)vec_xl(0, (const signed char*)ctrl);
735
+
736
+ // Gather from both LUT halves via vec_perm (VPERM)
737
+ __vector unsigned char from_lo =
738
+ vec_perm((__vector unsigned char)lut.reg.val[0],
739
+ (__vector unsigned char)lut.reg.val[1], perm);
740
+ __vector unsigned char from_hi =
741
+ vec_perm((__vector unsigned char)lut.reg.val[2],
742
+ (__vector unsigned char)lut.reg.val[3], perm);
743
+
744
+ // Build selection mask: 0xFF bytes for indices >= 8, 0x00 otherwise
745
+ uint8_t m0 = (n0 >= 8) ? 0xFF : 0x00;
746
+ uint8_t m1 = (n1 >= 8) ? 0xFF : 0x00;
747
+ uint8_t m2 = (n2 >= 8) ? 0xFF : 0x00;
748
+ uint8_t m3 = (n3 >= 8) ? 0xFF : 0x00;
749
+ alignas(16) uint8_t sel[16] = {
750
+ m0, m0, m0, m0, m1, m1, m1, m1, m2, m2, m2, m2, m3, m3, m3, m3,
751
+ };
752
+ __vector __bool char mask =
753
+ (__vector __bool char)vec_xl(0, (const signed char*)sel);
754
+
755
+ reg.val[v] = (__vector float)vec_sel(from_lo, from_hi, mask);
756
+ }
757
+ }
758
+
759
// FP8 stub: dead code on s390x (fp8 KV cache is x86-only), needed for
760
// load_b_pair_vec template to compile on all platforms.
761
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
@@ -820,31 +872,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
872
return vec_extract(m, 0);
873
}
874
823
- FP32Vec16 exp() const {
824
- FP32Vec8 lo(f32x4x2_t{reg.val[0], reg.val[1]});
825
- FP32Vec8 hi(f32x4x2_t{reg.val[2], reg.val[3]});
826
- auto lo_e = lo.exp();
827
- auto hi_e = hi.exp();
828
- return FP32Vec16(f32x4x4_t{lo_e.reg.val[0], lo_e.reg.val[1],
829
- hi_e.reg.val[0], hi_e.reg.val[1]});
830
- }
831
-
875
FP32Vec16 abs() const {
876
return FP32Vec16(f32x4x4_t({vec_abs(reg.val[0]), vec_abs(reg.val[1]),
877
vec_abs(reg.val[2]), vec_abs(reg.val[3])}));
878
}
879
837
- FP32Vec16 min(const FP32Vec16& b) const {
838
- return FP32Vec16(f32x4x4_t({vec_min(reg.val[0], b.reg.val[0]),
839
- vec_min(reg.val[1], b.reg.val[1]),
840
- vec_min(reg.val[2], b.reg.val[2]),
841
- vec_min(reg.val[3], b.reg.val[3])}));
842
- }
843
-
844
- FP32Vec16 clamp(const FP32Vec16& lo, const FP32Vec16& hi) const {
845
- return this->max(lo).min(hi);
846
- }
847
-
880
float reduce_min() const {
881
__vector float m = vec_min(vec_min(reg.val[0], reg.val[1]),
882
vec_min(reg.val[2], reg.val[3]));
@@ -1387,6 +1419,32 @@ FORCE_INLINE void prefetch(const void* addr) {
1419
__builtin_prefetch(addr, 0, 3); // 0=read, 3=high temporal locality
1420
}
1421
1422
+static void interleave_save(const BF16Vec16& vec0, const BF16Vec16& vec1,
1423
+ void* ptr) {
1424
+ alignas(16) uint16_t v0[BF16Vec16::VEC_ELEM_NUM];
1425
+ alignas(16) uint16_t v1[BF16Vec16::VEC_ELEM_NUM];
1426
+ vec0.save(v0);
1427
+ vec1.save(v1);
1428
+ auto* packed = reinterpret_cast<uint32_t*>(ptr);
1429
+ for (int i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) {
1430
+ packed[i] =
1431
+ static_cast<uint32_t>(v0[i]) | (static_cast<uint32_t>(v1[i]) << 16);
1432
+ }
1433
+}
1434
+
1435
+static void interleave_save(const FP16Vec16& vec0, const FP16Vec16& vec1,
1436
+ void* ptr) {
1437
+ alignas(16) uint16_t v0[FP16Vec16::VEC_ELEM_NUM];
1438
+ alignas(16) uint16_t v1[FP16Vec16::VEC_ELEM_NUM];
1439
+ vec0.save(v0);
1440
+ vec1.save(v1);
1441
+ auto* packed = reinterpret_cast<uint32_t*>(ptr);
1442
+ for (int i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) {
1443
+ packed[i] =
1444
+ static_cast<uint32_t>(v0[i]) | (static_cast<uint32_t>(v1[i]) << 16);
1445
+ }
1446
+}
1447
+
1448
}; // namespace vec_op
1449
1450
#endif
csrc/cpu/torch_bindings.cpp
+1
-1
@@ -605,7 +605,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
605
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
606
607
// WNA16
608
-#if defined(__AVX512F__) || defined(__riscv_v)
608
+#if defined(__AVX512F__) || defined(__riscv_v) || defined(__s390x__)
609
ops.def(
610
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
611
"Tensor scales, Tensor? zeros, Tensor? g_idx, 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 only)
316
- - GPTQ (x86 only)
315
+ - AWQ (x86, s390x)
316
+ - GPTQ (x86, s390x)
317
- compressed-tensor INT8 W8A8 (x86, s390x only)
318
319
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
docs/getting_started/installation/cpu.s390x.inc.md
+1
-1
@@ -3,7 +3,7 @@
3
4
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.
5
6
-Currently, the CPU implementation for s390x architecture supports FP32, BF16 and FP16.
6
+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.
7
8
--8<-- [end:installation]
9
--8<-- [start:requirements]