diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 581d6003f5..46ed96c600 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -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" diff --git a/csrc/cpu/cpu_types_vxe.hpp b/csrc/cpu/cpu_types_vxe.hpp index f705213a60..3e3bf9587b 100644 --- a/csrc/cpu/cpu_types_vxe.hpp +++ b/csrc/cpu/cpu_types_vxe.hpp @@ -704,6 +704,58 @@ 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 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(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 { 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(ptr); + for (int i = 0; i < BF16Vec16::VEC_ELEM_NUM; ++i) { + packed[i] = + static_cast(v0[i]) | (static_cast(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(ptr); + for (int i = 0; i < FP16Vec16::VEC_ELEM_NUM; ++i) { + packed[i] = + static_cast(v0[i]) | (static_cast(v1[i]) << 16); + } +} + }; // namespace vec_op #endif diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 7527ac4846..4bb0edf2b2 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -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 " diff --git a/docs/getting_started/installation/cpu.md b/docs/getting_started/installation/cpu.md index ce7562478f..539c28d709 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 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? diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md index 16fa13525a..88767e4c6e 100644 --- a/docs/getting_started/installation/cpu.s390x.inc.md +++ b/docs/getting_started/installation/cpu.s390x.inc.md @@ -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]