Skip to content

[CPU][PowerPC] Vectorize top-k and top-p sampling for PowerPC (and s390x) #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Akashcodes732 wants to merge 1 commit into
base: arm-neon-topk-topp
Choose a base branch
from
Open

[CPU][PowerPC] Vectorize top-k and top-p sampling for PowerPC (and s390x) #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 123 additions & 3 deletions csrc/cpu/topk_topp_sampler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#include "cpu_arch_macros.h"
#include "cpu_types.hpp"
#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>
Expand Down Expand Up @@ -94,6 +95,44 @@ struct Fp64Acc16 {
return vaddvq_f64(sum);
}
};
#elif defined(__powerpc__) || defined(__s390x__)
using vf32 = __vector float;
using vi32 = __vector signed int;

static inline int hsum_i32(vi32 v) { return v[0] + v[1] + v[2] + v[3]; }

static inline float hmin_f32(vf32 v) {
return std::min(std::min(v[0], v[1]), std::min(v[2], v[3]));
}

struct Fp64Acc16 {
__vector double acc[8] = {};
inline void add(const vec_op::FP32Vec16& v) {
for (int i = 0; i < 4; ++i) {
const vf32 x = v.reg.val[i];
// vec_doublee widens lanes 0 and 2; rotate by one lane for 1 and 3.
acc[2 * i] += vec_doublee(x);
acc[2 * i + 1] += vec_doublee(vec_sld(x, x, 4));
}
}
inline double reduce() const {
__vector double sum = acc[0];
for (int i = 1; i < 8; ++i) sum += acc[i];
return sum[0] + sum[1];
}
};

// VSX / VXE exp() clamps its input, so exp(-inf) yields ~1e-38 instead of 0.
// Zero PAD lanes to match the scalar and NEON paths.
static inline vec_op::FP32Vec16 zero_pad_lanes(const vec_op::FP32Vec16& v,
const vec_op::FP32Vec16& e) {
const vf32 pad = vec_splats(kPadSentinel);
const vf32 zero = vec_splats(0.0f);
vec_op::f32x4x4_t out;
for (int r = 0; r < 4; ++r)
out.val[r] = vec_sel(zero, e.reg.val[r], vec_cmpgt(v.reg.val[r], pad));
return vec_op::FP32Vec16(out);
}
#endif

