Skip to content

[Hardware][Power]Add Power VSX Attention Backend and fix l2 Cache Crash #40451

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

Merged
vllm-bot merged 11 commits into from
May 5, 2026
Merged

[Hardware][Power]Add Power VSX Attention Backend and fix l2 Cache Crash #40451

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
19a2605
Add vsx backend
Apr 20, 2026
5f5a58c
fix l2 cache error
Apr 20, 2026
e63aa24
fix pre-commit
Apr 21, 2026
f454e23
Address review comments
Apr 21, 2026
397a9e7
fix bf16 issue
Apr 21, 2026
e7bf4a2
fp8 not on power
Apr 30, 2026
40cff6d
fix make pair
Apr 30, 2026
5f111ad
fix l2 func
Apr 30, 2026
12dd1c9
fix
Apr 30, 2026
9eb18c0
Merge branch 'main' into feat/power_vsx_backend
bigPYJ1151 May 1, 2026
714bb6b
Merge branch 'main' into feat/power_vsx_backend
Akashcodes732 May 2, 2026
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
4 changes: 4 additions & 0 deletions csrc/cpu/cpu_attn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
isa = cpu_attention::ISA::NEON;
} else if (isa_hint == "vxe") {
isa = cpu_attention::ISA::VXE;
} else if (isa_hint == "vsx") {
isa = cpu_attention::ISA::VSX;
} else {
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
}
Expand Down Expand Up @@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache(
return cpu_attention::ISA::NEON;
} else if (isa == "vxe") {
return cpu_attention::ISA::VXE;
} else if (isa == "vsx") {
return cpu_attention::ISA::VSX;
} else {
TORCH_CHECK(false, "Invalid ISA type: " + isa);
}
Expand Down
5 changes: 4 additions & 1 deletion csrc/cpu/cpu_attn_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
#include "cpu/utils.hpp"

namespace cpu_attention {
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };

// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
enum class Fp8KVCacheDataType {
Expand Down Expand Up @@ -164,6 +164,9 @@ struct AttentionMetadata {
case ISA::VXE:
ss << "VXE, ";
break;
case ISA::VSX:
ss << "VSX, ";
break;
}
ss << "workitem_group_num: " << workitem_group_num
<< ", reduction_item_num: " << reduction_item_num
Expand Down
10 changes: 7 additions & 3 deletions csrc/cpu/cpu_attn_vec.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,11 @@ namespace {
// covering the lower 16 and upper 16 positions.
// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
// Note: FP8 KV cache is not supported on PowerPC (VSX).
template <typename kv_cache_t>
FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
const kv_cache_t* ptr) {
#if !defined(__powerpc__)
Comment thread
Akashcodes732 marked this conversation as resolved.
Outdated
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
// BF16 container, but values are in the FP16 exponent range (bias 15 not
// 127).
Expand All @@ -25,10 +27,12 @@ FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
vec_op::fp8_e5m2_tag{});
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
} else {
} else
#endif
{
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
return {vec_op::FP32Vec16(load_vec_t(ptr)),
vec_op::FP32Vec16(load_vec_t(ptr + 16))};
return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
vec_op::FP32Vec16(load_vec_t(ptr + 16)));
}
}

Expand Down
Loading
Loading