iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[HARDWARE][POWER] Enable fp16 support for PowerPC (#46135)

Browse Source 
Signed-off-by: Rukhaiya <bibirukhaiya123@gmail.com>
This commit is contained in:
Rukhaiya2004
2026-06-23 13:24:49 +00:00
committed by GitHub
parent 1bf149f334
commit 9f5117820f
6 changed files with 332 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/pos_encoding.cpp
+154 -1
View File
@@ -1,4 +1,3 @@
#include "cpu_types.hpp"
namespace {
@@ -97,6 +96,91 @@ void rotary_embedding_impl(
}
}
template <>
void rotary_embedding_impl<c10::Half>(
const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
const int num_heads, const int num_kv_heads, const int head_size,
const int num_tokens) {
using scalar_vec_t = vec_op::FP16Vec8;
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
const int embed_dim = rot_dim / 2;
bool flag = (embed_dim % VEC_ELEM_NUM == 0);
const int loop_upper = flag ? embed_dim : embed_dim - VEC_ELEM_NUM;
auto compute_loop = [&](const int64_t token_head, const c10::Half* cache_ptr,
c10::Half* qk) {
int j = 0;
for (; j < loop_upper; j += VEC_ELEM_NUM) {
const int rot_offset = j;
const int x_index = rot_offset;
const int y_index = embed_dim + rot_offset;
const int64_t out_x = token_head + x_index;
const int64_t out_y = token_head + y_index;
const vec_op::FP16Vec8 cos_fp16(cache_ptr + x_index);
const vec_op::FP16Vec8 sin_fp16(cache_ptr + y_index);
const vec_op::FP16Vec8 q_x_fp16(qk + out_x);
const vec_op::FP16Vec8 q_y_fp16(qk + out_y);
const vec_op::FP32Vec8 fp32_cos(cos_fp16);
const vec_op::FP32Vec8 fp32_sin(sin_fp16);
const vec_op::FP32Vec8 fp32_q_x(q_x_fp16);
const vec_op::FP32Vec8 fp32_q_y(q_y_fp16);
auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin;
auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin;
vec_op::FP16Vec8(out1).save(qk + out_x);
vec_op::FP16Vec8(out2).save(qk + out_y);
}
if (!flag) {
for (; j < embed_dim; ++j) {
const int x_index = j;
const int y_index = embed_dim + j;
const int64_t out_x = token_head + x_index;
const int64_t out_y = token_head + y_index;
const float fp32_cos = static_cast<float>(cache_ptr[x_index]);
const float fp32_sin = static_cast<float>(cache_ptr[y_index]);
const float fp32_q_x = static_cast<float>(qk[out_x]);
const float fp32_q_y = static_cast<float>(qk[out_y]);
qk[out_x] =
static_cast<c10::Half>(fp32_q_x * fp32_cos - fp32_q_y * fp32_sin);
qk[out_y] =
static_cast<c10::Half>(fp32_q_y * fp32_cos + fp32_q_x * fp32_sin);
}
}
};
#pragma omp parallel for
for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
int64_t pos = positions[token_idx];
const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
for (int i = 0; i < num_heads; ++i) {
const int head_idx = i;
const int64_t token_head =
token_idx * query_stride + head_idx * head_size;
compute_loop(token_head, cache_ptr, query);
}
if (key != nullptr) {
for (int i = 0; i < num_kv_heads; ++i) {
const int head_idx = i;
const int64_t token_head =
token_idx * key_stride + head_idx * head_size;
compute_loop(token_head, cache_ptr, key);
}
}
}
}
template <typename scalar_t>
void rotary_embedding_gptj_impl(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
@@ -174,6 +258,75 @@ void rotary_embedding_gptj_impl(
}
}
}
template <>
void rotary_embedding_gptj_impl<c10::Half>(
const int64_t* __restrict__ positions, c10::Half* __restrict__ query,
c10::Half* __restrict__ key, const c10::Half* __restrict__ cos_sin_cache,
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
const int num_heads, const int num_kv_heads, const int head_size,
const int num_tokens) {
const int embed_dim = rot_dim / 2;
#pragma omp parallel for collapse(2)
for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
for (int i = 0; i < num_heads; ++i) {
int64_t pos = positions[token_idx];
const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
const c10::Half* cos_cache_ptr = cache_ptr;
const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
const int head_idx = i;
const int64_t token_head =
token_idx * query_stride + head_idx * head_size;
c10::Half* head_query = token_head + query;
for (int j = 0; j < embed_dim; j += 1) {
const int rot_offset = j;
const int x_index = 2 * rot_offset;
const int y_index = 2 * rot_offset + 1;
const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
const float x = static_cast<float>(head_query[x_index]);
const float y = static_cast<float>(head_query[y_index]);
head_query[x_index] = static_cast<c10::Half>(x * cos - y * sin);
head_query[y_index] = static_cast<c10::Half>(y * cos + x * sin);
}
}
}
if (key == nullptr) {
return;
}
#pragma omp parallel for collapse(2)
for (int token_idx = 0; token_idx < num_tokens; ++token_idx) {
for (int i = 0; i < num_kv_heads; ++i) {
int64_t pos = positions[token_idx];
const c10::Half* cache_ptr = cos_sin_cache + pos * rot_dim;
const c10::Half* cos_cache_ptr = cache_ptr;
const c10::Half* sin_cache_ptr = cache_ptr + embed_dim;
const int head_idx = i;
const int64_t token_head = token_idx * key_stride + head_idx * head_size;
c10::Half* head_key = key + token_head;
for (int j = 0; j < embed_dim; j += 1) {
const int rot_offset = j;
const int x_index = 2 * rot_offset;
const int y_index = 2 * rot_offset + 1;
const float cos = static_cast<float>(cos_cache_ptr[rot_offset]);
const float sin = static_cast<float>(sin_cache_ptr[rot_offset]);
const float x = static_cast<float>(head_key[x_index]);
const float y = static_cast<float>(head_key[y_index]);
head_key[x_index] = static_cast<c10::Half>(x * cos - y * sin);
head_key[y_index] = static_cast<c10::Half>(y * cos + x * sin);
}
}
}
}
}; // namespace
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,