FORCE_INLINE static double sum_gt_to_double(const float* buf, int n,
Expand All @@ -113,6 +152,19 @@ FORCE_INLINE static double sum_gt_to_double(const float* buf, int n,
acc.add(vec_op::FP32Vec16(masked));
}
s = acc.reduce();
#elif defined(__powerpc__) || defined(__s390x__)
Fp64Acc16 acc;
const vf32 thr = vec_splats(threshold);
const vf32 zero = vec_splats(0.0f);
for (; j + 16 <= n; j += 16) {
vec_op::f32x4x4_t masked;
for (int r = 0; r < 4; ++r) {
const vf32 v = vec_xl(0, buf + j + 4 * r);
masked.val[r] = vec_sel(zero, v, vec_cmpgt(v, thr));
}
acc.add(vec_op::FP32Vec16(masked));
}
s = acc.reduce();
#endif
for (; j < n; ++j)
if (buf[j] > threshold) s += (double)buf[j];
Expand All @@ -133,6 +185,18 @@ FORCE_INLINE static int count_within_tol(const float* row, int V, float center,
n += (int)vaddvq_u32(vshrq_n_u32(mask, 31));
}
}
#elif defined(__powerpc__) || defined(__s390x__)
const vf32 c = vec_splats(center);
const vf32 t = vec_splats(tol);
// Mask lanes are all-ones (-1), so subtracting counts matches.
vi32 cnt = vec_splats(0);
for (; i + 16 <= V; i += 16) {
for (int r = 0; r < 4; ++r) {
const vf32 v = vec_xl(0, row + i + 4 * r);
cnt -= (vi32)vec_cmplt(vec_abs(v - c), t);
}
}
n += hsum_i32(cnt);
#endif
for (; i < V; ++i)
if (fabsf(row[i] - center) < tol) ++n;
Expand All @@ -152,6 +216,15 @@ FORCE_INLINE static void mask_write_below(float* row, int V, float threshold,
vst1q_f32(row + i + 4 * r, v);
}
}
#elif defined(__powerpc__) || defined(__s390x__)
const vf32 thr = vec_splats(threshold);
const vf32 fv = vec_splats(fill);
for (; i + 16 <= V; i += 16) {
for (int r = 0; r < 4; ++r) {
const vf32 v = vec_xl(0, row + i + 4 * r);
vec_xst(vec_sel(fv, v, vec_cmpgt(v, thr)), 0, row + i + 4 * r);
}
}
#endif
for (; i < V; ++i)
if (!(row[i] > threshold)) row[i] = fill;
Expand Down Expand Up @@ -187,6 +260,30 @@ FORCE_INLINE static void vec_max_min_with_pad_blend(
max_l = maxv.reduce_max();
min_l = minv.reduce_min();
}
#elif defined(__powerpc__) || defined(__s390x__)
{
const float pos_inf_val = std::numeric_limits<float>::infinity();
const vf32 pad = vec_splats(pad_sentinel);
const vf32 pos_inf = vec_splats(pos_inf_val);
vec_op::FP32Vec16 maxv(row[0]);
vec_op::FP32Vec16 minv(pos_inf_val);
vi32 cnt = vec_splats(0);
for (; i + 16 <= V; i += 16) {
vec_op::FP32Vec16 v16(row + i);
maxv = maxv.max(v16);
vec_op::f32x4x4_t safe_raw;
for (int r = 0; r < 4; ++r) {
const vf32 x = v16.reg.val[r];
const auto fin = vec_cmpgt(x, pad);
cnt -= (vi32)fin;
safe_raw.val[r] = vec_sel(pos_inf, x, fin);
}
minv = minv.min(vec_op::FP32Vec16(safe_raw));
}
n_finite = hsum_i32(cnt);
max_l = maxv.reduce_max();
min_l = minv.reduce_min();
}
#endif
*max_out = max_l;
*min_out = min_l;
Expand Down Expand Up @@ -341,6 +438,23 @@ static void scan_pivot_stats(const float* buf, int n, float pivot, float tol,
}
ma = vminvq_f32(vmin);
}
#elif defined(__powerpc__) || defined(__s390x__)
{
const vf32 vpivot = vec_splats(pivot);
const vf32 pos_inf = vec_splats(std::numeric_limits<float>::infinity());
vf32 vmin = pos_inf;
vi32 cnt = vec_splats(0);
for (; i + 16 <= n; i += 16) {
for (int r = 0; r < 4; ++r) {
const vf32 values = vec_xl(0, buf + i + 4 * r);
const auto above = vec_cmpgt(values, vpivot);
cnt -= (vi32)above;
vmin = vec_min(vmin, vec_sel(pos_inf, values, above));
}
}
na = hsum_i32(cnt);
ma = hmin_f32(vmin);
}
#endif
for (; i < n; ++i) {
if (buf[i] > pivot) {
Expand Down Expand Up @@ -482,18 +596,21 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val,
sigma = sigma + fabsf(sigma) * -0.25f;
float outlier_logit = st.avg + st.std_v * sigma;

#if defined(__aarch64__)
#if defined(__aarch64__) || defined(__powerpc__) || defined(__s390x__)
DEFINE_FAST_EXP
vec_op::FP32Vec16 base16(st.max_l);
#endif
double sum_exp_d = 0.0;
{
int i = 0;
#if defined(__aarch64__)
#if defined(__aarch64__) || defined(__powerpc__) || defined(__s390x__)
Fp64Acc16 acc;
for (; i + 16 <= V; i += 16) {
vec_op::FP32Vec16 v16(row + i);
vec_op::FP32Vec16 e16 = fast_exp(v16 - base16);
#if defined(__powerpc__) || defined(__s390x__)
e16 = zero_pad_lanes(v16, e16);
#endif
acc.add(e16);
}
sum_exp_d += acc.reduce();
Expand All @@ -510,11 +627,14 @@ static void _top_p_row_core(float* __restrict__ row, int V, float p_val,
{
const float inv_sum = 1.0f / sum_exp;
int i = 0;
#if defined(__aarch64__)
#if defined(__aarch64__) || defined(__powerpc__) || defined(__s390x__)
vec_op::FP32Vec16 isum16(inv_sum);
for (; i + 16 <= V; i += 16) {
vec_op::FP32Vec16 v16(row + i);
vec_op::FP32Vec16 p16 = fast_exp(v16 - base16) * isum16;
#if defined(__powerpc__) || defined(__s390x__)
p16 = zero_pad_lanes(v16, p16);
#endif
p16.save(scratch + i);
}
#endif
Expand Down
19 changes: 14 additions & 5 deletions vllm/v1/sample/ops/topk_topp_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,10 +260,8 @@ def __init__(
)
elif current_platform.is_cpu():
arch = current_platform.get_cpu_architecture()
# Fall back to native implementation for POWERPC and RISCV.
# On PowerPC argmax produces incorrect output with torch.compile.
# PR: https://github.com/vllm-project/vllm/pull/26987
if arch in (CpuArchEnum.RISCV, CpuArchEnum.POWERPC):
# Fall back to native implementation for RISCV.
if arch == CpuArchEnum.RISCV:
self.forward = self.forward_native
else:
self.forward = self.forward_cpu
Expand Down Expand Up @@ -366,7 +364,11 @@ def forward_cpu(
if k is not None or p is not None:
if HAS_TRITON:
logits = apply_top_k_top_p_triton(logits, k, p)
elif current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
elif current_platform.get_cpu_architecture() in (
CpuArchEnum.ARM,
CpuArchEnum.POWERPC,
CpuArchEnum.S390X,
):
logits = apply_top_k_top_p_cpu(logits, k, p)
else:
logits = apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=True)
Expand Down Expand Up @@ -532,6 +534,13 @@ def apply_top_k_top_p(
logits, k, p, has_topp_only_rows=has_topp_only_rows
)

if current_platform.is_cpu() and current_platform.get_cpu_architecture() in (
CpuArchEnum.ARM,
CpuArchEnum.POWERPC,
CpuArchEnum.S390X,
):
return apply_top_k_top_p_cpu(logits, k, p)

is_cpu = current_platform.is_cpu()
return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu)

Expand Down