[HARDWARE][POWER] optimize math functions of VSX power (#47321)
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akash.kaothalkar@ibm.com> Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Co-authored-by: Akash Kaothalkar <akash.kaothalkar@ibm.com>
Rukhaiya2004 committed
Jul 7, 2026 at 15:05 UTC
c85d72076ae2263e0634a6bdb7c1b8a4b0cbab11
5 files changed
+154
-54
csrc/cpu/cpu_attn_vsx.hpp
-2
@@ -323,8 +323,6 @@ class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
323
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
324
const int64_t block_size, const int64_t block_size_stride,
325
const float k_inv = 0.0f, const float v_inv = 0.0f) {
326
- // k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
327
- // PowerPC. The parameters are present to match the common interface.
326
#pragma omp parallel for collapse(2)
327
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
328
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
csrc/cpu/cpu_types.hpp
+2
-2
@@ -4,7 +4,7 @@
4
#if defined(__x86_64__)
5
// x86 implementation
6
#include "cpu_types_x86.hpp"
7
-#elif defined(__POWER9_VECTOR__)
7
+#elif defined(__powerpc__)
8
// ppc implementation
9
#include "cpu_types_vsx.hpp"
10
#elif defined(__s390x__)
@@ -41,4 +41,4 @@ inline int get_max_threads() {
41
}
42
} // namespace cpu_utils
43
44
-#endif
\ No newline at end of file
44
+#endif
csrc/cpu/cpu_types_vsx.hpp
+120
-40
@@ -344,53 +344,133 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
344
345
return result;
346
}
347
-
347
FP32Vec8 exp() const {
349
- // TODO: Vectorize this
350
- AliasReg ar;
351
- ar.reg = reg;
352
- f32x4x4_t ret;
353
- ret.val[0][0] = std::exp(ar.values[0]);
354
- ret.val[0][1] = std::exp(ar.values[1]);
355
- ret.val[0][2] = std::exp(ar.values[2]);
356
- ret.val[0][3] = std::exp(ar.values[3]);
357
- ret.val[1][0] = std::exp(ar.values[4]);
358
- ret.val[1][1] = std::exp(ar.values[5]);
359
- ret.val[1][2] = std::exp(ar.values[6]);
360
- ret.val[1][3] = std::exp(ar.values[7]);
361
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
348
+ f32x4x2_t out;
349
+ const __vector float log2e = vec_splats(1.44269504088896341f);
350
+ const __vector float one = vec_splats(1.0f);
351
+ const __vector float min_x = vec_splats(-87.3f);
352
+ const __vector float max_x = vec_splats(88.7f);
353
+
354
+ // 5th-degree minimax polynomial for 2^r (r in [0,1))
355
+ const __vector float c1 = vec_splats(0.6931471805599453f);
356
+ const __vector float c2 = vec_splats(0.240226506959101f);
357
+ const __vector float c3 = vec_splats(0.05550410866482158f);
358
+ const __vector float c4 = vec_splats(0.009618129107628477f);
359
+ const __vector float c5 = vec_splats(0.0013333558146428443f);
360
+
361
+ for (int i = 0; i < 2; i++) {
362
+ __vector float x = reg.val[i];
363
+ x = vec_max(x, min_x);
364
+ x = vec_min(x, max_x);
365
+
366
+ __vector float y = vec_mul(x, log2e);
367
+
368
+ __vector float kf = vec_floor(y);
369
+ __vector float r = vec_sub(y, kf);
370
+
371
+ // Convert float to signed integer. Use vec_cts for PowerPC AltiVec
372
+ // compatibility.
373
+ __vector signed int k = vec_cts(kf, 0);
374
+ const __vector signed int min_k = vec_splats((signed int)-126);
375
+ const __vector signed int max_k = vec_splats((signed int)127);
376
+ k = vec_min(vec_max(k, min_k), max_k);
377
+
378
+ // Build 2^k from exponent bits
379
+ __vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
380
+ __vector unsigned int bits = (__vector unsigned int)exp_int;
381
+ bits = vec_sl(bits, vec_splats((unsigned int)23));
382
+ __vector float pow2k = (__vector float)bits;
383
+
384
+ // Improved minimax polynomial
385
+ __vector float poly = vec_madd(c5, r, c4);
386
+ poly = vec_madd(poly, r, c3);
387
+ poly = vec_madd(poly, r, c2);
388
+ poly = vec_madd(poly, r, c1);
389
+ poly = vec_madd(poly, r, one);
390
+
391
+ out.val[i] = vec_mul(pow2k, poly);
392
+ }
393
+ return FP32Vec8(out);
394
}
395
396
FP32Vec8 tanh() const {
365
- // TODO: Vectorize this
366
- AliasReg ar;
367
- ar.reg = reg;
368
- f32x4x4_t ret;
369
- ret.val[0][0] = std::tanh(ar.values[0]);
370
- ret.val[0][1] = std::tanh(ar.values[1]);
371
- ret.val[0][2] = std::tanh(ar.values[2]);
372
- ret.val[0][3] = std::tanh(ar.values[3]);
373
- ret.val[1][0] = std::tanh(ar.values[4]);
374
- ret.val[1][1] = std::tanh(ar.values[5]);
375
- ret.val[1][2] = std::tanh(ar.values[6]);
376
- ret.val[1][3] = std::tanh(ar.values[7]);
377
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
397
+ const __vector float one = vec_splats(1.0f);
398
+ const __vector float two = vec_splats(2.0f);
399
+ const __vector float zero = vec_splats(0.0f);
400
+ const __vector float sat = vec_splats(9.0f);
401
+
402
+ f32x4x2_t out;
403
+
404
+ for (int i = 0; i < 2; i++) {
405
+ __vector float x = reg.val[i];
406
+ __vector float ax = vec_abs(x);
407
+
408
+ __vector bool int mask = vec_cmpge(x, zero);
409
+ __vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
410
+
411
+ __vector bool int saturated = vec_cmpge(ax, sat);
412
+
413
+ __vector float two_x = vec_mul(x, two);
414
+ f32x4x2_t tmp;
415
+ tmp.val[0] = two_x;
416
+ tmp.val[1] = two_x;
417
+ FP32Vec8 temp_vec(tmp);
418
+ vector float e = temp_vec.exp().reg.val[0];
419
+
420
+ vector float num = vec_sub(e, one);
421
+ vector float den = vec_add(e, one);
422
+ vector float t = vec_div(num, den);
423
+
424
+ out.val[i] = vec_sel(t, sign, saturated);
425
+ }
426
+ return FP32Vec8(out);
427
}
428
429
FP32Vec8 er() const {
381
- // TODO: Vectorize this
382
- AliasReg ar;
383
- ar.reg = reg;
384
- f32x4x4_t ret;
385
- ret.val[0][0] = std::erf(ar.values[0]);
386
- ret.val[0][1] = std::erf(ar.values[1]);
387
- ret.val[0][2] = std::erf(ar.values[2]);
388
- ret.val[0][3] = std::erf(ar.values[3]);
389
- ret.val[1][0] = std::erf(ar.values[4]);
390
- ret.val[1][1] = std::erf(ar.values[5]);
391
- ret.val[1][2] = std::erf(ar.values[6]);
392
- ret.val[1][3] = std::erf(ar.values[7]);
393
- return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
430
+ const vector float a1 = vec_splats(0.254829592f);
431
+ const vector float a2 = vec_splats(-0.284496736f);
432
+ const vector float a3 = vec_splats(1.421413741f);
433
+ const vector float a4 = vec_splats(-1.453152027f);
434
+ const vector float a5 = vec_splats(1.061405429f);
435
+ const vector float p = vec_splats(0.3275911f);
436
+ const vector float one = vec_splats(1.0f);
437
+ const vector float zero = vec_splats(0.0f);
438
+ const vector float sat = vec_splats(6.0f);
439
+
440
+ f32x4x2_t ret;
441
+
442
+ for (int i = 0; i < 2; i++) {
443
+ vector float x = reg.val[i];
444
+ vector float ax = vec_abs(x);
445
+
446
+ vector bool int mask = vec_cmpge(x, zero);
447
+ vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
448
+
449
+ vector bool int saturated = vec_cmpge(ax, sat);
450
+
451
+ vector float t = vec_div(one, vec_madd(p, ax, one));
452
+
453
+ vector float poly = a5;
454
+ poly = vec_madd(poly, t, a4);
455
+ poly = vec_madd(poly, t, a3);
456
+ poly = vec_madd(poly, t, a2);
457
+ poly = vec_madd(poly, t, a1);
458
+ poly = vec_mul(poly, t);
459
+
460
+ vector float x_squared = vec_mul(x, x);
461
+ vector float neg_x_squared = vec_mul(vec_splats(-1.0f), x_squared);
462
+ f32x4x2_t tmp;
463
+ tmp.val[0] = neg_x_squared;
464
+ tmp.val[1] = neg_x_squared;
465
+ FP32Vec8 exp_input(tmp);
466
+ vector float exp_term = exp_input.exp().reg.val[0];
467
+
468
+ vector float y = vec_nmsub(poly, exp_term, one);
469
+ vector float erf_val = vec_mul(sign, y);
470
+
471
+ ret.val[i] = vec_sel(erf_val, sign, saturated);
472
+ }
473
+ return FP32Vec8(ret);
474
}
475
476
FP32Vec8 operator*(const FP32Vec8& b) const {
csrc/cpu/utils.hpp
+2
-2
@@ -76,14 +76,14 @@ inline int64_t get_available_l2_size() {
76
if (l2_cache_size == 0) {
77
l2_cache_size = 256 * 1024;
78
}
79
- return static_cast<int64_t>(l2_cache_size) >> 1; // use 50% of L2 cache
79
+ return static_cast<int64_t>(l2_cache_size) >> 1;
80
}();
81
return size;
82
#else
83
static int64_t size = []() {
84
auto caps = at::cpu::get_cpu_capabilities();
85
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
86
- return l2_cache_size >> 1; // use 50% of L2 cache
86
+ return l2_cache_size >> 1;
87
}();
88
return size;
89
#endif
vllm/model_executor/layers/activation.py
+30
-8
@@ -129,10 +129,12 @@ class SiluAndMul(CustomOp):
129
130
def __init__(self, *, compile_native: bool = True):
131
super().__init__(compile_native=compile_native)
132
- if current_platform.is_cuda_alike() or current_platform.is_xpu():
132
+ if (
133
+ current_platform.is_cuda_alike()
134
+ or current_platform.is_cpu()
135
+ or current_platform.is_xpu()
136
+ ):
137
self.op = torch.ops._C.silu_and_mul
134
- elif current_platform.is_cpu():
135
- self._forward_method = self.forward_native
138
139
@staticmethod
140
def forward_native(x: torch.Tensor) -> torch.Tensor:
@@ -150,6 +152,11 @@ class SiluAndMul(CustomOp):
152
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
153
return self.forward_cuda(x)
154
155
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
156
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
157
+ return self.forward_cuda(x)
158
+ return self.forward_native(x)
159
+
160
161
@CustomOp.register("silu_and_mul_with_clamp")
162
class SiluAndMulWithClamp(CustomOp):
@@ -417,14 +424,14 @@ class GeluAndMul(CustomOp):
424
self.op(out, x)
425
return out
426
420
- def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
421
- if self.op:
422
- return self.forward_cuda(x)
423
- return self.native(x)
424
-
427
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
428
return self.forward_cuda(x)
429
430
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
431
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
432
+ return self.forward_cuda(x)
433
+ return self.forward_native(x)
434
+
435
def extra_repr(self) -> str:
436
return f"approximate={repr(self.approximate)}"
437
@@ -526,6 +533,11 @@ class NewGELU(CustomOp):
533
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
534
return self.forward_cuda(x)
535
536
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
537
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
538
+ return self.forward_cuda(x)
539
+ return self.forward_native(x)
540
+
541
542
# --8<-- [start:gelu_fast]
543
@CustomOp.register("gelu_fast")
@@ -553,6 +565,11 @@ class FastGELU(CustomOp):
565
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
566
return self.forward_cuda(x)
567
568
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
569
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
570
+ return self.forward_cuda(x)
571
+ return self.forward_native(x)
572
+
573
574
# --8<-- [start:quick_gelu]
575
@CustomOp.register("quick_gelu")
@@ -581,6 +598,11 @@ class QuickGELU(CustomOp):
598
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
599
return self.forward_cuda(x)
600
601
+ def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
602
+ if current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC:
603
+ return self.forward_cuda(x)
604
+ return self.forward_native(x)
605
+
606
607
# --8<-- [start:relu2]
608
@CustomOp.register("relu2")