chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,175 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
namespace {
|
||||
template <typename scalar_t, vec_op::FP32Vec8 (*func)(const vec_op::FP32Vec8&),
|
||||
bool is_gated>
|
||||
void activation_kernel(int num_tokens, int d, scalar_t* __restrict__ input,
|
||||
scalar_t* __restrict__ output) {
|
||||
using scalar_vec_t = vec_op::vec_t<scalar_t>;
|
||||
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
|
||||
|
||||
TORCH_CHECK(d % VEC_ELEM_NUM == 0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int i = 0; i < num_tokens; ++i) {
|
||||
for (int j = 0; j < d; j += VEC_ELEM_NUM) {
|
||||
int start = i * d;
|
||||
if constexpr (is_gated) {
|
||||
start *= 2;
|
||||
}
|
||||
|
||||
const scalar_vec_t x(input + start + j);
|
||||
const vec_op::FP32Vec8 f32_x(x);
|
||||
vec_op::FP32Vec8 f32_ans = func(f32_x);
|
||||
|
||||
if constexpr (is_gated) {
|
||||
const scalar_vec_t y(input + start + d + j);
|
||||
const vec_op::FP32Vec8 f32_y(y);
|
||||
f32_ans = f32_y * f32_ans;
|
||||
}
|
||||
|
||||
const scalar_vec_t result(f32_ans);
|
||||
result.save(output + i * d + j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 silu_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 zeros(0.0);
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
return x / (ones + (zeros - x).exp());
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 gelu_new_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
const vec_op::FP32Vec8 w1(0.79788456f);
|
||||
const vec_op::FP32Vec8 w2(0.044715f);
|
||||
const vec_op::FP32Vec8 w3(0.5);
|
||||
const vec_op::FP32Vec8 x3 = x * x * x;
|
||||
const vec_op::FP32Vec8 t = (w1 * (x + w2 * x3)).tanh();
|
||||
return w3 * x * (ones + t);
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 gelu_fast_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
const vec_op::FP32Vec8 w1(0.79788456f);
|
||||
const vec_op::FP32Vec8 w2(0.044715f);
|
||||
const vec_op::FP32Vec8 w3(0.5);
|
||||
const vec_op::FP32Vec8 t = (x * w1 * (ones + x * w2 * x)).tanh();
|
||||
return w3 * x * (ones + t);
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 gelu_quick_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 zeros(0.0);
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
const vec_op::FP32Vec8 w1(1.702f);
|
||||
return x / (ones + (zeros - w1 * x).exp());
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 gelu_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
const vec_op::FP32Vec8 w1(M_SQRT1_2);
|
||||
const vec_op::FP32Vec8 w2(0.5);
|
||||
return x * w2 * (ones + (x * w1).er());
|
||||
}
|
||||
|
||||
FORCE_INLINE vec_op::FP32Vec8 gelu_tanh_act(const vec_op::FP32Vec8& x) {
|
||||
const vec_op::FP32Vec8 ones(1.0);
|
||||
const vec_op::FP32Vec8 w1(M_SQRT2 * M_2_SQRTPI * 0.5);
|
||||
const vec_op::FP32Vec8 w2(0.5);
|
||||
const vec_op::FP32Vec8 w3(0.044715);
|
||||
const vec_op::FP32Vec8 x_3 = x * x * x;
|
||||
const vec_op::FP32Vec8 inner = w1 * (x + x_3 * w3);
|
||||
return x * w2 * (ones + inner.tanh());
|
||||
}
|
||||
}; // namespace
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1) / 2;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(silu_and_mul_impl)
|
||||
activation_kernel<scalar_t, silu_act, true>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(silu_and_mul_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1) / 2;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_and_mul_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_and_mul_impl)
|
||||
activation_kernel<scalar_t, gelu_act, true>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_and_mul_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1) / 2;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "gelu_tanh_and_mul_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_tanh_and_mul_impl)
|
||||
activation_kernel<scalar_t, gelu_tanh_act, true>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(),
|
||||
out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_tanh_and_mul_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
|
||||
activation_kernel<scalar_t, gelu_tanh_act, false>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_new(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_new_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_new_impl)
|
||||
activation_kernel<scalar_t, gelu_new_act, false>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_new_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_fast(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_fast_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_fast_impl)
|
||||
activation_kernel<scalar_t, gelu_fast_act, false>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_fast_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void gelu_quick(torch::Tensor& out, torch::Tensor& input) {
|
||||
int num_tokens = input.numel() / input.size(-1);
|
||||
int d = input.size(-1);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_quick_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(gelu_quick_impl)
|
||||
activation_kernel<scalar_t, gelu_quick_act, false>(
|
||||
num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
|
||||
CPU_KERNEL_GUARD_OUT(gelu_quick_impl)
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
|
||||
#include <ATen/ops/empty.h>
|
||||
#include <ATen/ops/gelu.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
|
||||
constexpr uint32_t ActivationLutSize = 1u << 16;
|
||||
|
||||
at::Tensor gelu_reference(const at::Tensor& x) { return at::gelu(x, "none"); }
|
||||
|
||||
void maybe_init_activation_lut_bf16(
|
||||
uint16_t* lut, std::once_flag& once,
|
||||
at::Tensor (*activation)(const at::Tensor&)) {
|
||||
std::call_once(once, [&]() {
|
||||
auto lut_input =
|
||||
at::empty({static_cast<int64_t>(ActivationLutSize)},
|
||||
at::TensorOptions().device(at::kCPU).dtype(at::kFloat));
|
||||
auto* lut_input_ptr = lut_input.data_ptr<float>();
|
||||
#pragma omp parallel for
|
||||
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
|
||||
lut_input_ptr[i] = c10::detail::f32_from_bits(static_cast<uint16_t>(i));
|
||||
}
|
||||
|
||||
auto lut_output = activation(lut_input);
|
||||
const auto* lut_output_ptr = lut_output.data_ptr<float>();
|
||||
#pragma omp parallel for
|
||||
for (uint32_t i = 0; i < ActivationLutSize; ++i) {
|
||||
lut[i] = c10::detail::round_to_nearest_even(lut_output_ptr[i]);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const uint16_t* lut, const char* op_name) {
|
||||
TORCH_CHECK(input.scalar_type() == at::kBFloat16, op_name,
|
||||
": input must be bfloat16");
|
||||
TORCH_CHECK(out.scalar_type() == at::kBFloat16, op_name,
|
||||
": out must be bfloat16");
|
||||
TORCH_CHECK(input.is_contiguous(), op_name, ": input must be contiguous");
|
||||
TORCH_CHECK(out.is_contiguous(), op_name, ": out must be contiguous");
|
||||
|
||||
const auto* src =
|
||||
reinterpret_cast<const uint16_t*>(input.data_ptr<at::BFloat16>());
|
||||
auto* dst = reinterpret_cast<uint16_t*>(out.data_ptr<at::BFloat16>());
|
||||
const int64_t n = input.numel();
|
||||
|
||||
CPU_KERNEL_GUARD_IN(activation_lut_bf16_impl)
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
dst[i] = lut[src[i]];
|
||||
}
|
||||
CPU_KERNEL_GUARD_OUT(activation_lut_bf16_impl)
|
||||
}
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const std::string& activation) {
|
||||
if (activation == "gelu") {
|
||||
static std::array<uint16_t, ActivationLutSize> lut{};
|
||||
static std::once_flag once;
|
||||
maybe_init_activation_lut_bf16(lut.data(), once, gelu_reference);
|
||||
activation_lut_bf16(out, input, lut.data(), "gelu_lut");
|
||||
return;
|
||||
}
|
||||
|
||||
TORCH_CHECK(false, "Unsupported activation: ", activation);
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
#ifndef CPU_ARCH_MACROS_H
|
||||
#define CPU_ARCH_MACROS_H
|
||||
|
||||
// x86_64
|
||||
#ifdef __x86_64__
|
||||
#define FAST_SPINNING _mm_pause();
|
||||
|
||||
#ifdef __AVX512F__
|
||||
#define DEFINE_FAST_EXP \
|
||||
const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); \
|
||||
const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); \
|
||||
const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); \
|
||||
const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); \
|
||||
const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); \
|
||||
const __m512 vec_exp_log2ef = \
|
||||
_mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); \
|
||||
const __m512 vec_half = _mm512_set1_ps(0.5f); \
|
||||
const __m512 vec_one = _mm512_set1_ps(1.f); \
|
||||
const __m512 vec_zero = _mm512_set1_ps(0.f); \
|
||||
const __m512 vec_two = _mm512_set1_ps(2.f); \
|
||||
const __m512 vec_ln2f = \
|
||||
_mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); \
|
||||
const __m512 vec_ln_flt_min = \
|
||||
_mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); \
|
||||
const __m512 vec_ln_flt_max = \
|
||||
_mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); \
|
||||
const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); \
|
||||
const int n_mantissa_bits = 23; \
|
||||
auto fast_exp = [&](const vec_op::FP32Vec16& vec) __attribute__(( \
|
||||
always_inline)) { \
|
||||
__m512 values = vec.reg; \
|
||||
auto less_ln_flt_min_mask = \
|
||||
_mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); \
|
||||
auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); \
|
||||
vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); \
|
||||
auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); \
|
||||
auto vec_fx_i = _mm512_cvt_roundps_epi32( \
|
||||
vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); \
|
||||
vec_fx = _mm512_cvtepi32_ps(vec_fx_i); \
|
||||
auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); \
|
||||
auto vec_res = \
|
||||
_mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); \
|
||||
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); \
|
||||
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); \
|
||||
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); \
|
||||
vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); \
|
||||
auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); \
|
||||
auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); \
|
||||
auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); \
|
||||
vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); \
|
||||
auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); \
|
||||
vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, \
|
||||
vec_two_pow_n, vec_zero); \
|
||||
vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); \
|
||||
vec_res = _mm512_mul_ps(vec_res, vec_two); \
|
||||
vec_op::FP32Vec16 res(vec_res); \
|
||||
return res; \
|
||||
};
|
||||
#endif
|
||||
|
||||
#endif
|
||||
|
||||
#ifdef __aarch64__
|
||||
// Implementation of neon_expf copied from Arm Optimized Routines (expf
|
||||
// AdvSIMD)
|
||||
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
|
||||
//
|
||||
// Additional fast exponential intended for cases where outputs will be
|
||||
// downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
|
||||
// for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
|
||||
// clamps inputs outside this range to 0 / inf. Implementation is similar to
|
||||
// exp_u20, but:
|
||||
// - uses a third degree polynomial approximation for exp(r) instead of a
|
||||
// fifth degree one, with coefficients re-tuned.
|
||||
// - does not split natural log (ln) into high / low parts
|
||||
// - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
|
||||
// exp(x) = 2^n (exp(r))
|
||||
// r = x - n*ln2, with n = round(x/ln2)
|
||||
// exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
|
||||
// n = round(x / ln2), r = x - n*ln2
|
||||
#include <limits>
|
||||
#define DEFINE_FAST_EXP \
|
||||
const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
|
||||
const float ln2_hi = 0x1.62e4p-1f; \
|
||||
const float ln2_lo = 0x1.7f7d1cp-20f; \
|
||||
const float c0 = 0x1.0e4020p-7f; \
|
||||
const float c2 = 0x1.555e66p-3f; \
|
||||
const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2}; \
|
||||
const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000); \
|
||||
const float32x4_t c1 = vdupq_n_f32(0x1.573e2ep-5f); \
|
||||
const float32x4_t c3 = vdupq_n_f32(0x1.fffdb6p-2f); \
|
||||
const float32x4_t c4 = vdupq_n_f32(0x1.ffffecp-1f); \
|
||||
const float32x4_t pos_special_bound = vdupq_n_f32(0x1.5d5e2ap+6f); \
|
||||
const float32x4_t neg_special_bound = vnegq_f32(pos_special_bound); \
|
||||
const float32x4_t inf = \
|
||||
vdupq_n_f32(std::numeric_limits<float>::infinity()); \
|
||||
const float32x4_t zero = vdupq_n_f32(0.0f); \
|
||||
auto neon_expf = [&](float32x4_t values) __attribute__((always_inline)) { \
|
||||
float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
|
||||
float32x4_t r = vfmsq_laneq_f32(values, n, ln2_c02, 0); \
|
||||
r = vfmsq_laneq_f32(r, n, ln2_c02, 1); \
|
||||
uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
|
||||
float32x4_t scale = vreinterpretq_f32_u32(vaddq_u32(e, exponent_bias)); \
|
||||
float32x4_t r2 = vmulq_f32(r, r); \
|
||||
float32x4_t p = vfmaq_laneq_f32(c1, r, ln2_c02, 2); \
|
||||
float32x4_t q = vfmaq_laneq_f32(c3, r, ln2_c02, 3); \
|
||||
q = vfmaq_f32(q, p, r2); \
|
||||
p = vmulq_f32(c4, r); \
|
||||
float32x4_t poly = vfmaq_f32(p, q, r2); \
|
||||
poly = vfmaq_f32(scale, poly, scale); \
|
||||
const uint32x4_t hi_mask = vcgeq_f32(values, pos_special_bound); \
|
||||
const uint32x4_t lo_mask = vcleq_f32(values, neg_special_bound); \
|
||||
poly = vbslq_f32(hi_mask, inf, poly); \
|
||||
return vbslq_f32(lo_mask, zero, poly); \
|
||||
}; \
|
||||
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { \
|
||||
float32x4x4_t result; \
|
||||
result.val[0] = neon_expf(vec.reg.val[0]); \
|
||||
result.val[1] = neon_expf(vec.reg.val[1]); \
|
||||
result.val[2] = neon_expf(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
}; \
|
||||
const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
|
||||
const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
|
||||
constexpr float ln2 = 0x1.62e43p-1f; \
|
||||
constexpr float f_c2 = 0x1.5592ecp-3f; \
|
||||
const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
|
||||
auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
|
||||
always_inline)) { \
|
||||
const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
|
||||
const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
|
||||
float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
|
||||
float32x4_t r = vfmsq_n_f32(values, n, ln2); \
|
||||
uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
|
||||
float32x4_t r2 = vmulq_f32(r, r); \
|
||||
float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
|
||||
float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
|
||||
float32x4_t p = vfmaq_f32(s, q, r2); \
|
||||
float32x4_t y = \
|
||||
vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
|
||||
y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
|
||||
y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
|
||||
return y; \
|
||||
}; \
|
||||
auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { \
|
||||
float32x4x4_t result; \
|
||||
result.val[0] = neon_expf_f16(vec.reg.val[0]); \
|
||||
result.val[1] = neon_expf_f16(vec.reg.val[1]); \
|
||||
result.val[2] = neon_expf_f16(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf_f16(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
};
|
||||
|
||||
#endif // __aarch64__
|
||||
|
||||
// RISC-V RVV
|
||||
#ifdef __riscv_v
|
||||
#include <riscv_vector.h>
|
||||
|
||||
#ifdef __riscv_zihintpause
|
||||
#define FAST_SPINNING __riscv_pause();
|
||||
#endif
|
||||
|
||||
// FP32Vec16::exp() in cpu_types_riscv.hpp already implements the full
|
||||
// polynomial approximation for RVV, so we simply delegate to it.
|
||||
#define DEFINE_FAST_EXP \
|
||||
auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { return vec.exp(); };
|
||||
|
||||
#endif // __riscv_v
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,251 @@
|
||||
#include "cpu_attn_dispatch_generated.h"
|
||||
|
||||
// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
|
||||
// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
|
||||
static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
|
||||
const std::string& kv_cache_dtype) {
|
||||
if (kv_cache_dtype == "fp8_e5m2")
|
||||
return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
|
||||
if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
|
||||
return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
|
||||
return cpu_attention::Fp8KVCacheDataType::kAuto;
|
||||
}
|
||||
|
||||
bool cpu_attn_has_isa(const std::string& isa) {
|
||||
if (isa == "rvv") {
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
||||
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
torch::Tensor get_scheduler_metadata(
|
||||
const int64_t num_req, const int64_t num_heads_q,
|
||||
const int64_t num_heads_kv, const int64_t head_dim,
|
||||
const torch::Tensor& seq_lens, at::ScalarType dtype,
|
||||
const torch::Tensor& query_start_loc, const bool causal,
|
||||
const int64_t window_size, const std::string& isa_hint,
|
||||
const bool enable_kv_split,
|
||||
const std::optional<torch::Tensor>& dynamic_causal) {
|
||||
cpu_attention::ISA isa;
|
||||
if (isa_hint == "amx") {
|
||||
isa = cpu_attention::ISA::AMX;
|
||||
} else if (isa_hint == "vec") {
|
||||
isa = cpu_attention::ISA::VEC;
|
||||
} else if (isa_hint == "vec16") {
|
||||
isa = cpu_attention::ISA::VEC16;
|
||||
} else if (isa_hint == "neon") {
|
||||
isa = cpu_attention::ISA::NEON;
|
||||
} else if (isa_hint == "vxe") {
|
||||
isa = cpu_attention::ISA::VXE;
|
||||
} else if (isa_hint == "rvv") {
|
||||
isa = cpu_attention::ISA::RVV;
|
||||
} else if (isa_hint == "vsx") {
|
||||
isa = cpu_attention::ISA::VSX;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
|
||||
}
|
||||
|
||||
cpu_attention::AttentionScheduler::ScheduleInput input;
|
||||
input.num_reqs = num_req;
|
||||
input.num_heads_q = num_heads_q;
|
||||
input.num_heads_kv = num_heads_kv;
|
||||
input.head_dim = head_dim;
|
||||
input.query_start_loc = query_start_loc.data_ptr<int32_t>();
|
||||
input.seq_lens = seq_lens.data_ptr<int32_t>();
|
||||
|
||||
input.sliding_window_size = window_size;
|
||||
input.causal = causal;
|
||||
input.isa = isa;
|
||||
input.enable_kv_split = enable_kv_split;
|
||||
input.dynamic_causal =
|
||||
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
|
||||
input.elem_size = sizeof(scalar_t);
|
||||
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
|
||||
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
|
||||
input.output_buffer_elem_size =
|
||||
sizeof(attn_impl::partial_output_buffer_t);
|
||||
input.max_num_q_per_iter = attn_impl::MaxQHeadNumPerIteration;
|
||||
input.kv_block_alignment = attn_impl::BlockSizeAlignment;
|
||||
});
|
||||
});
|
||||
|
||||
cpu_attention::AttentionScheduler scheduler;
|
||||
torch::Tensor metadata = scheduler.schedule(input);
|
||||
return metadata;
|
||||
}
|
||||
|
||||
void cpu_attn_reshape_and_cache(
|
||||
const torch::Tensor& key, // [token_num, head_num, head_size]
|
||||
const torch::Tensor& value, // [token_num, head_num, head_size]
|
||||
torch::Tensor&
|
||||
key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
torch::Tensor&
|
||||
value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
const torch::Tensor& slot_mapping, const std::string& isa,
|
||||
const double k_scale = 1.0, const double v_scale = 1.0,
|
||||
const std::string& kv_cache_dtype = "auto") {
|
||||
TORCH_CHECK_EQ(key.dim(), 3);
|
||||
TORCH_CHECK_EQ(value.dim(), 3);
|
||||
TORCH_CHECK_EQ(key_cache.dim(), 4);
|
||||
TORCH_CHECK_EQ(value_cache.dim(), 4);
|
||||
TORCH_CHECK_EQ(key.stride(2), 1);
|
||||
TORCH_CHECK_EQ(value.stride(2), 1);
|
||||
|
||||
const int64_t kv_cache_idx =
|
||||
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
|
||||
const bool is_fp8 = (kv_cache_idx != 0);
|
||||
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"key_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"value_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
|
||||
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
|
||||
}
|
||||
|
||||
const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
|
||||
const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
|
||||
|
||||
const int64_t token_num = key.size(0);
|
||||
const int64_t head_num = key.size(1);
|
||||
const int64_t head_dim = key.size(2);
|
||||
const int64_t num_blocks = key_cache.size(0);
|
||||
const int64_t num_blocks_stride = key_cache.stride(0);
|
||||
const int64_t cache_head_num_stride = key_cache.stride(1);
|
||||
const int64_t block_size = key_cache.size(2);
|
||||
const int64_t block_size_stride = key_cache.stride(2);
|
||||
|
||||
cpu_attention::ISA isa_tag = [&]() {
|
||||
if (isa == "amx") {
|
||||
return cpu_attention::ISA::AMX;
|
||||
} else if (isa == "vec") {
|
||||
return cpu_attention::ISA::VEC;
|
||||
} else if (isa == "vec16") {
|
||||
return cpu_attention::ISA::VEC16;
|
||||
} else if (isa == "neon") {
|
||||
return cpu_attention::ISA::NEON;
|
||||
} else if (isa == "vxe") {
|
||||
return cpu_attention::ISA::VXE;
|
||||
} else if (isa == "rvv") {
|
||||
return cpu_attention::ISA::RVV;
|
||||
} else if (isa == "vsx") {
|
||||
return cpu_attention::ISA::VSX;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid ISA type: " + isa);
|
||||
}
|
||||
}();
|
||||
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
|
||||
isa_tag == cpu_attention::ISA::VEC,
|
||||
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
|
||||
using kv_t = typename attn_impl::kv_cache_t;
|
||||
attn_impl::reshape_and_cache(
|
||||
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
|
||||
reinterpret_cast<kv_t*>(key_cache.data_ptr()),
|
||||
reinterpret_cast<kv_t*>(value_cache.data_ptr()),
|
||||
slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
|
||||
value.stride(0), head_num, key.stride(1), value.stride(1),
|
||||
num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
|
||||
block_size_stride, k_inv, v_inv);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
void cpu_attention_with_kv_cache(
|
||||
const torch::Tensor& query, // [num_tokens, num_heads, head_size]
|
||||
const torch::Tensor&
|
||||
key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
const torch::Tensor&
|
||||
value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
torch::Tensor& output, // [num_tokens, num_heads, head_size]
|
||||
const torch::Tensor& query_start_loc, // [num_tokens + 1]
|
||||
const torch::Tensor& seq_lens, // [num_tokens]
|
||||
const double scale, const bool causal,
|
||||
const std::optional<torch::Tensor>& alibi_slopes, // [num_heads]
|
||||
const int64_t sliding_window,
|
||||
const torch::Tensor& block_table, // [num_tokens, max_block_num]
|
||||
const double softcap, const torch::Tensor& scheduler_metadata,
|
||||
const std::optional<torch::Tensor>& s_aux, // [num_heads]
|
||||
const std::optional<torch::Tensor>& dynamic_causal, // [num_reqs]
|
||||
const double k_scale = 1.0, const double v_scale = 1.0,
|
||||
const std::string& kv_cache_dtype = "auto") {
|
||||
TORCH_CHECK_EQ(query.dim(), 3);
|
||||
TORCH_CHECK_EQ(query.stride(2), 1);
|
||||
TORCH_CHECK_EQ(key_cache.dim(), 4);
|
||||
TORCH_CHECK_EQ(value_cache.dim(), 4);
|
||||
|
||||
const int64_t kv_cache_idx =
|
||||
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
|
||||
const bool is_fp8 = (kv_cache_idx != 0);
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"key_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"value_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
|
||||
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
|
||||
}
|
||||
|
||||
cpu_attention::AttentionInput input;
|
||||
input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
|
||||
scheduler_metadata.data_ptr());
|
||||
input.num_tokens = query.size(0);
|
||||
input.num_heads = query.size(1);
|
||||
input.num_kv_heads = key_cache.size(1);
|
||||
input.block_size = key_cache.size(2);
|
||||
input.query = query.data_ptr();
|
||||
input.query_num_tokens_stride = query.stride(0);
|
||||
input.query_num_heads_stride = query.stride(1);
|
||||
input.cache_num_blocks_stride = key_cache.stride(0);
|
||||
input.cache_num_kv_heads_stride = key_cache.stride(1);
|
||||
input.blt_num_tokens_stride = block_table.stride(0);
|
||||
input.key_cache = key_cache.data_ptr();
|
||||
input.value_cache = value_cache.data_ptr();
|
||||
input.output = output.data_ptr();
|
||||
input.query_start_loc = query_start_loc.data_ptr<int32_t>();
|
||||
input.seq_lens = seq_lens.data_ptr<int32_t>();
|
||||
input.block_table = block_table.data_ptr<int32_t>();
|
||||
input.alibi_slopes =
|
||||
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
|
||||
input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
|
||||
input.dynamic_causal =
|
||||
dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
|
||||
input.scale = scale;
|
||||
input.causal = causal;
|
||||
input.sliding_window_size = sliding_window;
|
||||
input.softcap = static_cast<float>(softcap);
|
||||
|
||||
if (is_fp8) {
|
||||
input.k_scale_fp8 = static_cast<float>(k_scale);
|
||||
input.v_scale_fp8 = static_cast<float>(v_scale);
|
||||
TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
|
||||
input.metadata->isa == cpu_attention::ISA::VEC,
|
||||
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
|
||||
CPU_ATTN_DISPATCH(
|
||||
query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
|
||||
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
|
||||
0);
|
||||
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
|
||||
mainloop(&input);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,636 @@
|
||||
#ifndef CPU_ATTN_AMX_HPP
|
||||
#define CPU_ATTN_AMX_HPP
|
||||
|
||||
#include "cpu_attn_fp8.hpp"
|
||||
#include "cpu_attn_impl.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
namespace {
|
||||
// AMX specific
|
||||
constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
|
||||
constexpr static int64_t AMX_TILE_ROW_NUM = 16;
|
||||
constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
|
||||
|
||||
typedef struct __tile_config {
|
||||
uint8_t palette_id = 1;
|
||||
uint8_t start_row = 0;
|
||||
uint8_t reserved_0[14] = {0};
|
||||
uint16_t colsb[16] = {0};
|
||||
uint8_t rows[16] = {0};
|
||||
} __tilecfg;
|
||||
|
||||
// 2-2-4 pattern, for 16 < m <= 32
|
||||
// TILE 0, 1: load A matrix, row num should be 16, m - 16
|
||||
// TILE 2, 3: load B matrix, row num should be 16
|
||||
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
|
||||
// m - 16, m - 16
|
||||
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
|
||||
template <typename q_buffer_t, typename kv_cache_t>
|
||||
class TileGemm224 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
|
||||
void* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224");
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
TORCH_CHECK(false, "Unsupported kv cache type for TileGemm224");
|
||||
}
|
||||
};
|
||||
|
||||
// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
|
||||
for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
|
||||
.save(dst + r * 32);
|
||||
} else {
|
||||
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
|
||||
.save(dst + r * 32);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For FP8: dequant src into scratch and return scratch.
|
||||
// For BF16: return src directly (scratch is unused; the compiler elides it).
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
|
||||
c10::BFloat16* scratch) {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
|
||||
return scratch;
|
||||
} else {
|
||||
return reinterpret_cast<const c10::BFloat16*>(src);
|
||||
}
|
||||
}
|
||||
|
||||
// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm224<c10::BFloat16, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
|
||||
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
|
||||
|
||||
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
// BF16 path: scratch_elems=1 so the scratch array is eliminated by the
|
||||
// compiler.
|
||||
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
|
||||
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
c10::BFloat16* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
const int32_t k_times =
|
||||
dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
|
||||
|
||||
c10::BFloat16* __restrict__ a_tile_0 = a_tile;
|
||||
c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
|
||||
const int64_t a_tile_stride = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// q_buffer is prepacked
|
||||
return AMX_TILE_ROW_BYTES;
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// logits_buffer is row-major
|
||||
return lda * sizeof(c10::BFloat16);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
|
||||
kv_cache_t* __restrict__ b_tile_2 = b_tile;
|
||||
kv_cache_t* __restrict__ b_tile_3 = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// k_cache is prepacked
|
||||
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// v_cache is prepacked
|
||||
return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
// k_cache, v_cache are prepacked
|
||||
const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
|
||||
|
||||
// logits_buffer, output_buffer are not prepacked
|
||||
float* __restrict__ c_tile_4 = c_tile;
|
||||
float* __restrict__ c_tile_5 =
|
||||
c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
float* __restrict__ c_tile_6 = c_tile + AMX_TILE_ROW_NUM * ldc;
|
||||
float* __restrict__ c_tile_7 =
|
||||
c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
const int32_t c_tile_stride = ldc * sizeof(float);
|
||||
|
||||
if (accum_c) {
|
||||
_tile_loadd(4, c_tile_4, c_tile_stride);
|
||||
_tile_loadd(5, c_tile_5, c_tile_stride);
|
||||
_tile_loadd(6, c_tile_6, c_tile_stride);
|
||||
_tile_loadd(7, c_tile_7, c_tile_stride);
|
||||
} else {
|
||||
_tile_zero(4);
|
||||
_tile_zero(5);
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
|
||||
for (int32_t k = 0; k < k_times; ++k) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
|
||||
_tile_dpbf16ps(4, 0, 2);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
|
||||
_tile_dpbf16ps(5, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_dpbf16ps(6, 1, 2);
|
||||
_tile_dpbf16ps(7, 1, 3);
|
||||
|
||||
// update ptrs
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// Q buffer is prepacked
|
||||
a_tile_0 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// P buffer is not prepacked
|
||||
a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
}
|
||||
|
||||
_tile_stored(4, c_tile_4, c_tile_stride);
|
||||
_tile_stored(5, c_tile_5, c_tile_stride);
|
||||
_tile_stored(6, c_tile_6, c_tile_stride);
|
||||
_tile_stored(7, c_tile_7, c_tile_stride);
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
const int32_t m_0 = AMX_TILE_ROW_NUM;
|
||||
const int32_t m_1 = m - AMX_TILE_ROW_NUM;
|
||||
config.rows[0] = m_0;
|
||||
config.rows[1] = m_1;
|
||||
config.rows[2] = AMX_TILE_ROW_NUM;
|
||||
config.rows[3] = AMX_TILE_ROW_NUM;
|
||||
config.rows[4] = m_0;
|
||||
config.rows[5] = m_0;
|
||||
config.rows[6] = m_1;
|
||||
config.rows[7] = m_1;
|
||||
_tile_loadconfig(&config);
|
||||
}
|
||||
};
|
||||
|
||||
// 1-2-2 pattern, for 0 < m <= 16
|
||||
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
|
||||
// be m, m
|
||||
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
|
||||
// should be 16
|
||||
// TILE 6, 7: store results C matrix, row num should be m
|
||||
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
|
||||
template <typename q_buffer_t, typename kv_cache_t>
|
||||
class TileGemm122 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
|
||||
void* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122");
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
TORCH_CHECK(false, "Unsupported kv cache type for TileGemm122");
|
||||
}
|
||||
};
|
||||
|
||||
// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm122<c10::BFloat16, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
|
||||
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
|
||||
|
||||
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
|
||||
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
c10::BFloat16* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
c10::BFloat16* __restrict__ a_tile_0 = a_tile;
|
||||
c10::BFloat16* __restrict__ a_tile_1 = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// q_buffer is prepacked
|
||||
return a_tile + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// logits_buffer is row-major
|
||||
return a_tile + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
const int64_t a_tile_stride = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// q_buffer is prepacked
|
||||
return AMX_TILE_ROW_BYTES;
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// logits_buffer is row-major
|
||||
return lda * sizeof(c10::BFloat16);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
|
||||
kv_cache_t* __restrict__ b_tile_2 = b_tile;
|
||||
kv_cache_t* __restrict__ b_tile_3 = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
kv_cache_t* __restrict__ b_tile_4 =
|
||||
b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
kv_cache_t* __restrict__ b_tile_5 =
|
||||
b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
int64_t b_stride = AMX_TILE_ROW_BYTES;
|
||||
|
||||
float* __restrict__ c_tile_6 = c_tile;
|
||||
float* __restrict__ c_tile_7 = c_tile + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
int64_t c_stride = ldc * sizeof(float);
|
||||
|
||||
const int32_t k_times =
|
||||
dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
|
||||
const int32_t k_group_times = k_times / 2;
|
||||
const bool has_tail = (k_times % 2 == 1);
|
||||
|
||||
if (accum_c) {
|
||||
_tile_loadd(6, c_tile_6, c_stride);
|
||||
_tile_loadd(7, c_tile_7, c_stride);
|
||||
} else {
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_4[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_5[scratch_elems];
|
||||
for (int32_t k = 0; k < k_group_times; ++k) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
|
||||
const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
|
||||
_tile_dpbf16ps(6, 1, 4);
|
||||
_tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
|
||||
_tile_dpbf16ps(7, 1, 5);
|
||||
|
||||
// update ptrs
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// Q buffer is prepacked
|
||||
a_tile_0 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// P buffer is not prepacked
|
||||
a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
}
|
||||
b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
}
|
||||
|
||||
if (has_tail) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
}
|
||||
|
||||
_tile_stored(6, c_tile_6, c_stride);
|
||||
_tile_stored(7, c_tile_7, c_stride);
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
config.rows[0] = m;
|
||||
config.rows[1] = m;
|
||||
config.rows[2] = AMX_TILE_ROW_NUM;
|
||||
config.rows[3] = AMX_TILE_ROW_NUM;
|
||||
config.rows[4] = AMX_TILE_ROW_NUM;
|
||||
config.rows[5] = AMX_TILE_ROW_NUM;
|
||||
config.rows[6] = m;
|
||||
config.rows[7] = m;
|
||||
_tile_loadconfig(&config);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
|
||||
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = scalar_t;
|
||||
using kv_cache_t = kv_cache_scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = scalar_t;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment =
|
||||
32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
|
||||
constexpr static int64_t HeadDimAlignment =
|
||||
2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = 32;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::AMX;
|
||||
constexpr static bool scale_on_logits = true;
|
||||
|
||||
float k_scale = 1.0f;
|
||||
float v_scale = 1.0f;
|
||||
|
||||
public:
|
||||
AttentionImpl() : current_q_head_num_(0) {
|
||||
// Use all columns in AMX tiles
|
||||
vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
|
||||
}
|
||||
|
||||
~AttentionImpl() { _tile_release(); }
|
||||
|
||||
void init_from_input(const AttentionInput* input) {
|
||||
if constexpr (fp8_kv) {
|
||||
k_scale = input->k_scale_fp8;
|
||||
v_scale = input->v_scale_fp8;
|
||||
}
|
||||
}
|
||||
|
||||
float get_output_v_scale() const noexcept {
|
||||
if constexpr (fp8_kv) {
|
||||
// AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
|
||||
// Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
|
||||
// 2^112.
|
||||
constexpr float bias =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
|
||||
return v_scale * bias;
|
||||
}
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
if constexpr (fp8_kv) {
|
||||
// Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
|
||||
// shifts the exponent bias from FP8 to BF16 (127), so we multiply by
|
||||
// 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
|
||||
const float bias =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
|
||||
scale *= k_scale * bias;
|
||||
}
|
||||
if (q_head_num > AMX_TILE_ROW_NUM) {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
} else {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
}
|
||||
|
||||
// k_cache_token_group_stride: stride of K cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment * head_dim;
|
||||
}
|
||||
|
||||
// v_cache_token_group_stride: stride of V cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment * (AMX_TILE_ROW_BYTES / 4);
|
||||
}
|
||||
|
||||
// v_cache_head_group_stride: stride of V cache when move to next
|
||||
// HeadDimAlignment head dims in a block
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return block_size * HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(
|
||||
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
|
||||
scalar_t* __restrict__ q_buffer, const int32_t q_num,
|
||||
const int32_t q_heads_per_kv, const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, const float scale) {
|
||||
constexpr int64_t bytes_per_head = head_dim * sizeof(scalar_t);
|
||||
static_assert(bytes_per_head % AMX_TILE_ROW_BYTES == 0);
|
||||
constexpr int64_t head_size_block_num = bytes_per_head / AMX_TILE_ROW_BYTES;
|
||||
constexpr int64_t head_elem_num_pre_block =
|
||||
AMX_TILE_ROW_BYTES / sizeof(scalar_t);
|
||||
|
||||
int32_t idx = 0;
|
||||
int8_t* __restrict__ q_buffer_iter = reinterpret_cast<int8_t*>(q_buffer);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num;
|
||||
++q_num_idx, src += q_num_stride) {
|
||||
scalar_t* __restrict__ src_iter = src;
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv;
|
||||
++q_head_idx, src_iter += q_head_stride) {
|
||||
vec_op::unroll_loop<int32_t, head_size_block_num>(
|
||||
[&](int32_t head_size_block_idx) {
|
||||
// Use INT8Vec64 for 64 bytes block
|
||||
vec_op::INT8Vec64 vec(src_iter + head_size_block_idx *
|
||||
head_elem_num_pre_block);
|
||||
vec.save(q_buffer_iter + head_size_block_idx * AMX_TILE_BYTES);
|
||||
});
|
||||
|
||||
++idx;
|
||||
q_buffer_iter += AMX_TILE_ROW_BYTES;
|
||||
if ((idx & (AMX_TILE_ROW_NUM - 1)) == 0) {
|
||||
// head is in another amx tile
|
||||
q_buffer_iter -= AMX_TILE_ROW_NUM * AMX_TILE_ROW_BYTES;
|
||||
q_buffer_iter += head_size_block_num * AMX_TILE_BYTES;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reshape KV to AMX friendly layout
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
if constexpr (fp8_kv) {
|
||||
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
|
||||
reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
|
||||
key, value, reinterpret_cast<uint8_t*>(key_cache),
|
||||
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
|
||||
head_num, head_dim, block_size, key_token_num_stride,
|
||||
key_head_num_stride, value_token_num_stride, value_head_num_stride,
|
||||
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
|
||||
cache_head_num_stride, k_inv, v_inv);
|
||||
return;
|
||||
}
|
||||
|
||||
// For AMX 2D tiles, size of each line is 64 bytes
|
||||
constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
|
||||
// For AMX B matrix, N always is 16
|
||||
constexpr int64_t amx_b_tile_n_size = AMX_TILE_ROW_BYTES / 4;
|
||||
constexpr int64_t amx_b_tile_k_size = amx_tile_row_size / sizeof(scalar_t);
|
||||
// For now suppose block_size is divisible by amx_tile_column_num
|
||||
TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
|
||||
|
||||
scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
|
||||
scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
|
||||
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) {
|
||||
// skip
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
{
|
||||
// Write Key
|
||||
// Head elements should be packed as quand-words and stored in token
|
||||
// groups with (quadword_stride/4) tokens
|
||||
constexpr int64_t token_num_per_group = amx_tile_row_size / 4;
|
||||
static_assert(head_dim % (4 / sizeof(scalar_t)) == 0);
|
||||
constexpr int64_t quadword_num = head_dim / (4 / sizeof(scalar_t));
|
||||
const int32_t* key_start_quadword_ptr =
|
||||
reinterpret_cast<const int32_t*>(
|
||||
key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride);
|
||||
const int64_t group_idx = block_offset / token_num_per_group;
|
||||
const int64_t group_offset = block_offset % token_num_per_group;
|
||||
constexpr int64_t quadword_num_per_group =
|
||||
token_num_per_group * quadword_num;
|
||||
int32_t* key_cache_start_ptr =
|
||||
reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride) +
|
||||
group_idx * quadword_num_per_group + group_offset;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int64_t i = 0, j = 0; j < quadword_num;
|
||||
i += token_num_per_group, ++j) {
|
||||
key_cache_start_ptr[i] = key_start_quadword_ptr[j];
|
||||
}
|
||||
}
|
||||
{
|
||||
// Write Value
|
||||
// Different from Key, block_size dimension is packed rather than
|
||||
// head_size dimension block_size dimension is packed as quand-words;
|
||||
constexpr int64_t token_num_per_sub_group = 4 / sizeof(scalar_t);
|
||||
const int64_t token_num_per_group = block_size;
|
||||
constexpr int64_t head_elems_per_group = amx_b_tile_n_size;
|
||||
const int64_t group_size = token_num_per_group * head_elems_per_group;
|
||||
// For now suppose head_dim is divisible by amx_b_tile_n_size
|
||||
static_assert(head_dim % head_elems_per_group == 0);
|
||||
constexpr int64_t group_num = head_dim / head_elems_per_group;
|
||||
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
|
||||
const int64_t sub_group_offset =
|
||||
block_offset % token_num_per_sub_group;
|
||||
|
||||
const scalar_t* value_start_ptr = value +
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
vc + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
|
||||
sub_group_offset;
|
||||
|
||||
for (int64_t i = 0; i < group_num; ++i) {
|
||||
#pragma GCC unroll head_elems_per_group
|
||||
for (int64_t j = 0, k = 0; j < head_elems_per_group;
|
||||
++j, k += token_num_per_sub_group) {
|
||||
value_cache_start_ptr[k] = value_start_ptr[j];
|
||||
}
|
||||
value_start_ptr += head_elems_per_group;
|
||||
value_cache_start_ptr += group_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
alignas(64) __tilecfg amx_tile_config_;
|
||||
int32_t current_q_head_num_;
|
||||
};
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,214 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#pragma once
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <type_traits>
|
||||
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
|
||||
typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
|
||||
typedef float __attribute__((__may_alias__)) f32_alias_t;
|
||||
|
||||
// Reference scalar dequant — used to verify vectorized AMX dequant.
|
||||
inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
|
||||
// NaN encoding in E4M3
|
||||
if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
|
||||
uint32_t b_u32 = static_cast<uint32_t>(b);
|
||||
uint32_t sign = (b_u32 & 0x80) << 24;
|
||||
uint32_t payload = (b_u32 & 0x7F) << 20;
|
||||
uint32_t bits = sign | payload;
|
||||
float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
|
||||
float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
|
||||
return b_f32_scaled;
|
||||
}
|
||||
|
||||
inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
|
||||
v *= inv_scale;
|
||||
constexpr float fp8_max = 448.0f;
|
||||
v = std::max(-fp8_max, std::min(fp8_max, v));
|
||||
if (v == 0.0f) return 0;
|
||||
|
||||
// Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
|
||||
// bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
|
||||
float v_f32_unscaled = v * 0x1p-120f;
|
||||
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
|
||||
uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
|
||||
uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
|
||||
if (payload == 0) return sign;
|
||||
payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding
|
||||
return static_cast<uint8_t>(sign | payload);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AMX reshape impl — parameterised on the quantisation function.
|
||||
// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
|
||||
// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
|
||||
// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
|
||||
// block_size must be divisible by 32.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
|
||||
inline void reshape_and_cache_fp8_amx_impl(
|
||||
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
|
||||
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
|
||||
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
|
||||
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
|
||||
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
|
||||
float v_inv) {
|
||||
constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM
|
||||
const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16
|
||||
const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
|
||||
constexpr int64_t head_elems_per_group = 16;
|
||||
constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16)
|
||||
const int64_t group_num = head_dim / head_elems_per_group;
|
||||
const int64_t group_size = block_size * head_elems_per_group;
|
||||
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t tok = 0; tok < token_num; ++tok) {
|
||||
for (int64_t h = 0; h < head_num; ++h) {
|
||||
const int64_t slot = slot_ptr[tok];
|
||||
if (slot < 0) continue;
|
||||
const int64_t block_idx = slot / block_size;
|
||||
const int64_t block_offset = slot % block_size;
|
||||
|
||||
// Key: halfword-packed, 2 FP8 per uint16
|
||||
{
|
||||
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
|
||||
const int64_t group_idx = block_offset / token_num_per_group;
|
||||
const int64_t group_offset = block_offset % token_num_per_group;
|
||||
uint16_t* kdst =
|
||||
reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
|
||||
h * kc_stride1) +
|
||||
group_idx * halfword_num_per_group + group_offset;
|
||||
for (int64_t j = 0; j < halfword_num; ++j) {
|
||||
uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
|
||||
uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
|
||||
uint8_t bytes[2] = {fp8_0, fp8_1};
|
||||
uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
|
||||
kdst[j * token_num_per_group] = hw;
|
||||
}
|
||||
}
|
||||
|
||||
// Value: sub-group packing (token_num_per_sub_group = 2)
|
||||
{
|
||||
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
|
||||
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
|
||||
const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
|
||||
uint8_t* vdst =
|
||||
value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
|
||||
sub_group_idx * token_num_per_sub_group * head_elems_per_group +
|
||||
sub_group_offset;
|
||||
for (int64_t i = 0; i < group_num; ++i) {
|
||||
for (int64_t j = 0; j < head_elems_per_group; ++j)
|
||||
vdst[j * token_num_per_sub_group] =
|
||||
quant_fn(static_cast<float>(vsrc[j]), v_inv);
|
||||
vsrc += head_elems_per_group;
|
||||
vdst += group_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// FP8 E5M2 scalar helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Reference scalar dequant — used to verify vectorized AMX dequant.
|
||||
// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
|
||||
// Byte b → FP16 bits = b << 8 (no bias correction needed).
|
||||
inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
|
||||
const uint8_t exp_bits = (b >> 2) & 0x1F;
|
||||
const uint8_t mant_bits = b & 0x03;
|
||||
// NaN: exp=11111, mant!=00
|
||||
if (exp_bits == 0x1F && mant_bits != 0)
|
||||
return std::numeric_limits<float>::quiet_NaN();
|
||||
const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
|
||||
if (exp_bits == 0x1F)
|
||||
return sign ? -std::numeric_limits<float>::infinity()
|
||||
: std::numeric_limits<float>::infinity();
|
||||
if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
|
||||
if (mant_bits == 0) return 0.0f;
|
||||
float v = mant_bits * 0x1p-16f;
|
||||
return (sign ? -v : v) * scale;
|
||||
}
|
||||
// Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
|
||||
uint32_t fp32_bits = sign |
|
||||
((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
|
||||
(static_cast<uint32_t>(mant_bits) << 21);
|
||||
float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
|
||||
return val * scale;
|
||||
}
|
||||
|
||||
inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
|
||||
v *= inv_scale;
|
||||
constexpr float fp8_e5m2_max = 57344.0f;
|
||||
v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
|
||||
if (v == 0.0f) return 0;
|
||||
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
|
||||
const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
|
||||
const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
|
||||
const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
|
||||
if (exp_fp32 < -14) { // subnormal in E5M2
|
||||
const int shift = -14 - exp_fp32;
|
||||
if (shift + 21 >= 32)
|
||||
return sign; // underflow: too small for E5M2 subnormal
|
||||
const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
|
||||
return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
|
||||
}
|
||||
const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
|
||||
return sign | (exp5 << 2) | mant2;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Select the FP8 quant function at compile time based on kv_cache_t.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename kv_cache_t>
|
||||
constexpr auto select_fp8_quant_fn() {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
|
||||
return float_to_fp8e5m2_scalar;
|
||||
else
|
||||
return float_to_fp8e4m3_scalar;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// VEC reshape impl — parameterised on the quantisation function.
|
||||
// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
|
||||
// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
|
||||
// #pragma cannot appear inside variadic macro arguments.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
|
||||
inline void reshape_and_cache_fp8_vec_impl(
|
||||
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
|
||||
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
|
||||
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
|
||||
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
|
||||
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
|
||||
float v_inv) {
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t tok = 0; tok < token_num; ++tok) {
|
||||
for (int64_t h = 0; h < head_num; ++h) {
|
||||
const int64_t slot = slot_ptr[tok];
|
||||
if (slot < 0) continue;
|
||||
const int64_t block_idx = slot / block_size;
|
||||
const int64_t block_offset = slot % block_size;
|
||||
|
||||
// Key layout: column-major within block
|
||||
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
|
||||
uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
|
||||
block_offset;
|
||||
for (int64_t i = 0; i < head_dim; ++i)
|
||||
kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
|
||||
|
||||
// Value layout: row-major within block (contiguous head_dim bytes)
|
||||
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
|
||||
uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
|
||||
h * vc_stride1 + block_offset * head_dim;
|
||||
for (int64_t i = 0; i < head_dim; ++i)
|
||||
vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,402 @@
|
||||
#ifndef CPU_ATTN_NEON_HPP
|
||||
#define CPU_ATTN_NEON_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include <arm_neon.h>
|
||||
#include <type_traits>
|
||||
#ifdef ARM_BF16_SUPPORT
|
||||
#include "cpu_attn_neon_bfmmla.hpp"
|
||||
#endif
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
// These do not use vectorized class for loading / converting
|
||||
// because csrc/cpu/cpu_types_arm.hpp does not have fallback options
|
||||
// for vec_op::BF16Vec* / vec_op::BF16Vec* on Arm HW that
|
||||
// doesn't support BF16.
|
||||
// We don't use vec_op::FP32Vec* or vec_op::FP16Vec* for consistency.
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, float32x4_t& b0,
|
||||
float32x4_t& b1);
|
||||
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, float32x4_t& b0,
|
||||
float32x4_t& b1) {
|
||||
b0 = vld1q_f32(p + 0);
|
||||
b1 = vld1q_f32(p + 4);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
|
||||
float32x4_t& b0,
|
||||
float32x4_t& b1) {
|
||||
const float16_t* h = reinterpret_cast<const float16_t*>(p);
|
||||
float16x8_t v = vld1q_f16(h);
|
||||
b0 = vcvt_f32_f16(vget_low_f16(v));
|
||||
b1 = vcvt_f32_f16(vget_high_f16(v));
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
|
||||
float32x4_t& b0,
|
||||
float32x4_t& b1) {
|
||||
const uint16_t* u = reinterpret_cast<const uint16_t*>(p);
|
||||
#ifdef ARM_BF16_SUPPORT
|
||||
uint16x8_t u0 = vld1q_u16(u);
|
||||
bfloat16x8_t bf0 = vreinterpretq_bf16_u16(u0);
|
||||
b0 = vcvtq_low_f32_bf16(bf0);
|
||||
b1 = vcvtq_high_f32_bf16(bf0);
|
||||
#else
|
||||
uint16x8_t x0 = vld1q_u16(u);
|
||||
uint32x4_t lo = vshlq_n_u32(vmovl_u16(vget_low_u16(x0)), 16);
|
||||
uint32x4_t hi = vshlq_n_u32(vmovl_u16(vget_high_u16(x0)), 16);
|
||||
b0 = vreinterpretq_f32_u32(lo);
|
||||
b1 = vreinterpretq_f32_u32(hi);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Mx8, with 1 <= M <= 8 , K streamed, unroll-by-4 with ASIMD FMLAs
|
||||
// #Loads = (K // 4) * (M + 4 * sizeof(kv_cache_t) / 2)
|
||||
// #FMLAs = (K // 4) * (4 * 2 * M)
|
||||
// We have (4 * 2 * M) FMLAs for (M + 4 * sizeof(kv_cache_t) / 2) loads
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_neon_fmla_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K],
|
||||
const kv_cache_t* __restrict B, // [K x 8],
|
||||
float* __restrict C, // [M x 8],
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
// kernel supports max M of 8, as it'd spill for larger M
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
// helpers for per-M codegen
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
// A row base pointers
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
// declare 2 accumulators per row of M
|
||||
#define DECL_ACC(i) float32x4_t acc##i##_0, acc##i##_1;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
// initialize accumulators
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i##_0 = vld1q_f32(C + (i) * ldc + 0); \
|
||||
acc##i##_1 = vld1q_f32(C + (i) * ldc + 4); \
|
||||
} else { \
|
||||
acc##i##_0 = vdupq_n_f32(0.f); \
|
||||
acc##i##_1 = vdupq_n_f32(0.f); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
// K unrolled by 4
|
||||
for (; k + 3 < K; k += 4) {
|
||||
// load A[k..k+3] for each active row (M)
|
||||
#define LOAD_A4(i) \
|
||||
float32x4_t a##i##v; \
|
||||
IF_M(i) a##i##v = vld1q_f32(a##i + k);
|
||||
ROWS_APPLY(LOAD_A4)
|
||||
#undef LOAD_A4
|
||||
|
||||
// helper: FMA lane L from aiv
|
||||
#define FMAS_LANE(i, aiv, L) \
|
||||
IF_M(i) { \
|
||||
acc##i##_0 = vfmaq_laneq_f32(acc##i##_0, b0, aiv, L); \
|
||||
acc##i##_1 = vfmaq_laneq_f32(acc##i##_1, b1, aiv, L); \
|
||||
}
|
||||
|
||||
// k + 0
|
||||
{
|
||||
float32x4_t b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
|
||||
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
// k + 1
|
||||
{
|
||||
float32x4_t b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
|
||||
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
// k + 2
|
||||
{
|
||||
float32x4_t b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
|
||||
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
// k + 3
|
||||
{
|
||||
float32x4_t b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
|
||||
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
#undef FMAS_LANE
|
||||
}
|
||||
|
||||
// K tail
|
||||
for (; k < K; ++k) {
|
||||
float32x4_t b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
float32x4_t ai = vdupq_n_f32(*(a##i + k)); \
|
||||
acc##i##_0 = vfmaq_f32(acc##i##_0, b0, ai); \
|
||||
acc##i##_1 = vfmaq_f32(acc##i##_1, b1, ai); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
// store accumulators to C
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { \
|
||||
vst1q_f32(C + (i) * ldc + 0, acc##i##_0); \
|
||||
vst1q_f32(C + (i) * ldc + 4, acc##i##_1); \
|
||||
}
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_neon_fmla_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda,
|
||||
int64_t ldb, int64_t ldc,
|
||||
bool accumulate) {
|
||||
// micro kernel is Mx8
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_neon_fmla_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_neon_fmla_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_neon_fmla_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_neon_fmla_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
// no tail loop for N as it's guaranteed to be a multiple of 8
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmNeonFMLA {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_neon_fmla_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_neon_fmla_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// this is similar to "ISA::VEC" at the moment
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment =
|
||||
BLOCK_SIZE_ALIGNMENT; // KV token num unit of QK and PV phases
|
||||
constexpr static int64_t HeadDimAlignment =
|
||||
HEAD_SIZE_ALIGNMENT; // headdim num unit of PV phase
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::NEON;
|
||||
constexpr static bool scale_on_logits = false; // apply scale on q_buffer
|
||||
|
||||
static_assert(HeadDim % HeadDimAlignment == 0);
|
||||
// the gemm micro kernel is Mx8
|
||||
static_assert(HeadDimAlignment % 8 == 0);
|
||||
static_assert(BlockSizeAlignment % 8 == 0);
|
||||
|
||||
public:
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmNeonFMLA<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// k_cache_token_group_stride: stride of K cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // layout of k_cache block is [head_dim,
|
||||
// block_size], row-major
|
||||
}
|
||||
|
||||
// v_cache_token_group_stride: stride of V cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
|
||||
// head_dim], row-major
|
||||
}
|
||||
|
||||
// v_cache_head_group_stride: stride of V cache when move to next
|
||||
// HeadDimAlignment head dims in a block
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
|
||||
// row-major
|
||||
}
|
||||
|
||||
// Copy q to q_buffer and cast it to fp32
|
||||
static void copy_q_heads_tile(
|
||||
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
|
||||
float* __restrict__ q_buffer, const int32_t q_num,
|
||||
const int32_t q_heads_per_kv, const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
static_assert(head_dim % 16 == 0);
|
||||
constexpr int32_t unroll_size = head_dim / 16;
|
||||
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
vec_op::FP32Vec16 scale_vec(scale);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
|
||||
scalar_t* __restrict__ curr_q =
|
||||
src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
|
||||
float* __restrict__ curr_q_buffer =
|
||||
q_buffer + q_num_idx * q_heads_per_kv * head_dim +
|
||||
q_head_idx * head_dim;
|
||||
|
||||
vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
|
||||
load_vec_t vec(curr_q);
|
||||
vec_op::FP32Vec16 fp32_vec(vec);
|
||||
fp32_vec = fp32_vec * scale_vec;
|
||||
fp32_vec.save(curr_q_buffer);
|
||||
|
||||
curr_q += 16;
|
||||
curr_q_buffer += 16;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reshape K as column-major and V as row-major
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) {
|
||||
// skip
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
{
|
||||
// Write Key
|
||||
const scalar_t* key_start_ptr = key +
|
||||
token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_cache_start_ptr =
|
||||
key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_cache_start_ptr[j] = key_start_ptr[i];
|
||||
}
|
||||
}
|
||||
{
|
||||
// Write Value
|
||||
const scalar_t* value_start_ptr = value +
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset * head_dim;
|
||||
std::memcpy(value_cache_start_ptr, value_start_ptr,
|
||||
sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
#ifdef ARM_BF16_SUPPORT
|
||||
// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
|
||||
template <int64_t head_dim>
|
||||
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
|
||||
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
|
||||
head_dim> {};
|
||||
#endif
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif // #ifndef CPU_ATTN_ASIMD_HPP
|
||||
@@ -0,0 +1,683 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#ifndef CPU_ATTN_NEON_BFMMLA_HPP
|
||||
#define CPU_ATTN_NEON_BFMMLA_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
|
||||
#include <arm_neon.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// BFMMLA tile dimensions
|
||||
constexpr int32_t TILE_ROWS = 2; // M dimension
|
||||
constexpr int32_t TILE_K = 4; // K reduction
|
||||
constexpr int32_t TILE_COLS = 2; // N dimension (column-pair)
|
||||
|
||||
// Derived constants
|
||||
constexpr int32_t OUTPUT_COLS_PER_BLOCK = 8; // 4 column-pairs
|
||||
constexpr int32_t K_TOKENS_PER_GROUP = 8; // Tokens grouped in K cache
|
||||
constexpr int32_t V_TOKENS_PER_ROW_BLOCK = 4; // Tokens per V cache row block
|
||||
constexpr int32_t K_INNER_STRIDE = K_TOKENS_PER_GROUP * TILE_K;
|
||||
constexpr int32_t V_INNER_STRIDE = V_TOKENS_PER_ROW_BLOCK * TILE_COLS;
|
||||
constexpr int32_t PACK_ELEMENTS_PER_K_CHUNK = TILE_ROWS * TILE_K; // A packing
|
||||
|
||||
// Matrix Packing and Accumulator
|
||||
// Reshape two rows of Q into BFMMLA-friendly interleaved
|
||||
// Input: row0 = [a0,a1,a2,a3], row1 = [b0,b1,b2,b3]
|
||||
// Output: [a0,a1,a2,a3,b0,b1,b2,b3, a4,a5,a6,a7,b4,b5,b6,b7]
|
||||
// For K tail (K % TILE_K != 0): pads with zeros to complete the final chunk
|
||||
FORCE_INLINE void reshape_Q_2xK_for_bfmmla(const c10::BFloat16* __restrict r0,
|
||||
const c10::BFloat16* __restrict r1,
|
||||
c10::BFloat16* __restrict dst,
|
||||
int32_t K) {
|
||||
const uint16_t* s0 = reinterpret_cast<const uint16_t*>(r0);
|
||||
const uint16_t* s1 = reinterpret_cast<const uint16_t*>(r1);
|
||||
uint16_t* d = reinterpret_cast<uint16_t*>(dst);
|
||||
|
||||
// Process TILE_K elements at a time (PACK_ELEMENTS_PER_K_CHUNK output)
|
||||
int32_t k = 0;
|
||||
for (; k + TILE_K <= K; k += TILE_K, d += PACK_ELEMENTS_PER_K_CHUNK) {
|
||||
vst1q_u16(d, vcombine_u16(vld1_u16(s0 + k), vld1_u16(s1 + k)));
|
||||
}
|
||||
|
||||
// Handle K tail: pack remaining elements with zero-padding
|
||||
const int32_t tail = K - k;
|
||||
if (tail > 0) {
|
||||
// Pack remaining tail elements: [r0[k..k+tail-1], pad, r1[k..k+tail-1],
|
||||
// pad]
|
||||
for (int32_t t = 0; t < tail; ++t) {
|
||||
d[t] = s0[k + t];
|
||||
d[t + TILE_K] = s1[k + t];
|
||||
}
|
||||
// Zero-pad the rest
|
||||
for (int32_t t = tail; t < TILE_K; ++t) {
|
||||
d[t] = 0;
|
||||
d[t + TILE_K] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 2x2 accumulator load/store with compile-time row count
|
||||
template <int32_t m_rows>
|
||||
FORCE_INLINE float32x4_t load_acc_2x2(float* base, int64_t ldc, int col_off) {
|
||||
static_assert(m_rows == 1 || m_rows == 2);
|
||||
float32x2_t row0 = vld1_f32(base + col_off);
|
||||
float32x2_t row1 =
|
||||
(m_rows == 2) ? vld1_f32(base + ldc + col_off) : vdup_n_f32(0.f);
|
||||
return vcombine_f32(row0, row1);
|
||||
}
|
||||
|
||||
template <int32_t m_rows>
|
||||
FORCE_INLINE void store_acc_2x2(float32x4_t acc, float* base, int64_t ldc,
|
||||
int col_off) {
|
||||
static_assert(m_rows == 1 || m_rows == 2);
|
||||
vst1_f32(base + col_off, vget_low_f32(acc));
|
||||
if constexpr (m_rows == 2) {
|
||||
vst1_f32(base + ldc + col_off, vget_high_f32(acc));
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize 4 column-pair accumulators for 2 rows (8 columns total)
|
||||
#define INIT_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows, accum) \
|
||||
do { \
|
||||
if (accum) { \
|
||||
if (m_rows == 2) { \
|
||||
a0 = load_acc_2x2<2>(Crow, ldc, 0); \
|
||||
a1 = load_acc_2x2<2>(Crow, ldc, 2); \
|
||||
a2 = load_acc_2x2<2>(Crow, ldc, 4); \
|
||||
a3 = load_acc_2x2<2>(Crow, ldc, 6); \
|
||||
} else { \
|
||||
a0 = load_acc_2x2<1>(Crow, ldc, 0); \
|
||||
a1 = load_acc_2x2<1>(Crow, ldc, 2); \
|
||||
a2 = load_acc_2x2<1>(Crow, ldc, 4); \
|
||||
a3 = load_acc_2x2<1>(Crow, ldc, 6); \
|
||||
} \
|
||||
} else { \
|
||||
a0 = a1 = a2 = a3 = vdupq_n_f32(0.f); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Store 4 column-pair accumulators back to C matrix
|
||||
#define STORE_ACC_ROWPAIR_4(a0, a1, a2, a3, Crow, ldc, m_rows) \
|
||||
do { \
|
||||
if (m_rows == 2) { \
|
||||
store_acc_2x2<2>(a0, Crow, ldc, 0); \
|
||||
store_acc_2x2<2>(a1, Crow, ldc, 2); \
|
||||
store_acc_2x2<2>(a2, Crow, ldc, 4); \
|
||||
store_acc_2x2<2>(a3, Crow, ldc, 6); \
|
||||
} else { \
|
||||
store_acc_2x2<1>(a0, Crow, ldc, 0); \
|
||||
store_acc_2x2<1>(a1, Crow, ldc, 2); \
|
||||
store_acc_2x2<1>(a2, Crow, ldc, 4); \
|
||||
store_acc_2x2<1>(a3, Crow, ldc, 6); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
// Perform 4 BFMMLA operations: acc += A @ B for 4 column-pairs
|
||||
#define BFMMLA_COMPUTE_4(r0, r1, r2, r3, a, b0, b1, b2, b3) \
|
||||
do { \
|
||||
r0 = vbfmmlaq_f32(r0, a, b0); \
|
||||
r1 = vbfmmlaq_f32(r1, a, b1); \
|
||||
r2 = vbfmmlaq_f32(r2, a, b2); \
|
||||
r3 = vbfmmlaq_f32(r3, a, b3); \
|
||||
} while (0)
|
||||
|
||||
// Micro-kernel: updates a small fixed tile using BFMMLA.
|
||||
// RP = number of row-pairs (1,2,4)
|
||||
// Computes C[TILE_ROWS*RP, OUTPUT_COLS_PER_BLOCK] += A_packed @ B.
|
||||
// A_packed interleaves RP row-pairs; B layout is driven by the attention phase:
|
||||
// - AttentionGemmPhase::QK -> token-column layout (Q @ K^T)
|
||||
// - AttentionGemmPhase::PV -> token-row layout (P @ V)
|
||||
// K_static < 0 enables runtime K (PV only)
|
||||
template <int32_t RP, int32_t K_static, AttentionGemmPhase phase>
|
||||
FORCE_INLINE void gemm_rowpairs_x8_bfmmla_neon(
|
||||
const bfloat16_t* const* __restrict A_packed_rp,
|
||||
const int32_t* __restrict m_rows_rp, const bfloat16_t* __restrict B_blk,
|
||||
float* __restrict C, int64_t ldc, bool accumulate, int64_t b_stride,
|
||||
int32_t K_runtime = 0) {
|
||||
static_assert(RP == 1 || RP == 2 || RP == 4, "RP must be 1,2,4");
|
||||
static_assert(K_static < 0 || K_static % TILE_K == 0,
|
||||
"K must be divisible by TILE_K");
|
||||
static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
|
||||
"Runtime K only supported for PV");
|
||||
|
||||
constexpr bool runtime_k = (K_static < 0);
|
||||
const int32_t K_iters =
|
||||
runtime_k ? (K_runtime / TILE_K) : (K_static / TILE_K);
|
||||
const int32_t K_tail = runtime_k ? (K_runtime % TILE_K) : 0;
|
||||
|
||||
if (!runtime_k) {
|
||||
// Help the compiler fold away unused K_runtime when K is compile-time
|
||||
(void)K_runtime;
|
||||
}
|
||||
|
||||
auto* C_al = C;
|
||||
const auto* B_al = B_blk;
|
||||
|
||||
// Setup A pointers
|
||||
const bfloat16_t* a_ptr[4] = {
|
||||
A_packed_rp[0],
|
||||
(RP >= 2) ? A_packed_rp[1] : nullptr,
|
||||
(RP >= 4) ? A_packed_rp[2] : nullptr,
|
||||
(RP >= 4) ? A_packed_rp[3] : nullptr,
|
||||
};
|
||||
|
||||
// Setup B pointers based on layout
|
||||
const bfloat16_t* b_ptr[4];
|
||||
if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
b_ptr[0] = B_blk + 0 * b_stride;
|
||||
b_ptr[1] = B_blk + 1 * b_stride;
|
||||
b_ptr[2] = B_blk + 2 * b_stride;
|
||||
b_ptr[3] = B_blk + 3 * b_stride;
|
||||
}
|
||||
|
||||
float32x4_t acc[4][4];
|
||||
|
||||
// Initialize accumulators
|
||||
#define INIT_RP(rp) \
|
||||
if constexpr (RP > rp) { \
|
||||
INIT_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
|
||||
C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp], accumulate); \
|
||||
}
|
||||
INIT_RP(0);
|
||||
INIT_RP(1);
|
||||
INIT_RP(2);
|
||||
INIT_RP(3);
|
||||
#undef INIT_RP
|
||||
|
||||
// Main compute loop
|
||||
for (int32_t ki = 0; ki < K_iters; ++ki) {
|
||||
bfloat16x8_t b0, b1, b2, b3;
|
||||
if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
b0 = vld1q_bf16(b_ptr[0] + ki * V_INNER_STRIDE);
|
||||
b1 = vld1q_bf16(b_ptr[1] + ki * V_INNER_STRIDE);
|
||||
b2 = vld1q_bf16(b_ptr[2] + ki * V_INNER_STRIDE);
|
||||
b3 = vld1q_bf16(b_ptr[3] + ki * V_INNER_STRIDE);
|
||||
} else {
|
||||
const bfloat16_t* b_base = B_al + ki * b_stride;
|
||||
b0 = vld1q_bf16(b_base + 0 * V_INNER_STRIDE);
|
||||
b1 = vld1q_bf16(b_base + 1 * V_INNER_STRIDE);
|
||||
b2 = vld1q_bf16(b_base + 2 * V_INNER_STRIDE);
|
||||
b3 = vld1q_bf16(b_base + 3 * V_INNER_STRIDE);
|
||||
}
|
||||
|
||||
#define COMPUTE_RP(rp) \
|
||||
if constexpr (RP > rp) { \
|
||||
bfloat16x8_t a = vld1q_bf16(a_ptr[rp] + ki * PACK_ELEMENTS_PER_K_CHUNK); \
|
||||
BFMMLA_COMPUTE_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], a, b0, \
|
||||
b1, b2, b3); \
|
||||
}
|
||||
COMPUTE_RP(0);
|
||||
COMPUTE_RP(1);
|
||||
COMPUTE_RP(2);
|
||||
COMPUTE_RP(3);
|
||||
#undef COMPUTE_RP
|
||||
}
|
||||
|
||||
// K tail for runtime PV: fallback path
|
||||
if constexpr (runtime_k) {
|
||||
if (K_tail > 0) {
|
||||
const int32_t tail_offset = K_iters * V_INNER_STRIDE;
|
||||
const int32_t a_tail_offset = K_iters * PACK_ELEMENTS_PER_K_CHUNK;
|
||||
for (int32_t kt = 0; kt < K_tail; ++kt) {
|
||||
float32x4_t b_vecs[4];
|
||||
for (int32_t p = 0; p < 4; ++p) {
|
||||
const bfloat16_t* bp = b_ptr[p] + tail_offset + kt * TILE_COLS;
|
||||
const float b0 = vcvtah_f32_bf16(bp[0]);
|
||||
const float b1 = vcvtah_f32_bf16(bp[1]);
|
||||
const float32x2_t b_pair = vset_lane_f32(b1, vdup_n_f32(b0), 1);
|
||||
b_vecs[p] = vcombine_f32(b_pair, b_pair);
|
||||
}
|
||||
|
||||
#define TAIL_RP(rp) \
|
||||
if constexpr (RP > rp) { \
|
||||
const bfloat16_t* ap = A_packed_rp[rp] + a_tail_offset; \
|
||||
float a_row0 = vcvtah_f32_bf16(ap[kt]); \
|
||||
float a_row1 = \
|
||||
(m_rows_rp[rp] == 2) ? vcvtah_f32_bf16(ap[kt + TILE_K]) : 0.0f; \
|
||||
const float32x4_t a_vec = \
|
||||
vcombine_f32(vdup_n_f32(a_row0), vdup_n_f32(a_row1)); \
|
||||
for (int32_t p = 0; p < 4; ++p) { \
|
||||
acc[rp][p] = vmlaq_f32(acc[rp][p], a_vec, b_vecs[p]); \
|
||||
} \
|
||||
}
|
||||
TAIL_RP(0);
|
||||
TAIL_RP(1);
|
||||
TAIL_RP(2);
|
||||
TAIL_RP(3);
|
||||
#undef TAIL_RP
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Store results
|
||||
#define STORE_RP(rp) \
|
||||
if constexpr (RP > rp) { \
|
||||
STORE_ACC_ROWPAIR_4(acc[rp][0], acc[rp][1], acc[rp][2], acc[rp][3], \
|
||||
C_al + (rp * 2) * ldc, ldc, m_rows_rp[rp]); \
|
||||
}
|
||||
STORE_RP(0);
|
||||
STORE_RP(1);
|
||||
STORE_RP(2);
|
||||
STORE_RP(3);
|
||||
#undef STORE_RP
|
||||
}
|
||||
|
||||
// Meso-kernel: packs a small MBxK slice of A, then tiles over N and calls the
|
||||
// micro-kernel for each OUTPUT_COLS_PER_BLOCK chunk. K_static < 0 enables
|
||||
// runtime K (PV only).
|
||||
template <int32_t MB, int32_t N, int32_t K_static, AttentionGemmPhase phase>
|
||||
FORCE_INLINE void gemm_packA_compute_MB_xN(
|
||||
const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
|
||||
float* __restrict C, int32_t K_runtime, int64_t lda, int64_t ldc,
|
||||
int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
|
||||
static_assert(MB >= 1 && MB <= 8, "MB must be in [1,8]");
|
||||
static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
|
||||
"N must be a multiple of OUTPUT_COLS_PER_BLOCK");
|
||||
static_assert(K_static < 0 || K_static % TILE_K == 0,
|
||||
"K must be divisible by TILE_K");
|
||||
static_assert(K_static >= 0 || phase == AttentionGemmPhase::PV,
|
||||
"Runtime K only supported for PV");
|
||||
|
||||
constexpr bool runtime_k = (K_static < 0);
|
||||
const int32_t K_val = runtime_k ? K_runtime : K_static;
|
||||
|
||||
// Keep small packs on-stack to avoid heap churn
|
||||
constexpr int32_t STACK_PACK_STRIDE =
|
||||
(1024 / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
|
||||
|
||||
constexpr int32_t ROW_PAIRS = (MB + 1) / TILE_ROWS;
|
||||
const int32_t pack_stride =
|
||||
runtime_k ? ((K_val + TILE_K - 1) / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK
|
||||
: (K_static / TILE_K) * PACK_ELEMENTS_PER_K_CHUNK;
|
||||
|
||||
alignas(64) c10::BFloat16 A_packed_stack[ROW_PAIRS * STACK_PACK_STRIDE];
|
||||
std::vector<c10::BFloat16> A_packed_heap;
|
||||
c10::BFloat16* A_packed =
|
||||
(pack_stride <= STACK_PACK_STRIDE)
|
||||
? A_packed_stack
|
||||
: (A_packed_heap.resize(ROW_PAIRS * pack_stride),
|
||||
A_packed_heap.data());
|
||||
|
||||
for (int32_t rp = 0; rp < ROW_PAIRS; ++rp) {
|
||||
const int32_t m = rp * TILE_ROWS;
|
||||
const int32_t m_rows = (m + 1 < MB) ? TILE_ROWS : 1;
|
||||
const c10::BFloat16* A0 = A + m * lda;
|
||||
const c10::BFloat16* A1 = (m_rows == TILE_ROWS) ? (A + (m + 1) * lda) : A0;
|
||||
reshape_Q_2xK_for_bfmmla(A0, A1, A_packed + rp * pack_stride, K_val);
|
||||
}
|
||||
|
||||
for (int32_t n = 0; n < N; n += OUTPUT_COLS_PER_BLOCK) {
|
||||
const c10::BFloat16* B_blk_c10 =
|
||||
(phase == AttentionGemmPhase::PV)
|
||||
? (B + (n / TILE_COLS) * b_layout_stride)
|
||||
: (B + (n / OUTPUT_COLS_PER_BLOCK) * b_layout_stride);
|
||||
const bfloat16_t* B_blk = reinterpret_cast<const bfloat16_t*>(B_blk_c10);
|
||||
|
||||
// Process row-pairs in groups of 4, 2, then 1
|
||||
int32_t row_pair_idx = 0;
|
||||
|
||||
#define PROCESS_RP_GROUP(group_size) \
|
||||
for (; row_pair_idx + (group_size - 1) < ROW_PAIRS; \
|
||||
row_pair_idx += group_size) { \
|
||||
const bfloat16_t* Ap[group_size]; \
|
||||
int32_t mr[group_size]; \
|
||||
for (int32_t i = 0; i < group_size; ++i) { \
|
||||
Ap[i] = reinterpret_cast<const bfloat16_t*>( \
|
||||
A_packed + (row_pair_idx + i) * pack_stride); \
|
||||
mr[i] = (((row_pair_idx + i) * TILE_ROWS + 1) < MB) ? TILE_ROWS : 1; \
|
||||
} \
|
||||
float* C_blk = C + (row_pair_idx * TILE_ROWS) * ldc + n; \
|
||||
if constexpr (runtime_k) { \
|
||||
gemm_rowpairs_x8_bfmmla_neon<group_size, -1, phase>( \
|
||||
Ap, mr, B_blk, C_blk, ldc, accumulate, b_layout_stride, K_val); \
|
||||
} else { \
|
||||
gemm_rowpairs_x8_bfmmla_neon<group_size, K_static, phase>( \
|
||||
Ap, mr, B_blk, C_blk, ldc, accumulate, \
|
||||
(phase == AttentionGemmPhase::PV) ? b_layout_stride \
|
||||
: b_reduction_stride); \
|
||||
} \
|
||||
}
|
||||
|
||||
PROCESS_RP_GROUP(4);
|
||||
PROCESS_RP_GROUP(2);
|
||||
PROCESS_RP_GROUP(1);
|
||||
#undef PROCESS_RP_GROUP
|
||||
}
|
||||
}
|
||||
|
||||
// Macro-kernel: iterates over M in MB={8,4,2,1} chunks.
|
||||
// Supports compile-time K specialization when K >= 0; otherwise uses runtime K
|
||||
// (runtime K path is only supported for PV).
|
||||
template <AttentionGemmPhase phase, int32_t N, int32_t K = -1>
|
||||
FORCE_INLINE void gemm_macro_neon_bfmmla(
|
||||
const c10::BFloat16* __restrict A, const c10::BFloat16* __restrict B,
|
||||
float* __restrict C, int32_t M, int32_t K_runtime, int64_t lda, int64_t ldc,
|
||||
int64_t b_layout_stride, int64_t b_reduction_stride, bool accumulate) {
|
||||
static_assert(N % OUTPUT_COLS_PER_BLOCK == 0,
|
||||
"N must be a multiple of OUTPUT_COLS_PER_BLOCK");
|
||||
|
||||
if constexpr (K >= 0) {
|
||||
static_assert(K % TILE_K == 0, "K must be divisible by TILE_K");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
const int32_t rem = M - m;
|
||||
const c10::BFloat16* A_blk = A + m * lda;
|
||||
float* C_blk = C + m * ldc;
|
||||
|
||||
#define DISPATCH_MB(mb) \
|
||||
gemm_packA_compute_MB_xN<mb, N, K, phase>(A_blk, B, C_blk, 0, lda, ldc, \
|
||||
b_layout_stride, \
|
||||
b_reduction_stride, accumulate)
|
||||
|
||||
if (rem >= 8) {
|
||||
DISPATCH_MB(8);
|
||||
m += 8;
|
||||
} else if (rem >= 4) {
|
||||
DISPATCH_MB(4);
|
||||
m += 4;
|
||||
} else if (rem >= 2) {
|
||||
DISPATCH_MB(2);
|
||||
m += 2;
|
||||
} else {
|
||||
DISPATCH_MB(1);
|
||||
m += 1;
|
||||
}
|
||||
#undef DISPATCH_MB
|
||||
}
|
||||
} else {
|
||||
static_assert(phase == AttentionGemmPhase::PV,
|
||||
"Runtime K specialization only supported for PV.");
|
||||
const int32_t K_val = K_runtime;
|
||||
|
||||
for (int32_t m = 0; m < M;) {
|
||||
const int32_t rem = M - m;
|
||||
const c10::BFloat16* A_blk = A + m * lda;
|
||||
float* C_blk = C + m * ldc;
|
||||
|
||||
#define DISPATCH_MB_RUNTIME(mb) \
|
||||
gemm_packA_compute_MB_xN<mb, N, -1, phase>(A_blk, B, C_blk, K_val, lda, ldc, \
|
||||
b_layout_stride, \
|
||||
b_reduction_stride, accumulate)
|
||||
|
||||
if (rem >= 8) {
|
||||
DISPATCH_MB_RUNTIME(8);
|
||||
m += 8;
|
||||
} else if (rem >= 4) {
|
||||
DISPATCH_MB_RUNTIME(4);
|
||||
m += 4;
|
||||
} else if (rem >= 2) {
|
||||
DISPATCH_MB_RUNTIME(2);
|
||||
m += 2;
|
||||
} else {
|
||||
DISPATCH_MB_RUNTIME(1);
|
||||
m += 1;
|
||||
}
|
||||
#undef DISPATCH_MB_RUNTIME
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#undef INIT_ACC_ROWPAIR_4
|
||||
#undef STORE_ACC_ROWPAIR_4
|
||||
#undef BFMMLA_COMPUTE_4
|
||||
|
||||
} // namespace
|
||||
|
||||
// TileGemm Adapter for Attention
|
||||
|
||||
template <typename kv_cache_t, int32_t BlockTokens, int32_t HeadDim>
|
||||
class TileGemmNEONBFMMLA {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t head_dim_ct>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size, void* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
[[maybe_unused]] const int64_t ldb,
|
||||
const int64_t ldc,
|
||||
[[maybe_unused]] const int32_t block_size,
|
||||
[[maybe_unused]] const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
static_assert(BlockTokens % OUTPUT_COLS_PER_BLOCK == 0);
|
||||
// BFMMLA kernels require compile-time head_dim; keep head_dim_ct only for
|
||||
// API parity with other tile_gemm implementations.
|
||||
if constexpr (head_dim_ct >= 0) {
|
||||
static_assert(head_dim_ct == HeadDim,
|
||||
"BFMMLA expects head_dim_ct to match HeadDim; PV passes "
|
||||
"-1 for API parity.");
|
||||
}
|
||||
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
const int64_t b_reduction_stride = K_INNER_STRIDE;
|
||||
const int64_t b_token_block_stride = (HeadDim / TILE_K) * K_INNER_STRIDE;
|
||||
|
||||
gemm_macro_neon_bfmmla<AttentionGemmPhase::QK, BlockTokens, HeadDim>(
|
||||
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
|
||||
m_size, 0, lda, ldc, b_token_block_stride, b_reduction_stride,
|
||||
accum_c);
|
||||
} else {
|
||||
const int64_t b_pair_stride =
|
||||
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
|
||||
|
||||
// PV gemm with runtime K specialization
|
||||
switch (dynamic_k_size) {
|
||||
case 32:
|
||||
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 32>(
|
||||
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
|
||||
m_size, 32, lda, ldc, b_pair_stride, 0, accum_c);
|
||||
break;
|
||||
case 128:
|
||||
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 128>(
|
||||
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
|
||||
m_size, 128, lda, ldc, b_pair_stride, 0, accum_c);
|
||||
break;
|
||||
case 256:
|
||||
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim, 256>(
|
||||
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
|
||||
m_size, 256, lda, ldc, b_pair_stride, 0, accum_c);
|
||||
break;
|
||||
default:
|
||||
gemm_macro_neon_bfmmla<AttentionGemmPhase::PV, HeadDim>(
|
||||
reinterpret_cast<const c10::BFloat16*>(a_tile), b_tile, c_tile,
|
||||
m_size, dynamic_k_size, lda, ldc, b_pair_stride, 0, accum_c);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Shared ASIMD BFMMLA implementation (BF16 only). The block size alignment and
|
||||
// ISA tag are template parameters so we can reuse the same kernels for
|
||||
// different NEON configurations.
|
||||
template <int64_t block_size_alignment, ISA isa_type, int64_t head_dim>
|
||||
class AttentionImplNEONBFMMLA {
|
||||
public:
|
||||
using query_t = c10::BFloat16;
|
||||
using q_buffer_t = c10::BFloat16;
|
||||
using kv_cache_t = c10::BFloat16;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = c10::BFloat16;
|
||||
|
||||
static constexpr int64_t BlockSizeAlignment = block_size_alignment;
|
||||
// HeadDimAlignment equals head_dim so that the PV phase processes
|
||||
// the full head dimension in a single gemm call.
|
||||
static constexpr int64_t HeadDimAlignment = head_dim;
|
||||
static constexpr int64_t MaxQHeadNumPerIteration = 16;
|
||||
static constexpr int64_t HeadDim = head_dim;
|
||||
static constexpr ISA ISAType = isa_type;
|
||||
static constexpr bool scale_on_logits = false;
|
||||
|
||||
static_assert(HeadDim % OUTPUT_COLS_PER_BLOCK == 0);
|
||||
static_assert(BlockSizeAlignment % OUTPUT_COLS_PER_BLOCK == 0);
|
||||
static_assert(HeadDim % TILE_K == 0, "HeadDim must be a multiple of TILE_K");
|
||||
|
||||
public:
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<
|
||||
TileGemmNEONBFMMLA<kv_cache_t, static_cast<int32_t>(BlockSizeAlignment),
|
||||
static_cast<int32_t>(HeadDim)>>
|
||||
attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// Key cache stride per token group (TokenColumn layout; QK)
|
||||
static constexpr int64_t k_cache_token_group_stride(
|
||||
[[maybe_unused]] const int32_t block_size) {
|
||||
static_assert(BlockSizeAlignment % K_TOKENS_PER_GROUP == 0);
|
||||
return (BlockSizeAlignment / K_TOKENS_PER_GROUP) *
|
||||
((head_dim / TILE_K) * K_INNER_STRIDE);
|
||||
}
|
||||
|
||||
// Value cache stride per token group (TokenRow layout; PV)
|
||||
static constexpr int64_t v_cache_token_group_stride(
|
||||
[[maybe_unused]] const int32_t block_size) {
|
||||
static_assert(BlockSizeAlignment % V_TOKENS_PER_ROW_BLOCK == 0);
|
||||
return (BlockSizeAlignment / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
|
||||
}
|
||||
|
||||
// The stride to move to the "next" head_dim group
|
||||
// is the full V cache size per head, since HeadDimAlignment == head_dim.
|
||||
// Hence, the stride is not used in this case
|
||||
static constexpr int64_t v_cache_head_group_stride(
|
||||
[[maybe_unused]] const int32_t block_size) {
|
||||
return head_dim * block_size;
|
||||
}
|
||||
|
||||
// Convert Q heads to BF16 and apply scale factor using native BF16 intrinsics
|
||||
static void copy_q_heads_tile(c10::BFloat16* __restrict__ src,
|
||||
c10::BFloat16* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
constexpr int32_t dim = static_cast<int32_t>(head_dim);
|
||||
const float32x4_t scale_vec = vdupq_n_f32(scale);
|
||||
|
||||
for (int32_t qi = 0; qi < q_num; ++qi) {
|
||||
for (int32_t hi = 0; hi < q_heads_per_kv; ++hi) {
|
||||
c10::BFloat16* __restrict__ curr_q =
|
||||
src + qi * q_num_stride + hi * q_head_stride;
|
||||
c10::BFloat16* __restrict__ dst =
|
||||
q_buffer + qi * q_heads_per_kv * head_dim + hi * head_dim;
|
||||
|
||||
for (int32_t i = 0; i < dim; i += OUTPUT_COLS_PER_BLOCK) {
|
||||
bfloat16x8_t in8 =
|
||||
vld1q_bf16(reinterpret_cast<const bfloat16_t*>(curr_q + i));
|
||||
float32x4_t lo = vmulq_f32(vcvtq_low_f32_bf16(in8), scale_vec);
|
||||
float32x4_t hi = vmulq_f32(vcvtq_high_f32_bf16(in8), scale_vec);
|
||||
|
||||
bfloat16x4_t lo_b = vcvt_bf16_f32(lo);
|
||||
bfloat16x4_t hi_b = vcvt_bf16_f32(hi);
|
||||
bfloat16x8_t out = vcombine_bf16(lo_b, hi_b);
|
||||
vst1q_bf16(reinterpret_cast<bfloat16_t*>(dst + i), out);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
// Reshape and cache K/V into BFMMLA-optimized layouts
|
||||
// K cache:
|
||||
// [block_size/K_TOKENS_PER_GROUP][head_dim/TILE_K][K_INNER_STRIDE]
|
||||
// - TokenColumn
|
||||
// V cache:
|
||||
// [head_dim/TILE_COLS][block_size/V_TOKENS_PER_ROW_BLOCK][V_INNER_STRIDE]
|
||||
// - TokenRows
|
||||
static void reshape_and_cache(
|
||||
const c10::BFloat16* __restrict__ key,
|
||||
const c10::BFloat16* __restrict__ value,
|
||||
c10::BFloat16* __restrict__ key_cache,
|
||||
c10::BFloat16* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride,
|
||||
[[maybe_unused]] const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size,
|
||||
[[maybe_unused]] const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
|
||||
const int64_t v_pair_stride =
|
||||
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) continue;
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
|
||||
// Key cache: TokenColumn QK
|
||||
{
|
||||
const c10::BFloat16* __restrict key_src =
|
||||
key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
|
||||
c10::BFloat16* __restrict key_base = key_cache +
|
||||
block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride;
|
||||
|
||||
const int64_t block_in_block = block_offset / K_TOKENS_PER_GROUP;
|
||||
const int64_t pair_in_block =
|
||||
(block_offset % K_TOKENS_PER_GROUP) / TILE_COLS;
|
||||
const int64_t lane_base = (block_offset & 1) ? TILE_K : 0;
|
||||
|
||||
c10::BFloat16* __restrict block_base =
|
||||
key_base + block_in_block * k_block_stride;
|
||||
|
||||
for (int64_t hd4 = 0; hd4 < head_dim / TILE_K; ++hd4) {
|
||||
uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(
|
||||
block_base + hd4 * K_INNER_STRIDE +
|
||||
pair_in_block * V_INNER_STRIDE + lane_base);
|
||||
const uint16_t* src_u16 =
|
||||
reinterpret_cast<const uint16_t*>(key_src + hd4 * TILE_K);
|
||||
vst1_u16(dst_u16, vld1_u16(src_u16));
|
||||
}
|
||||
}
|
||||
|
||||
// Value cache: TokenRow PV
|
||||
{
|
||||
const c10::BFloat16* __restrict value_src =
|
||||
value + token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
|
||||
c10::BFloat16* __restrict value_base =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride;
|
||||
|
||||
const int64_t row_block = block_offset / V_TOKENS_PER_ROW_BLOCK;
|
||||
const int64_t lane = block_offset & (V_TOKENS_PER_ROW_BLOCK - 1);
|
||||
|
||||
c10::BFloat16* __restrict row_block_base =
|
||||
value_base + row_block * V_INNER_STRIDE;
|
||||
|
||||
for (int64_t hd2 = 0; hd2 < head_dim / TILE_COLS; ++hd2) {
|
||||
c10::BFloat16* __restrict dst_val =
|
||||
row_block_base + hd2 * v_pair_stride;
|
||||
|
||||
const uint16_t* src_u16 =
|
||||
reinterpret_cast<const uint16_t*>(value_src);
|
||||
uint16_t* dst_u16 = reinterpret_cast<uint16_t*>(dst_val);
|
||||
dst_u16[lane] = src_u16[hd2 * TILE_COLS + 0];
|
||||
dst_u16[lane + V_TOKENS_PER_ROW_BLOCK] =
|
||||
src_u16[hd2 * TILE_COLS + 1];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif // CPU_ATTN_ASIMD_BFMMLA_HPP
|
||||
@@ -0,0 +1,412 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_ATTN_RVV_HPP
|
||||
#define CPU_ATTN_RVV_HPP
|
||||
|
||||
// RVV attention kernel using VLEN-agnostic RVVI() macros from
|
||||
// cpu_types_riscv_defs.hpp. The Mx8 tile GEMM uses 8 FP32 elements
|
||||
// per vector (LMUL_256 bits of FP32 data), which maps to:
|
||||
// VLEN=128: m2 (256 bits = 8 x FP32)
|
||||
// VLEN=256: m1 (256 bits = 8 x FP32)
|
||||
// Only VLEN=128 and VLEN=256 are supported; other VLENs (512, 1024)
|
||||
// and scalar RISC-V builds fall back to VEC/VEC16.
|
||||
#if defined(__riscv_v_min_vlen) && \
|
||||
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include "cpu_types_riscv_defs.hpp"
|
||||
#include <riscv_vector.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
// ============================================================================
|
||||
// B-matrix row loading: load 8 elements as FP32
|
||||
// ============================================================================
|
||||
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32(const kv_cache_t* p);
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<float>(const float* p) {
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(p, 8);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_B_as_f32<c10::Half>(const c10::Half* p) {
|
||||
#ifdef __riscv_zvfh
|
||||
fixed_fp16x8_t h = RVVI(__riscv_vle16_v_f16, LMUL_128)(
|
||||
reinterpret_cast<const _Float16*>(p), 8);
|
||||
return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(h, 8);
|
||||
#else
|
||||
alignas(16) float tmp[8];
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
tmp[i] = static_cast<float>(p[i]);
|
||||
}
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(tmp, 8);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p) {
|
||||
#ifdef __riscv_zvfbfmin
|
||||
fixed_bf16x8_t bf = RVVI(__riscv_vle16_v_bf16, LMUL_128)(
|
||||
reinterpret_cast<const __bf16*>(p), 8);
|
||||
return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(bf, 8);
|
||||
#else
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(p), 8);
|
||||
fixed_u32x8_t wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, 8);
|
||||
fixed_u32x8_t shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, 8);
|
||||
return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Micro kernel: Mx8 tile, K unrolled by 4, RVV scalar-broadcast FMA
|
||||
// ============================================================================
|
||||
//
|
||||
// RVV has no lane-indexed FMA; instead we load A elements as scalars and
|
||||
// use vfmacc_vf (scalar * vector + accumulator).
|
||||
//
|
||||
// The 8-column tile uses LMUL_256 bits of FP32 data:
|
||||
// VLEN=128: m2 (2 regs per accumulator), M=8 => 18 of 32 regs
|
||||
// VLEN=256: m1 (1 reg per accumulator), M=8 => 9 of 32 regs
|
||||
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_fma_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K]
|
||||
const kv_cache_t* __restrict B, // [K x 8]
|
||||
float* __restrict C, // [M x 8]
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
constexpr size_t vl = 8;
|
||||
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
#define DECL_ACC(i) fixed_fp32x8_t acc##i;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(C + (i) * ldc, vl); \
|
||||
} else { \
|
||||
acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.f, vl); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
for (; k + 3 < K; k += 4) {
|
||||
{
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb);
|
||||
#define STEP_K0(i) \
|
||||
IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 0), \
|
||||
b, vl); \
|
||||
}
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
{
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb);
|
||||
#define STEP_K1(i) \
|
||||
IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 1), \
|
||||
b, vl); \
|
||||
}
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
{
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb);
|
||||
#define STEP_K2(i) \
|
||||
IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 2), \
|
||||
b, vl); \
|
||||
}
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
{
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb);
|
||||
#define STEP_K3(i) \
|
||||
IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k + 3), \
|
||||
b, vl); \
|
||||
}
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
}
|
||||
|
||||
for (; k < K; ++k) {
|
||||
fixed_fp32x8_t b = load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
acc##i = \
|
||||
RVVI(__riscv_vfmacc_vf_f32, LMUL_256)(acc##i, *(a##i + k), b, vl); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { RVVI(__riscv_vse32_v_f32, LMUL_256)(C + (i) * ldc, acc##i, vl); }
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Macro kernel: dispatch M tiles of {8,4,2,1}, step N by 8
|
||||
// ============================================================================
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_rvv_fma_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda,
|
||||
int64_t ldb, int64_t ldc,
|
||||
bool accumulate) {
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_rvv_fma_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_rvv_fma_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_rvv_fma_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_rvv_fma_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TileGemm wrapper — plugs into AttentionMainLoop
|
||||
// ============================================================================
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmRVV {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_rvv_fma_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_rvv_fma_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// ============================================================================
|
||||
// AttentionImpl<ISA::RVV> — mirrors ISA::NEON specialization
|
||||
// ============================================================================
|
||||
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::RVV, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::RVV;
|
||||
constexpr static bool scale_on_logits = false;
|
||||
|
||||
static_assert(HeadDim % HeadDimAlignment == 0);
|
||||
static_assert(HeadDimAlignment % 8 == 0);
|
||||
static_assert(BlockSizeAlignment % 8 == 0);
|
||||
|
||||
public:
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmRVV<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
static_assert(head_dim % 16 == 0);
|
||||
constexpr int32_t unroll_size = head_dim / 16;
|
||||
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
vec_op::FP32Vec16 scale_vec(scale);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
|
||||
scalar_t* __restrict__ curr_q =
|
||||
src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
|
||||
float* __restrict__ curr_q_buffer =
|
||||
q_buffer + q_num_idx * q_heads_per_kv * head_dim +
|
||||
q_head_idx * head_dim;
|
||||
|
||||
vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
|
||||
load_vec_t vec(curr_q);
|
||||
vec_op::FP32Vec16 fp32_vec(vec);
|
||||
fp32_vec = fp32_vec * scale_vec;
|
||||
fp32_vec.save(curr_q_buffer);
|
||||
|
||||
curr_q += 16;
|
||||
curr_q_buffer += 16;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
{
|
||||
const scalar_t* key_start_ptr = key +
|
||||
token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_cache_start_ptr =
|
||||
key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
{
|
||||
const ptrdiff_t byte_stride = block_size * sizeof(scalar_t);
|
||||
int64_t i = 0;
|
||||
for (; i < head_dim;) {
|
||||
size_t vl;
|
||||
if constexpr (std::is_same_v<scalar_t, float>) {
|
||||
vl = __riscv_vsetvl_e32m2(head_dim - i);
|
||||
vfloat32m2_t v = __riscv_vle32_v_f32m2(
|
||||
reinterpret_cast<const float*>(key_start_ptr + i), vl);
|
||||
__riscv_vsse32_v_f32m2(
|
||||
reinterpret_cast<float*>(key_cache_start_ptr +
|
||||
i * block_size),
|
||||
byte_stride, v, vl);
|
||||
} else {
|
||||
vl = __riscv_vsetvl_e16m1(head_dim - i);
|
||||
vuint16m1_t v = __riscv_vle16_v_u16m1(
|
||||
reinterpret_cast<const uint16_t*>(key_start_ptr + i), vl);
|
||||
__riscv_vsse16_v_u16m1(
|
||||
reinterpret_cast<uint16_t*>(key_cache_start_ptr +
|
||||
i * block_size),
|
||||
byte_stride, v, vl);
|
||||
}
|
||||
i += vl;
|
||||
}
|
||||
}
|
||||
}
|
||||
{
|
||||
const scalar_t* value_start_ptr = value +
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset * head_dim;
|
||||
std::memcpy(value_cache_start_ptr, value_start_ptr,
|
||||
sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif // __riscv_v_min_vlen == 128 || 256
|
||||
|
||||
#endif // CPU_ATTN_RVV_HPP
|
||||
@@ -0,0 +1,325 @@
|
||||
#ifndef CPU_ATTN_VEC_HPP
|
||||
#define CPU_ATTN_VEC_HPP
|
||||
|
||||
#include "cpu_attn_fp8.hpp"
|
||||
#include "cpu_attn_impl.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
|
||||
// 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.
|
||||
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 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).
|
||||
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
|
||||
vec_op::fp8_e4m3_tag{});
|
||||
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
|
||||
} else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
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 {
|
||||
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
|
||||
return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
|
||||
vec_op::FP32Vec16(load_vec_t(ptr + 16)));
|
||||
}
|
||||
}
|
||||
|
||||
// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm82 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
switch (m_size) {
|
||||
case 1:
|
||||
gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 3:
|
||||
case 4:
|
||||
gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 5:
|
||||
case 6:
|
||||
gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 7:
|
||||
case 8:
|
||||
gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t M>
|
||||
static void gemm_micro(float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size, const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
|
||||
float* __restrict__ curr_c_0 = c_tile;
|
||||
float* __restrict__ curr_c_1 = c_tile + 16;
|
||||
|
||||
vec_op::FP32Vec16 c_regs[M * 2];
|
||||
if (accum_c) {
|
||||
float* __restrict__ curr_m_c_0 = curr_c_0;
|
||||
float* __restrict__ curr_m_c_1 = curr_c_1;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
|
||||
c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
|
||||
|
||||
// update
|
||||
curr_m_c_0 += ldc;
|
||||
curr_m_c_1 += ldc;
|
||||
});
|
||||
}
|
||||
|
||||
float* __restrict__ curr_a = a_tile;
|
||||
kv_cache_t* __restrict__ curr_b = b_tile;
|
||||
|
||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
vec_op::FP32Vec16 a_reg(*curr_m_a);
|
||||
c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
|
||||
c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
|
||||
|
||||
// update
|
||||
curr_m_a += lda;
|
||||
});
|
||||
|
||||
// update
|
||||
curr_a += 1;
|
||||
curr_b += ldb;
|
||||
}
|
||||
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i * 2].save(curr_c_0);
|
||||
c_regs[i * 2 + 1].save(curr_c_1);
|
||||
|
||||
// update
|
||||
curr_c_0 += ldc;
|
||||
curr_c_1 += ldc;
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// This is a general but naive implementation based on vector instructions
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
|
||||
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = kv_cache_scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment =
|
||||
32; // KV token num unit of QK and PV phases
|
||||
constexpr static int64_t HeadDimAlignment =
|
||||
32; // headdim num unit of PV phase
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = 8;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VEC;
|
||||
constexpr static bool scale_on_logits = fp8_kv;
|
||||
|
||||
float k_scale = 1.0f;
|
||||
float v_scale = 1.0f;
|
||||
|
||||
public:
|
||||
void init_from_input(const AttentionInput* input) {
|
||||
if constexpr (fp8_kv) {
|
||||
k_scale = input->k_scale_fp8;
|
||||
v_scale = input->v_scale_fp8;
|
||||
}
|
||||
}
|
||||
|
||||
float get_output_v_scale() const noexcept {
|
||||
if constexpr (fp8_kv) {
|
||||
// VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
|
||||
// E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
|
||||
// so no correction.
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
return v_scale;
|
||||
} else {
|
||||
return v_scale * 0x1p8f;
|
||||
}
|
||||
}
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
if constexpr (fp8_kv) {
|
||||
// Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
|
||||
// uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
|
||||
// correction.
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
scale *= k_scale;
|
||||
} else {
|
||||
scale *= k_scale * 0x1p8f;
|
||||
}
|
||||
}
|
||||
attention<TileGemm82<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// k_cache_token_group_stride: stride of K cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // layout of k_cache block is [head_dim,
|
||||
// block_size], row-major
|
||||
}
|
||||
|
||||
// v_cache_token_group_stride: stride of V cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
|
||||
// head_dim], row-major
|
||||
}
|
||||
|
||||
// v_cache_head_group_stride: stride of V cache when move to next
|
||||
// HeadDimAlignment head dims in a block
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
|
||||
// row-major
|
||||
}
|
||||
|
||||
// Copy q to q_buffer and cast it to fp32.
|
||||
// FP8: QK scale is folded into execute_attention; copy Q unscaled here.
|
||||
void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer, const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
static_assert(head_dim % 16 == 0);
|
||||
constexpr int32_t unroll_size = head_dim / 16;
|
||||
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
const float effective_scale = fp8_kv ? 1.0f : scale;
|
||||
vec_op::FP32Vec16 scale_vec(effective_scale);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
|
||||
scalar_t* __restrict__ curr_q =
|
||||
src + q_num_idx * q_num_stride + q_head_idx * q_head_stride;
|
||||
float* __restrict__ curr_q_buffer =
|
||||
q_buffer + q_num_idx * q_heads_per_kv * head_dim +
|
||||
q_head_idx * head_dim;
|
||||
|
||||
vec_op::unroll_loop<int32_t, unroll_size>([&](int32_t i) {
|
||||
load_vec_t vec(curr_q);
|
||||
vec_op::FP32Vec16 fp32_vec(vec);
|
||||
fp32_vec = fp32_vec * scale_vec;
|
||||
fp32_vec.save(curr_q_buffer);
|
||||
|
||||
curr_q += 16;
|
||||
curr_q_buffer += 16;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reshape K as column-major and V as row-major
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
if constexpr (fp8_kv) {
|
||||
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
|
||||
reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
|
||||
key, value, reinterpret_cast<uint8_t*>(key_cache),
|
||||
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
|
||||
head_num, head_dim, block_size, key_token_num_stride,
|
||||
key_head_num_stride, value_token_num_stride, value_head_num_stride,
|
||||
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
|
||||
cache_head_num_stride, k_inv, v_inv);
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) {
|
||||
// skip
|
||||
continue;
|
||||
}
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
{
|
||||
// Write Key as column-major
|
||||
const scalar_t* key_start_ptr = key +
|
||||
token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_cache_start_ptr =
|
||||
reinterpret_cast<scalar_t*>(key_cache) +
|
||||
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
|
||||
block_offset;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_cache_start_ptr[j] = key_start_ptr[i];
|
||||
}
|
||||
}
|
||||
{
|
||||
// Write Value as row-major
|
||||
const scalar_t* value_start_ptr = value +
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
reinterpret_cast<scalar_t*>(value_cache) +
|
||||
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
std::memcpy(value_cache_start_ptr, value_start_ptr,
|
||||
sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,171 @@
|
||||
#ifndef CPU_ATTN_VEC16_HPP
|
||||
#define CPU_ATTN_VEC16_HPP
|
||||
|
||||
#include "cpu_attn_vec.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
// 16-1-16 pattern, 16 regs for A, 1 regs for B, 16 regs for C, [16, K] @ [k,
|
||||
// 16]
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm161 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
switch (m_size) {
|
||||
case 1:
|
||||
gemm_micro<1>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro<2>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 3:
|
||||
case 4:
|
||||
gemm_micro<4>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 5:
|
||||
case 6:
|
||||
gemm_micro<6>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 7:
|
||||
case 8:
|
||||
gemm_micro<8>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 9:
|
||||
case 10:
|
||||
case 11:
|
||||
case 12:
|
||||
gemm_micro<12>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
case 13:
|
||||
case 14:
|
||||
case 15:
|
||||
case 16:
|
||||
gemm_micro<16>(a_tile, b_tile, c_tile, lda, ldb, ldc, block_size,
|
||||
dynamic_k_size, accum_c);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t M>
|
||||
static void gemm_micro(float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size, const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 16);
|
||||
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
|
||||
|
||||
kv_cache_t* __restrict__ curr_b_0 = b_tile;
|
||||
float* __restrict__ curr_c_0 = c_tile;
|
||||
|
||||
vec_op::FP32Vec16 c_regs[M];
|
||||
if (accum_c) {
|
||||
float* __restrict__ curr_m_c_0 = curr_c_0;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i] = vec_op::FP32Vec16(curr_m_c_0);
|
||||
|
||||
// update
|
||||
curr_m_c_0 += ldc;
|
||||
});
|
||||
}
|
||||
|
||||
float* __restrict__ curr_a = a_tile;
|
||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||
load_vec_t b_0_reg(curr_b_0);
|
||||
vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
float v = *curr_m_a;
|
||||
vec_op::FP32Vec16 a_reg(v);
|
||||
c_regs[i] = c_regs[i] + a_reg * fp32_b_0_reg;
|
||||
|
||||
// update
|
||||
curr_m_a += lda;
|
||||
});
|
||||
|
||||
// update
|
||||
curr_a += 1;
|
||||
curr_b_0 += ldb;
|
||||
}
|
||||
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i].save(curr_c_0);
|
||||
|
||||
// update
|
||||
curr_c_0 += ldc;
|
||||
});
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
// This is a general but naive implementation based on vector instructions
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
|
||||
: public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment =
|
||||
16; // KV token num unit of QK and PV phases
|
||||
constexpr static int64_t HeadDimAlignment =
|
||||
16; // headdim num unit of PV phase
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = 16;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VEC16;
|
||||
constexpr static bool scale_on_logits = false; // apply scale on q_buffer
|
||||
|
||||
public:
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemm161<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// k_cache_token_group_stride: stride of K cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // layout of k_cache block is [head_dim,
|
||||
// block_size], row-major
|
||||
}
|
||||
|
||||
// v_cache_token_group_stride: stride of V cache when move to next
|
||||
// BlockSizeAlignment tokens in a block
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment; // layout of v_cache is [block_size,
|
||||
// head_dim], row-major
|
||||
}
|
||||
|
||||
// v_cache_head_group_stride: stride of V cache when move to next
|
||||
// HeadDimAlignment head dims in a block
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment; // layout of v_cache is [block_size, head_dim],
|
||||
// row-major
|
||||
}
|
||||
};
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,366 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#ifndef CPU_ATTN_VSX_HPP
|
||||
#define CPU_ATTN_VSX_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include <altivec.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// ppc64le Vector = 16 bytes (128 bits)
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
|
||||
__vector float& b1);
|
||||
|
||||
// [1] Float Specialization
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
|
||||
__vector float& b1) {
|
||||
b0 = vec_xl(0, const_cast<float*>(p));
|
||||
b1 = vec_xl(0, const_cast<float*>(p + 4));
|
||||
}
|
||||
|
||||
// [2] BFloat16 Specialization (Little Endian ppc64le)
|
||||
// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32.
|
||||
// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)]
|
||||
// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed.
|
||||
// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]}
|
||||
// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair:
|
||||
// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs
|
||||
// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits
|
||||
// Cast to float32 gives exactly (bf16_bits << 16) per element.
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
__vector unsigned short raw = vec_xl(
|
||||
0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p)));
|
||||
__vector unsigned short zeros = vec_splat_u16(0);
|
||||
|
||||
// LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32
|
||||
b0 = (__vector float)vec_mergeh(zeros, raw);
|
||||
b1 = (__vector float)vec_mergel(zeros, raw);
|
||||
}
|
||||
|
||||
// [3] Half (FP16) Specialization
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
vec_op::FP16Vec8 fp16_vec(p);
|
||||
vec_op::FP32Vec8 fp32_vec(fp16_vec);
|
||||
b0 = fp32_vec.reg.val[0];
|
||||
b1 = fp32_vec.reg.val[1];
|
||||
}
|
||||
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K]
|
||||
const kv_cache_t* __restrict B, // [K x 8]
|
||||
float* __restrict C, // [M x 8]
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
// 1. Define A pointers
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
// 2. Define Accumulators (2 vectors covers 8 columns)
|
||||
#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
// 3. Initialize Accumulators (Load C or Zero)
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \
|
||||
acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \
|
||||
} else { \
|
||||
acc##i##_0 = vec_splats(0.0f); \
|
||||
acc##i##_1 = vec_splats(0.0f); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
for (; k + 3 < K; k += 4) {
|
||||
// Load 4 values of A for each Row M: A[k...k+3]
|
||||
#define LOAD_A4(i) \
|
||||
__vector float a##i##v; \
|
||||
IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k));
|
||||
ROWS_APPLY(LOAD_A4)
|
||||
#undef LOAD_A4
|
||||
|
||||
// FMA for specific lane L of A
|
||||
// ppc64le: vec_madd(b, vec_splat(a, lane), acc)
|
||||
#define FMAS_LANE(i, aiv, L) \
|
||||
IF_M(i) { \
|
||||
__vector float a_broad = vec_splat(aiv, L); \
|
||||
acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
|
||||
}
|
||||
|
||||
// Unroll K=0..3
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
|
||||
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
|
||||
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
|
||||
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
|
||||
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
#undef FMAS_LANE
|
||||
}
|
||||
|
||||
for (; k < K; ++k) {
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
__vector float ai = vec_splats(*(a##i + k)); \
|
||||
acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { \
|
||||
vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
|
||||
vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
|
||||
}
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda,
|
||||
int64_t ldb, int64_t ldc,
|
||||
bool accumulate) {
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmPPC64 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_ppc64le_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_ppc64le_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VSX;
|
||||
constexpr static bool scale_on_logits =
|
||||
false; // Scale is applied to Q during copy
|
||||
|
||||
public:
|
||||
AttentionImpl() {}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmPPC64<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// Strides for Memory Layout
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // [head_dim, block_size] layout
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
__vector float scale_vec = vec_splats(scale);
|
||||
constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
|
||||
|
||||
for (int32_t i = 0; i < q_num; ++i) {
|
||||
for (int32_t h = 0; h < q_heads_per_kv; ++h) {
|
||||
scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
|
||||
float* curr_dst =
|
||||
q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
|
||||
|
||||
int32_t d = 0;
|
||||
for (; d <= head_dim - 8; d += 8) {
|
||||
__vector float v0, v1;
|
||||
load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
}
|
||||
|
||||
for (; d < head_dim; ++d) {
|
||||
float val = static_cast<float>(curr_src[d]);
|
||||
curr_dst[d] = val * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) continue;
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
|
||||
{
|
||||
const scalar_t* key_src = key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_dst[j] = key_src[i];
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
const scalar_t* val_src = value + token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
|
||||
std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif // CPU_ATTN_VSX_HPP
|
||||
@@ -0,0 +1,387 @@
|
||||
#ifndef CPU_ATTN_VXE_HPP
|
||||
#define CPU_ATTN_VXE_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include <vecintrin.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// s390x Vector = 16 bytes (128 bits)
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
|
||||
__vector float& b1);
|
||||
|
||||
// [1] Float Specialization
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
|
||||
__vector float& b1) {
|
||||
// Explicitly cast to long long for offset, and float* for pointer
|
||||
b0 = vec_xl((long long)0, const_cast<float*>(p));
|
||||
b1 = vec_xl((long long)0, const_cast<float*>(p + 4));
|
||||
}
|
||||
|
||||
// [2] BFloat16 Specialization (Big Endian Fix)
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
// 1. Load 8 BF16s (16 bytes) into one vector
|
||||
// Explicit cast to unsigned short* for vec_xl to return vector unsigned short
|
||||
__vector unsigned short raw = vec_xl((long long)0, (unsigned short*)p);
|
||||
|
||||
// 2. Prepare Zero vector
|
||||
__vector unsigned short zeros = vec_splat_u16(0);
|
||||
|
||||
// 3. Merge High/Low to expand BF16 -> Float32
|
||||
// On Big Endian, a float is [BF16_bits | 16_zero_bits]
|
||||
b0 = (__vector float)vec_mergeh(raw, zeros);
|
||||
b1 = (__vector float)vec_mergel(raw, zeros);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::Half>(const c10::Half* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
alignas(16) float tmp[8];
|
||||
|
||||
// Manual unroll / conversion
|
||||
tmp[0] = static_cast<float>(p[0]);
|
||||
tmp[1] = static_cast<float>(p[1]);
|
||||
tmp[2] = static_cast<float>(p[2]);
|
||||
tmp[3] = static_cast<float>(p[3]);
|
||||
tmp[4] = static_cast<float>(p[4]);
|
||||
tmp[5] = static_cast<float>(p[5]);
|
||||
tmp[6] = static_cast<float>(p[6]);
|
||||
tmp[7] = static_cast<float>(p[7]);
|
||||
|
||||
// Explicit arguments for intrinsic: (long long offset, float* ptr)
|
||||
b0 = vec_xl((long long)0, (float*)tmp);
|
||||
b1 = vec_xl((long long)0, (float*)(tmp + 4));
|
||||
}
|
||||
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_s390x_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K]
|
||||
const kv_cache_t* __restrict B, // [K x 8]
|
||||
float* __restrict C, // [M x 8]
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
// Helper macros to unroll codegen for M rows
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
// 1. Define A pointers
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
// 2. Define Accumulators (2 vectors covers 8 columns)
|
||||
#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
// 3. Initialize Accumulators (Load C or Zero)
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i##_0 = \
|
||||
vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 0)); \
|
||||
acc##i##_1 = \
|
||||
vec_xl((long long)0, const_cast<float*>(C + (i) * ldc + 4)); \
|
||||
} else { \
|
||||
acc##i##_0 = vec_splats(0.0f); \
|
||||
acc##i##_1 = vec_splats(0.0f); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
for (; k + 3 < K; k += 4) {
|
||||
// Load 4 values of A for each Row M: A[k...k+3]
|
||||
#define LOAD_A4(i) \
|
||||
__vector float a##i##v; \
|
||||
IF_M(i) a##i##v = vec_xl((long long)0, const_cast<float*>(a##i + k));
|
||||
ROWS_APPLY(LOAD_A4)
|
||||
#undef LOAD_A4
|
||||
|
||||
// Helper: FMA for specific lane L of A
|
||||
// s390x: vec_madd(b, vec_splat(a, lane), acc)
|
||||
#define FMAS_LANE(i, aiv, L) \
|
||||
IF_M(i) { \
|
||||
__vector float a_broad = vec_splat(aiv, L); \
|
||||
acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
|
||||
}
|
||||
|
||||
// Unroll K=0..3
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
|
||||
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
|
||||
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
|
||||
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
|
||||
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
#undef FMAS_LANE
|
||||
}
|
||||
|
||||
for (; k < K; ++k) {
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
__vector float ai = vec_splats(*(a##i + k)); \
|
||||
acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { \
|
||||
vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
|
||||
vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
|
||||
}
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_s390x_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda, int64_t ldb,
|
||||
int64_t ldc, bool accumulate) {
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_s390x_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_s390x_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_s390x_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_s390x_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc, K,
|
||||
accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmS390X {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_s390x_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_s390x_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VXE;
|
||||
constexpr static bool scale_on_logits =
|
||||
false; // Scale is applied to Q during copy
|
||||
|
||||
public:
|
||||
AttentionImpl() {}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmS390X<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// Strides for Memory Layout
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // [head_dim, block_size] layout
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
__vector float scale_vec = vec_splats(scale);
|
||||
constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
|
||||
|
||||
// Process 8 elements at a time (32 bytes of float output)
|
||||
for (int32_t i = 0; i < q_num; ++i) {
|
||||
for (int32_t h = 0; h < q_heads_per_kv; ++h) {
|
||||
scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
|
||||
float* curr_dst =
|
||||
q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
|
||||
|
||||
int32_t d = 0;
|
||||
for (; d <= head_dim - 8; d += 8) {
|
||||
if constexpr (is_bf16) {
|
||||
__vector float v0, v1;
|
||||
// Reuse our Big-Endian-Safe loader
|
||||
load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
} else {
|
||||
__vector float v0 = vec_xl((long long)0, (float*)curr_src + d);
|
||||
__vector float v1 = vec_xl((long long)0, (float*)curr_src + d + 4);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
}
|
||||
}
|
||||
|
||||
for (; d < head_dim; ++d) {
|
||||
float val = static_cast<float>(curr_src[d]);
|
||||
curr_dst[d] = val * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) continue;
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
|
||||
{
|
||||
const scalar_t* key_src = key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_dst[j] = key_src[i];
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
const scalar_t* val_src = value + token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
|
||||
std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,884 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
#include "cpu/utils.hpp"
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
|
||||
#include "cpu/cpu_arch_macros.h"
|
||||
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
|
||||
#define AMX_DISPATCH(...) \
|
||||
case cpu_utils::ISA::AMX: { \
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<cpu_utils::ISA::AMX, scalar_t>; \
|
||||
return __VA_ARGS__(); \
|
||||
}
|
||||
#else
|
||||
#define AMX_DISPATCH(...) case cpu_utils::ISA::AMX:
|
||||
#endif
|
||||
|
||||
#if defined(ARM_BF16_SUPPORT)
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_neon.hpp"
|
||||
#define NEON_DISPATCH(...) \
|
||||
case cpu_utils::ISA::NEON: { \
|
||||
using gemm_t = \
|
||||
cpu_micro_gemm::MicroGemm<cpu_utils::ISA::NEON, scalar_t>; \
|
||||
return __VA_ARGS__(); \
|
||||
}
|
||||
#else
|
||||
#define NEON_DISPATCH(...) case cpu_utils::ISA::NEON:
|
||||
#endif
|
||||
|
||||
#define CPU_ISA_DISPATCH_IMPL(ISA_TYPE, ...) \
|
||||
[&] { \
|
||||
switch (ISA_TYPE) { \
|
||||
AMX_DISPATCH(__VA_ARGS__) \
|
||||
case cpu_utils::ISA::VEC: { \
|
||||
using gemm_t = \
|
||||
cpu_micro_gemm::MicroGemm<cpu_utils::ISA::VEC, scalar_t>; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
NEON_DISPATCH(__VA_ARGS__) \
|
||||
default: { \
|
||||
TORCH_CHECK(false, "Invalid CPU ISA type."); \
|
||||
} \
|
||||
} \
|
||||
}()
|
||||
|
||||
namespace {
|
||||
enum class FusedMOEAct {
|
||||
SiluAndMul,
|
||||
SwigluOAIAndMul,
|
||||
GeluAndMul,
|
||||
GeluTanhAndMul,
|
||||
};
|
||||
|
||||
FusedMOEAct get_act_type(const std::string& act) {
|
||||
if (act == "silu") {
|
||||
return FusedMOEAct::SiluAndMul;
|
||||
} else if (act == "swigluoai") {
|
||||
return FusedMOEAct::SwigluOAIAndMul;
|
||||
} else if (act == "gelu") {
|
||||
return FusedMOEAct::GeluAndMul;
|
||||
} else if (act == "gelu_tanh") {
|
||||
return FusedMOEAct::GeluTanhAndMul;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid act type: " + act);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
#if !defined(__aarch64__)
|
||||
// For GPT-OSS interleaved gate-up weights
|
||||
alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
|
||||
16, 18, 20, 22, 24, 26, 28, 30};
|
||||
vec_op::INT32Vec16 index_vec(index);
|
||||
#endif
|
||||
vec_op::FP32Vec16 gate_up_max_vec(7.0);
|
||||
vec_op::FP32Vec16 up_min_vec(-7.0);
|
||||
vec_op::FP32Vec16 alpha_vec(1.702);
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < n_size; n += 32) {
|
||||
// Note: AdvSIMD does not support gather loads
|
||||
#if defined(__aarch64__)
|
||||
vec_op::FP32Vec16 gate_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16 up_vec(vec_op::uninit);
|
||||
vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
|
||||
#else
|
||||
vec_op::FP32Vec16 gate_vec(input + n, index_vec);
|
||||
vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
|
||||
#endif
|
||||
gate_vec = gate_vec.min(gate_up_max_vec);
|
||||
up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
|
||||
auto glu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = (one_vec + up_vec) * glu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n / 2);
|
||||
}
|
||||
input += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
|
||||
auto silu = gate_vec * sigmoid_vec;
|
||||
auto gated_output_fp32 = up_vec * silu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride, const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
alignas(64) float temp[16];
|
||||
|
||||
DEFINE_FAST_EXP
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto er_input_vec = gate_vec * w1_vec;
|
||||
|
||||
er_input_vec.save(temp);
|
||||
for (int32_t i = 0; i < 16; ++i) {
|
||||
temp[i] = std::erf(temp[i]);
|
||||
}
|
||||
vec_op::FP32Vec16 er_vec(temp);
|
||||
auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
|
||||
const int32_t m_size, const int32_t n_size,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
const int32_t dim = n_size / 2;
|
||||
float* __restrict__ gate = input;
|
||||
float* __restrict__ up = input + dim;
|
||||
vec_op::FP32Vec16 one_vec(1.0);
|
||||
vec_op::FP32Vec16 w1_vec(0.7978845608028654);
|
||||
vec_op::FP32Vec16 w2_vec(0.5);
|
||||
vec_op::FP32Vec16 w3_vec(0.044715);
|
||||
|
||||
for (int32_t m = 0; m < m_size; ++m) {
|
||||
for (int32_t n = 0; n < dim; n += 16) {
|
||||
vec_op::FP32Vec16 gate_vec(gate + n);
|
||||
vec_op::FP32Vec16 up_vec(up + n);
|
||||
auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
|
||||
auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
|
||||
// Note: can't use fast_exp form because diffusiongemma will generate
|
||||
// wrong results
|
||||
auto tanh_vec = inner_vec.tanh();
|
||||
auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
|
||||
auto gated_output_fp32 = up_vec * gelu_tanh;
|
||||
scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
|
||||
gated_output.save(output + n);
|
||||
}
|
||||
gate += input_stride;
|
||||
up += input_stride;
|
||||
output += output_stride;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
|
||||
float* __restrict__ input,
|
||||
scalar_t* __restrict__ output,
|
||||
const int32_t m, const int32_t n,
|
||||
const int32_t input_stride,
|
||||
const int32_t output_stride) {
|
||||
switch (act) {
|
||||
case FusedMOEAct::SwigluOAIAndMul:
|
||||
swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::SiluAndMul:
|
||||
silu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluAndMul:
|
||||
gelu_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
case FusedMOEAct::GeluTanhAndMul:
|
||||
gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
|
||||
return;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unsupported act type.");
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename gemm_t>
|
||||
void prepack_moe_weight_impl(scalar_t* __restrict__ weight_ptr,
|
||||
scalar_t* __restrict__ packed_weight_ptr,
|
||||
const int32_t expert_num,
|
||||
const int32_t output_size,
|
||||
const int32_t input_size,
|
||||
const int64_t expert_stride) {
|
||||
#pragma omp parallel for
|
||||
for (int32_t e_idx = 0; e_idx < expert_num; ++e_idx) {
|
||||
gemm_t::pack_weight(weight_ptr + expert_stride * e_idx,
|
||||
packed_weight_ptr + expert_stride * e_idx, output_size,
|
||||
input_size);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename w_t, typename gemm_t>
|
||||
void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
|
||||
w_t* __restrict__ w13, w_t* __restrict__ w2,
|
||||
w_t* __restrict__ w13_bias, w_t* __restrict__ w2_bias,
|
||||
float* __restrict__ topk_weights,
|
||||
int32_t* __restrict__ topk_id, FusedMOEAct act_type,
|
||||
const int32_t token_num, const int32_t expert_num,
|
||||
const int32_t topk_num, const int32_t input_size_13,
|
||||
const int32_t output_size_13, const int32_t input_size_2,
|
||||
const int32_t output_size_2, const bool skip_weighted) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
|
||||
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
|
||||
constexpr int32_t min_w13_n_tile_size = 2 * gemm_n_tile_size;
|
||||
constexpr bool pack_a = gemm_t::PackA;
|
||||
static_assert(gemm_n_tile_size % 16 == 0);
|
||||
|
||||
TORCH_CHECK_EQ(output_size_13 % min_w13_n_tile_size, 0);
|
||||
TORCH_CHECK_EQ(output_size_2 % gemm_n_tile_size, 0);
|
||||
TORCH_CHECK_EQ(output_size_13 / 2, input_size_2);
|
||||
|
||||
const int32_t thread_num = cpu_utils::get_max_threads();
|
||||
|
||||
const int32_t w13_input_buffer_size = cpu_utils::round_up<64>(
|
||||
gemm_m_tile_size * input_size_13 * sizeof(scalar_t));
|
||||
|
||||
const int32_t w13_n_tile_size = [&]() {
|
||||
const int64_t cache_size = cpu_utils::get_available_l2_size();
|
||||
// input buffer + output buffer + weight
|
||||
const int32_t n_size_cache_limit =
|
||||
(cache_size - w13_input_buffer_size) /
|
||||
(gemm_m_tile_size * sizeof(float) + input_size_13 * sizeof(scalar_t));
|
||||
const int32_t n_size_thread_limit =
|
||||
output_size_13 / std::max(1, thread_num / topk_num);
|
||||
const int32_t n_size = cpu_utils::round_down<min_w13_n_tile_size>(
|
||||
std::min(n_size_cache_limit, n_size_thread_limit));
|
||||
return std::max(n_size, min_w13_n_tile_size);
|
||||
}();
|
||||
|
||||
const int32_t w2_input_tile_size = cpu_utils::round_up<64>(
|
||||
gemm_m_tile_size * input_size_2 * sizeof(scalar_t));
|
||||
// use w2 input buffer only when we need to pack input
|
||||
const int32_t w2_input_buffer_size =
|
||||
pack_a ? cpu_utils::round_up<64>(gemm_m_tile_size * input_size_2 *
|
||||
sizeof(scalar_t))
|
||||
: 0;
|
||||
|
||||
const int32_t w2_n_tile_size = [&]() {
|
||||
const int64_t cache_size = cpu_utils::get_available_l2_size();
|
||||
// input tile + optional packed input + weight
|
||||
const int32_t n_size_cache_limit =
|
||||
(cache_size - (pack_a ? w2_input_buffer_size : w2_input_tile_size)) /
|
||||
(input_size_2 * sizeof(scalar_t));
|
||||
const int32_t n_size_thread_limit =
|
||||
output_size_2 / std::max(1, thread_num / topk_num);
|
||||
const int32_t n_size = cpu_utils::round_down<gemm_n_tile_size>(
|
||||
std::min(n_size_cache_limit, n_size_thread_limit));
|
||||
return std::max(n_size, gemm_n_tile_size);
|
||||
}();
|
||||
|
||||
// allocate buffers
|
||||
int32_t common_buffer_offset = 0;
|
||||
int32_t w13_thread_buffer_offset = 0;
|
||||
int32_t ws_thread_buffer_offset = 0;
|
||||
|
||||
// common buffers
|
||||
const int32_t token_num_per_group_buffer_size =
|
||||
cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
|
||||
const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += token_num_per_group_buffer_size;
|
||||
|
||||
const int32_t cu_token_num_per_group_buffer_size =
|
||||
cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
|
||||
const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += cu_token_num_per_group_buffer_size;
|
||||
|
||||
const int32_t expand_token_id_buffer_size =
|
||||
cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
|
||||
const int32_t expand_token_id_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += expand_token_id_buffer_size;
|
||||
|
||||
const int32_t expand_token_id_index_buffer_size =
|
||||
cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
|
||||
const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += expand_token_id_index_buffer_size;
|
||||
|
||||
const int32_t w13_gemm_output_buffer_size = cpu_utils::round_up<64>(
|
||||
token_num * topk_num * (output_size_13 / 2) * sizeof(scalar_t));
|
||||
const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += w13_gemm_output_buffer_size;
|
||||
|
||||
const int32_t w2_gemm_output_buffer_size = cpu_utils::round_up<64>(
|
||||
token_num * topk_num * output_size_2 * sizeof(float));
|
||||
const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
|
||||
common_buffer_offset += w2_gemm_output_buffer_size;
|
||||
|
||||
// w13 GEMM thread buffers
|
||||
const int32_t w13_input_buffer_offset = w13_thread_buffer_offset;
|
||||
w13_thread_buffer_offset += w13_input_buffer_size;
|
||||
|
||||
const int32_t w13_output_buffer_size = cpu_utils::round_up<64>(
|
||||
gemm_m_tile_size * w13_n_tile_size * sizeof(float));
|
||||
const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;
|
||||
w13_thread_buffer_offset += w13_output_buffer_size;
|
||||
|
||||
const int32_t w2_input_buffer_offset = w13_thread_buffer_offset;
|
||||
w13_thread_buffer_offset += w2_input_buffer_size;
|
||||
|
||||
// Weighted sum thread buffer
|
||||
const int32_t ws_output_buffer_size =
|
||||
cpu_utils::round_up<64>(output_size_2 * sizeof(float));
|
||||
const int32_t ws_output_buffer_offset = ws_thread_buffer_offset;
|
||||
ws_thread_buffer_offset += ws_output_buffer_size;
|
||||
|
||||
const int32_t buffer_size =
|
||||
common_buffer_offset +
|
||||
std::max(w13_thread_buffer_offset, ws_thread_buffer_offset) * thread_num;
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
|
||||
uint8_t* common_buffer_start =
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()
|
||||
->get_data<uint8_t>();
|
||||
uint8_t* thread_buffer_start = common_buffer_start + common_buffer_offset;
|
||||
|
||||
int32_t* __restrict__ token_num_per_group_buffer = reinterpret_cast<int32_t*>(
|
||||
common_buffer_start + token_num_per_group_buffer_offset);
|
||||
int32_t* __restrict__ cu_token_num_per_group_buffer =
|
||||
reinterpret_cast<int32_t*>(common_buffer_start +
|
||||
cu_token_num_per_group_buffer_offset);
|
||||
int32_t* __restrict__ expand_token_id_buffer = reinterpret_cast<int32_t*>(
|
||||
common_buffer_start + expand_token_id_buffer_offset);
|
||||
int32_t* __restrict__ expand_token_id_index_buffer =
|
||||
reinterpret_cast<int32_t*>(common_buffer_start +
|
||||
expand_token_id_index_buffer_offset);
|
||||
|
||||
// prepare token-expert mappings
|
||||
{
|
||||
std::memset(token_num_per_group_buffer, 0, expert_num * sizeof(int32_t));
|
||||
for (int32_t i = 0; i < token_num * topk_num; ++i) {
|
||||
int32_t curr_expert_id = topk_id[i];
|
||||
++token_num_per_group_buffer[curr_expert_id];
|
||||
}
|
||||
|
||||
int32_t token_num_sum = 0;
|
||||
cu_token_num_per_group_buffer[0] = 0;
|
||||
int32_t* token_index_buffer = cu_token_num_per_group_buffer + 1;
|
||||
for (int32_t i = 0; i < expert_num; ++i) {
|
||||
token_index_buffer[i] = token_num_sum;
|
||||
token_num_sum += token_num_per_group_buffer[i];
|
||||
}
|
||||
|
||||
for (int32_t i = 0; i < token_num; ++i) {
|
||||
int32_t* curr_topk_id = topk_id + i * topk_num;
|
||||
int32_t* curr_index_buffer = expand_token_id_index_buffer + i * topk_num;
|
||||
for (int32_t j = 0; j < topk_num; ++j) {
|
||||
int32_t curr_expert_id = curr_topk_id[j];
|
||||
int32_t curr_index = token_index_buffer[curr_expert_id];
|
||||
++token_index_buffer[curr_expert_id];
|
||||
expand_token_id_buffer[curr_index] = i;
|
||||
curr_index_buffer[j] = curr_index;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// w13 GEMM + act
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
const int32_t task_num_per_expert =
|
||||
(output_size_13 + w13_n_tile_size - 1) / w13_n_tile_size;
|
||||
const int32_t task_num = task_num_per_expert * expert_num;
|
||||
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * w13_thread_buffer_offset;
|
||||
scalar_t* __restrict__ w13_input_buffer =
|
||||
reinterpret_cast<scalar_t*>(thread_buffer + w13_input_buffer_offset);
|
||||
float* __restrict__ w13_output_buffer =
|
||||
reinterpret_cast<float*>(thread_buffer + w13_output_buffer_offset);
|
||||
scalar_t* __restrict__ w13_gemm_output_buffer =
|
||||
reinterpret_cast<scalar_t*>(common_buffer_start +
|
||||
w13_gemm_output_buffer_offset);
|
||||
|
||||
gemm_t gemm;
|
||||
|
||||
const int32_t input_size_13_bytes = input_size_13 * sizeof(scalar_t);
|
||||
const int32_t w13_n_group_stride =
|
||||
gemm_t::WeightOCGroupSize * input_size_13;
|
||||
const int32_t w13_n_tile_stride = gemm_n_tile_size * input_size_13;
|
||||
|
||||
for (;;) {
|
||||
int32_t task_id = counter_ptr->acquire_counter();
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
|
||||
const int32_t curr_expert_id = task_id / task_num_per_expert;
|
||||
const int32_t curr_output_group_id = task_id % task_num_per_expert;
|
||||
const int32_t curr_token_num =
|
||||
token_num_per_group_buffer[curr_expert_id];
|
||||
if (curr_token_num == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t actual_n_tile_size =
|
||||
std::min(w13_n_tile_size,
|
||||
output_size_13 - curr_output_group_id * w13_n_tile_size);
|
||||
const int32_t* __restrict__ curr_expand_token_id_buffer =
|
||||
expand_token_id_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id];
|
||||
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
|
||||
w13_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] *
|
||||
(output_size_13 / 2) +
|
||||
curr_output_group_id * w13_n_tile_size / 2;
|
||||
|
||||
w_t* __restrict__ w13_weight_ptr_0 = nullptr;
|
||||
w_t* __restrict__ w13_weight_ptr_1 = nullptr;
|
||||
w_t* __restrict__ w13_bias_ptr_0 = nullptr;
|
||||
w_t* __restrict__ w13_bias_ptr_1 = nullptr;
|
||||
if (act_type == FusedMOEAct::SwigluOAIAndMul) {
|
||||
// For SwigluOAIAndMul, up and down weights are interleaved
|
||||
w13_weight_ptr_0 =
|
||||
w13 + curr_expert_id * input_size_13 * output_size_13 +
|
||||
curr_output_group_id * w13_n_tile_size * input_size_13;
|
||||
w13_weight_ptr_1 =
|
||||
w13_weight_ptr_0 + actual_n_tile_size / 2 * input_size_13;
|
||||
if (w13_bias != nullptr) {
|
||||
w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 +
|
||||
curr_output_group_id * w13_n_tile_size;
|
||||
w13_bias_ptr_1 = w13_bias_ptr_0 + actual_n_tile_size / 2;
|
||||
}
|
||||
} else {
|
||||
w13_weight_ptr_0 =
|
||||
w13 + curr_expert_id * input_size_13 * output_size_13 +
|
||||
curr_output_group_id * (w13_n_tile_size / 2) * input_size_13;
|
||||
w13_weight_ptr_1 =
|
||||
w13_weight_ptr_0 + output_size_13 / 2 * input_size_13;
|
||||
if (w13_bias != nullptr) {
|
||||
w13_bias_ptr_0 = w13_bias + curr_expert_id * output_size_13 +
|
||||
curr_output_group_id * (w13_n_tile_size / 2);
|
||||
w13_bias_ptr_1 = w13_bias_ptr_0 + output_size_13 / 2;
|
||||
}
|
||||
}
|
||||
|
||||
scalar_t* __restrict__ curr_w13_input_buffer = w13_input_buffer;
|
||||
for (int32_t token_idx = 0; token_idx < curr_token_num;
|
||||
token_idx += gemm_m_tile_size) {
|
||||
const int32_t actual_token_num =
|
||||
std::min(gemm_m_tile_size, curr_token_num - token_idx);
|
||||
|
||||
scalar_t* __restrict__ curr_w13_gemm_input_buffer = nullptr;
|
||||
if constexpr (pack_a) {
|
||||
// copy and pack inputs
|
||||
curr_w13_gemm_input_buffer = w13_input_buffer;
|
||||
const scalar_t* w13_input_rows[gemm_m_tile_size];
|
||||
for (int32_t i = 0; i < actual_token_num; ++i) {
|
||||
w13_input_rows[i] =
|
||||
input + curr_expand_token_id_buffer[i] * input_size_13;
|
||||
}
|
||||
gemm_t::pack_input_from_rows(w13_input_rows,
|
||||
curr_w13_gemm_input_buffer,
|
||||
actual_token_num, input_size_13);
|
||||
curr_expand_token_id_buffer += actual_token_num;
|
||||
} else {
|
||||
// copy inputs
|
||||
curr_w13_gemm_input_buffer = curr_w13_input_buffer;
|
||||
scalar_t* __restrict__ curr_w13_input_buffer_iter =
|
||||
curr_w13_input_buffer;
|
||||
for (int32_t i = 0; i < actual_token_num; ++i) {
|
||||
const int32_t curr_token_id = curr_expand_token_id_buffer[i];
|
||||
int8_t* __restrict__ curr_input_iter = reinterpret_cast<int8_t*>(
|
||||
input + curr_token_id * input_size_13);
|
||||
int8_t* __restrict__ curr_output_iter =
|
||||
reinterpret_cast<int8_t*>(curr_w13_input_buffer_iter);
|
||||
int32_t j = 0;
|
||||
for (; j < input_size_13_bytes - 64; j += 64) {
|
||||
vec_op::INT8Vec64 vec(curr_input_iter);
|
||||
vec.save(curr_output_iter);
|
||||
curr_input_iter += 64;
|
||||
curr_output_iter += 64;
|
||||
}
|
||||
vec_op::INT8Vec64 vec(curr_input_iter);
|
||||
vec.save(curr_output_iter, input_size_13_bytes - j);
|
||||
|
||||
// update
|
||||
curr_w13_input_buffer_iter += input_size_13;
|
||||
}
|
||||
// update
|
||||
curr_expand_token_id_buffer += actual_token_num;
|
||||
}
|
||||
|
||||
// gemm + act
|
||||
{
|
||||
scalar_t* __restrict__ w13_weight_ptr_0_iter = w13_weight_ptr_0;
|
||||
scalar_t* __restrict__ w13_weight_ptr_1_iter = w13_weight_ptr_1;
|
||||
scalar_t* __restrict__ w13_bias_ptr_0_iter = w13_bias_ptr_0;
|
||||
scalar_t* __restrict__ w13_bias_ptr_1_iter = w13_bias_ptr_1;
|
||||
float* __restrict__ w13_output_buffer_0_iter = w13_output_buffer;
|
||||
float* __restrict__ w13_output_buffer_1_iter =
|
||||
w13_output_buffer + actual_n_tile_size / 2;
|
||||
for (int32_t i = 0; i < actual_n_tile_size;
|
||||
i += min_w13_n_tile_size) {
|
||||
gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_0_iter,
|
||||
w13_output_buffer_0_iter, actual_token_num,
|
||||
input_size_13, input_size_13, w13_n_group_stride,
|
||||
actual_n_tile_size, false);
|
||||
|
||||
if (w13_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
w13_output_buffer_0_iter, w13_output_buffer_0_iter,
|
||||
w13_bias_ptr_0_iter, actual_token_num, actual_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
w13_bias_ptr_0_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
gemm.gemm(curr_w13_gemm_input_buffer, w13_weight_ptr_1_iter,
|
||||
w13_output_buffer_1_iter, actual_token_num,
|
||||
input_size_13, input_size_13, w13_n_group_stride,
|
||||
actual_n_tile_size, false);
|
||||
|
||||
if (w13_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
w13_output_buffer_1_iter, w13_output_buffer_1_iter,
|
||||
w13_bias_ptr_1_iter, actual_token_num, actual_n_tile_size,
|
||||
actual_n_tile_size);
|
||||
w13_bias_ptr_1_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
// update
|
||||
w13_weight_ptr_0_iter += w13_n_tile_stride;
|
||||
w13_weight_ptr_1_iter += w13_n_tile_stride;
|
||||
w13_output_buffer_0_iter += gemm_n_tile_size;
|
||||
w13_output_buffer_1_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
apply_gated_act(act_type, w13_output_buffer,
|
||||
curr_w13_gemm_output_buffer, actual_token_num,
|
||||
actual_n_tile_size, actual_n_tile_size,
|
||||
output_size_13 / 2);
|
||||
|
||||
// update
|
||||
curr_w13_gemm_output_buffer +=
|
||||
gemm_m_tile_size * (output_size_13 / 2);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// w2 GEMM
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
const int32_t task_num_per_expert =
|
||||
(output_size_2 + w2_n_tile_size - 1) / w2_n_tile_size;
|
||||
const int32_t task_num = task_num_per_expert * expert_num;
|
||||
scalar_t* __restrict__ w13_gemm_output_buffer =
|
||||
reinterpret_cast<scalar_t*>(common_buffer_start +
|
||||
w13_gemm_output_buffer_offset);
|
||||
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
|
||||
common_buffer_start + w2_gemm_output_buffer_offset);
|
||||
|
||||
gemm_t gemm;
|
||||
|
||||
const int32_t w2_n_tile_stride = gemm_n_tile_size * input_size_2;
|
||||
const int32_t w2_n_group_stride =
|
||||
gemm_t::WeightOCGroupSize * input_size_2;
|
||||
|
||||
for (;;) {
|
||||
int32_t task_id = counter_ptr->acquire_counter();
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
|
||||
const int32_t curr_expert_id = task_id / task_num_per_expert;
|
||||
const int32_t curr_output_group_id = task_id % task_num_per_expert;
|
||||
const int32_t curr_token_num =
|
||||
token_num_per_group_buffer[curr_expert_id];
|
||||
if (curr_token_num == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t actual_n_tile_size =
|
||||
std::min(w2_n_tile_size,
|
||||
output_size_2 - curr_output_group_id * w2_n_tile_size);
|
||||
scalar_t* __restrict__ curr_w13_gemm_output_buffer =
|
||||
w13_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] * input_size_2;
|
||||
float* __restrict__ curr_w2_gemm_output_buffer =
|
||||
w2_gemm_output_buffer +
|
||||
cu_token_num_per_group_buffer[curr_expert_id] * output_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size;
|
||||
scalar_t* __restrict__ w2_weight_ptr =
|
||||
w2 + curr_expert_id * output_size_2 * input_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size * input_size_2;
|
||||
scalar_t* __restrict__ w2_bias_ptr = nullptr;
|
||||
if (w2_bias != nullptr) {
|
||||
w2_bias_ptr = w2_bias + curr_expert_id * output_size_2 +
|
||||
curr_output_group_id * w2_n_tile_size;
|
||||
}
|
||||
|
||||
for (int32_t token_idx = 0; token_idx < curr_token_num;
|
||||
token_idx += gemm_m_tile_size) {
|
||||
const int32_t actual_token_num =
|
||||
std::min(gemm_m_tile_size, curr_token_num - token_idx);
|
||||
scalar_t* __restrict__ curr_w2_gemm_input_buffer =
|
||||
curr_w13_gemm_output_buffer;
|
||||
if constexpr (pack_a) {
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * w13_thread_buffer_offset;
|
||||
scalar_t* __restrict__ w2_input_buffer =
|
||||
reinterpret_cast<scalar_t*>(thread_buffer +
|
||||
w2_input_buffer_offset);
|
||||
curr_w2_gemm_input_buffer = w2_input_buffer;
|
||||
const scalar_t* w2_input_rows[gemm_m_tile_size];
|
||||
for (int32_t i = 0; i < actual_token_num; ++i) {
|
||||
w2_input_rows[i] = curr_w13_gemm_output_buffer + i * input_size_2;
|
||||
}
|
||||
gemm_t::pack_input_from_rows(w2_input_rows,
|
||||
curr_w2_gemm_input_buffer,
|
||||
actual_token_num, input_size_2);
|
||||
}
|
||||
|
||||
scalar_t* __restrict__ w2_weight_ptr_iter = w2_weight_ptr;
|
||||
scalar_t* __restrict__ w2_bias_ptr_iter = w2_bias_ptr;
|
||||
float* __restrict__ curr_w2_gemm_output_buffer_iter =
|
||||
curr_w2_gemm_output_buffer;
|
||||
for (int32_t i = 0; i < actual_n_tile_size; i += gemm_n_tile_size) {
|
||||
gemm.gemm(curr_w2_gemm_input_buffer, w2_weight_ptr_iter,
|
||||
curr_w2_gemm_output_buffer_iter, actual_token_num,
|
||||
input_size_2, input_size_2, w2_n_group_stride,
|
||||
output_size_2, false);
|
||||
|
||||
if (w2_bias != nullptr) {
|
||||
cpu_micro_gemm::add_bias_epilogue<gemm_n_tile_size>(
|
||||
curr_w2_gemm_output_buffer_iter,
|
||||
curr_w2_gemm_output_buffer_iter, w2_bias_ptr_iter,
|
||||
actual_token_num, output_size_2, output_size_2);
|
||||
w2_bias_ptr_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
w2_weight_ptr_iter += w2_n_tile_stride;
|
||||
curr_w2_gemm_output_buffer_iter += gemm_n_tile_size;
|
||||
}
|
||||
|
||||
// update
|
||||
curr_w13_gemm_output_buffer += gemm_m_tile_size * input_size_2;
|
||||
curr_w2_gemm_output_buffer += gemm_m_tile_size * output_size_2;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// weighted sum
|
||||
{
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
const int32_t task_num = token_num;
|
||||
uint8_t* __restrict__ thread_buffer =
|
||||
thread_buffer_start + thread_id * ws_thread_buffer_offset;
|
||||
float* __restrict__ ws_output_buffer =
|
||||
reinterpret_cast<float*>(thread_buffer + ws_output_buffer_offset);
|
||||
float* __restrict__ w2_gemm_output_buffer = reinterpret_cast<float*>(
|
||||
common_buffer_start + w2_gemm_output_buffer_offset);
|
||||
|
||||
for (;;) {
|
||||
int32_t task_id = counter_ptr->acquire_counter();
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
|
||||
int32_t token_id = task_id;
|
||||
int32_t* __restrict__ curr_expand_token_id_index_buffer =
|
||||
expand_token_id_index_buffer + token_id * topk_num;
|
||||
float* __restrict__ curr_weight = topk_weights + token_id * topk_num;
|
||||
scalar_t* __restrict__ curr_output_buffer =
|
||||
output + token_id * output_size_2;
|
||||
|
||||
if (skip_weighted) {
|
||||
// Only for topk_num == 1
|
||||
*curr_weight = 1.0f;
|
||||
}
|
||||
|
||||
if (topk_num > 1) {
|
||||
{
|
||||
int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
|
||||
float* __restrict__ w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(curr_weight[0]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec = vec * weight_vec;
|
||||
vec.save(ws_output_buffer_iter);
|
||||
|
||||
// update
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
for (int32_t idx = 1; idx < topk_num - 1; ++idx) {
|
||||
int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx];
|
||||
float* __restrict__ w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(curr_weight[idx]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
|
||||
sum = sum + vec * weight_vec;
|
||||
sum.save(ws_output_buffer_iter);
|
||||
|
||||
// update
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
int32_t idx = topk_num - 1;
|
||||
int32_t w2_output_idx = curr_expand_token_id_index_buffer[idx];
|
||||
float* __restrict__ w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
float* __restrict__ ws_output_buffer_iter = ws_output_buffer;
|
||||
scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(curr_weight[idx]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec_op::FP32Vec16 sum(ws_output_buffer_iter);
|
||||
sum = sum + vec * weight_vec;
|
||||
scalar_vec_t out_vec(sum);
|
||||
out_vec.save(curr_output_buffer_iter);
|
||||
|
||||
// update
|
||||
w2_output_iter += 16;
|
||||
ws_output_buffer_iter += 16;
|
||||
curr_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
int32_t w2_output_idx = curr_expand_token_id_index_buffer[0];
|
||||
float* __restrict__ w2_output_iter =
|
||||
w2_gemm_output_buffer + w2_output_idx * output_size_2;
|
||||
scalar_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
|
||||
vec_op::FP32Vec16 weight_vec(curr_weight[0]);
|
||||
for (int32_t i = 0; i < output_size_2; i += 16) {
|
||||
vec_op::FP32Vec16 vec(w2_output_iter);
|
||||
vec = vec * weight_vec;
|
||||
scalar_vec_t out_vec(vec);
|
||||
out_vec.save(curr_output_buffer_iter);
|
||||
|
||||
// update
|
||||
w2_output_iter += 16;
|
||||
curr_output_buffer_iter += 16;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void prepack_moe_weight(
|
||||
const torch::Tensor& weight, // [expert_num, output_size, input_size]
|
||||
torch::Tensor& packed_weight, const std::string& isa) {
|
||||
TORCH_CHECK(weight.is_contiguous());
|
||||
const int32_t expert_num = weight.size(0);
|
||||
const int32_t output_size = weight.size(1);
|
||||
const int32_t input_size = weight.size(2);
|
||||
TORCH_CHECK_EQ(output_size % 32, 0);
|
||||
const int64_t expert_stride = weight.stride(0);
|
||||
cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
weight.scalar_type(), "prepack_moe_weight", [&]() {
|
||||
CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
|
||||
scalar_t* weight_ptr = weight.data_ptr<scalar_t>();
|
||||
scalar_t* packed_weight_ptr = packed_weight.data_ptr<scalar_t>();
|
||||
prepack_moe_weight_impl<scalar_t, gemm_t>(
|
||||
weight_ptr, packed_weight_ptr, expert_num, output_size,
|
||||
input_size, expert_stride);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
void cpu_fused_moe(
|
||||
torch::Tensor& output, // [token_num, output_size_2]
|
||||
const torch::Tensor& input, // [token_num, input_size_13]
|
||||
const torch::Tensor&
|
||||
w13, // [expert_num, output_size_13, input_size_13], packed
|
||||
const torch::Tensor&
|
||||
w2, // [expert_num, output_size_2, input_size_2], packed
|
||||
const std::optional<torch::Tensor>&
|
||||
w13_bias, // [expert_num, output_size_13]
|
||||
const std::optional<torch::Tensor>& w2_bias, // [expert_num, output_size_2]
|
||||
const torch::Tensor& topk_weights, // [token_num, k], float32
|
||||
const torch::Tensor& topk_id, // [token_num, k], int32
|
||||
const bool skip_weighted, const std::string& act, const std::string& isa) {
|
||||
const int32_t token_num = input.size(0);
|
||||
const int32_t input_size_13 = input.size(1);
|
||||
const int64_t input_stride = input.stride(0);
|
||||
TORCH_CHECK_EQ(input_stride, input_size_13);
|
||||
const int32_t expert_num = w13.size(0);
|
||||
const int32_t output_size_13 = w13.size(1);
|
||||
const int32_t input_size_2 = w2.size(2);
|
||||
const int32_t output_size_2 = w2.size(1);
|
||||
const int32_t topk_num = topk_id.size(1);
|
||||
const FusedMOEAct act_type = get_act_type(act);
|
||||
cpu_utils::ISA isa_type = cpu_utils::get_isa(isa);
|
||||
TORCH_CHECK(!skip_weighted || topk_num == 1,
|
||||
"skip_weighted is only supported for topk=1 on CPU");
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(w13.scalar_type(), "cpu_fused_moe", [&]() {
|
||||
CPU_ISA_DISPATCH_IMPL(isa_type, [&]() {
|
||||
fused_moe_impl<scalar_t, scalar_t, gemm_t>(
|
||||
output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
|
||||
w13.data_ptr<scalar_t>(), w2.data_ptr<scalar_t>(),
|
||||
w13_bias.has_value() ? w13_bias->data_ptr<scalar_t>() : nullptr,
|
||||
w2_bias.has_value() ? w2_bias->data_ptr<scalar_t>() : nullptr,
|
||||
topk_weights.data_ptr<float>(), topk_id.data_ptr<int32_t>(), act_type,
|
||||
token_num, expert_num, topk_num, input_size_13, output_size_13,
|
||||
input_size_2, output_size_2, skip_weighted);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#ifndef CPU_TANHF_NEON_HPP
|
||||
#define CPU_TANHF_NEON_HPP
|
||||
|
||||
#include <cstdint>
|
||||
#include <arm_neon.h>
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
namespace {
|
||||
|
||||
struct TanhfConstants {
|
||||
float32x4_t special_bound;
|
||||
float32x4_t two;
|
||||
float32x4_t c0;
|
||||
float32x4_t c2;
|
||||
int32x4_t exponent_bias;
|
||||
float c1;
|
||||
float c3;
|
||||
float two_over_ln2;
|
||||
float c4;
|
||||
float ln2_hi;
|
||||
float ln2_lo;
|
||||
};
|
||||
|
||||
const TanhfConstants kTanhfConstants = {
|
||||
// 9.01, above which tanhf rounds to 1 (or -1 for negative).
|
||||
.special_bound = vdupq_n_f32(0x1.205966p+3f),
|
||||
.two = vdupq_n_f32(0x1.0p+1f),
|
||||
.c0 = vdupq_n_f32(0x1.fffffep-2f),
|
||||
.c2 = vdupq_n_f32(0x1.555736p-5f),
|
||||
.exponent_bias = vdupq_n_s32(0x3f800000),
|
||||
.c1 = 0x1.5554aep-3f,
|
||||
.c3 = 0x1.12287cp-7f,
|
||||
.two_over_ln2 = 0x1.715476p+1f,
|
||||
.c4 = 0x1.6b55a2p-10f,
|
||||
.ln2_hi = 0x1.62e4p-1f,
|
||||
.ln2_lo = 0x1.7f7d1cp-20f,
|
||||
};
|
||||
|
||||
// Return the ptr but hide it's value from the compiler so accesses
|
||||
// through it can't be optimised based on contents.
|
||||
template <typename T>
|
||||
inline const T* ptr_barrier(const T* ptr) {
|
||||
const T* opaque_ptr = ptr;
|
||||
__asm__("" : "+r"(opaque_ptr));
|
||||
return opaque_ptr;
|
||||
}
|
||||
|
||||
// Check whether any lanes in the mask are set
|
||||
inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
|
||||
|
||||
// e^2x - 1 inline helper
|
||||
inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
|
||||
float32x2_t ln2 = vld1_f32(&d->ln2_hi);
|
||||
float32x4_t lane_consts = vld1q_f32(&d->c1);
|
||||
|
||||
// Reduce argument: f in [-ln2/2, ln2/2], i is exact.
|
||||
float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
|
||||
int32x4_t i = vcvtq_s32_f32(j);
|
||||
float32x4_t f = vaddq_f32(x, x);
|
||||
f = vfmsq_lane_f32(f, j, ln2, 0);
|
||||
f = vfmsq_lane_f32(f, j, ln2, 1);
|
||||
|
||||
// Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
|
||||
float32x4_t f2 = vmulq_f32(f, f);
|
||||
float32x4_t f4 = vmulq_f32(f2, f2);
|
||||
float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
|
||||
float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
|
||||
float32x4_t poly = vfmaq_f32(p01, f2, p23);
|
||||
poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
|
||||
poly = vfmaq_f32(f, f2, poly);
|
||||
|
||||
// scale = 2^i
|
||||
int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
|
||||
float32x4_t scale = vreinterpretq_f32_s32(u);
|
||||
return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
|
||||
}
|
||||
|
||||
// Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
|
||||
inline float32x4_t special_case(float32x4_t x, float32x4_t q,
|
||||
uint32x4_t special) {
|
||||
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
|
||||
|
||||
float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
|
||||
uint32x4_t ix = vreinterpretq_u32_f32(x);
|
||||
uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
|
||||
uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
|
||||
uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
|
||||
float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
|
||||
return vbslq_f32(special, special_y, y);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Implementation of tanhf adapted from Arm Optimized Routines (tanhf
|
||||
// AdvSIMD)
|
||||
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
|
||||
//
|
||||
// Approximation for single-precision vector tanh(x), using a simplified
|
||||
// version of expm1f. The maximum error is 2.08 + 0.5 ULP:
|
||||
// _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
|
||||
inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
|
||||
const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
|
||||
|
||||
// tanh(x) = (e^2x - 1) / (e^2x + 1)
|
||||
// q = e^2x -1
|
||||
float32x4_t q = e2xm1f_inline(x, d);
|
||||
|
||||
// Check for special cases
|
||||
uint32x4_t special = vcagtq_f32(x, d->special_bound);
|
||||
|
||||
// Fall back to vectorised special case for any lanes which would cause
|
||||
// expm1 to overflow
|
||||
if (any_u32(special)) {
|
||||
return special_case(x, q, special);
|
||||
}
|
||||
|
||||
// Complete fast path if no special lanes
|
||||
// tanh(x) = q / (q+2)
|
||||
return vdivq_f32(q, vaddq_f32(q, d->two));
|
||||
}
|
||||
|
||||
} // namespace vec_op
|
||||
|
||||
#endif // CPU_TANHF_NEON_HPP
|
||||
@@ -0,0 +1,44 @@
|
||||
#ifndef CPU_TYPES_HPP
|
||||
#define CPU_TYPES_HPP
|
||||
|
||||
#if defined(__x86_64__)
|
||||
// x86 implementation
|
||||
#include "cpu_types_x86.hpp"
|
||||
#elif defined(__powerpc__)
|
||||
// ppc implementation
|
||||
#include "cpu_types_vsx.hpp"
|
||||
#elif defined(__s390x__)
|
||||
// s390 implementation
|
||||
#include "cpu_types_vxe.hpp"
|
||||
#elif defined(__aarch64__)
|
||||
// arm implementation
|
||||
#include "cpu_types_arm.hpp"
|
||||
#elif defined(__riscv_v)
|
||||
// riscv implementation
|
||||
#include "cpu_types_riscv.hpp"
|
||||
#else
|
||||
#warning "unsupported vLLM cpu implementation, vLLM will compile with scalar"
|
||||
#include "cpu_types_scalar.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef _OPENMP
|
||||
#include <omp.h>
|
||||
#endif
|
||||
|
||||
#include <c10/util/Exception.h>
|
||||
|
||||
namespace cpu_utils {
|
||||
// Without OpenMP the omp pragmas compile to serial loops, so report 1: kernels
|
||||
// that barrier on the thread count would otherwise deadlock.
|
||||
inline int get_max_threads() {
|
||||
#ifdef _OPENMP
|
||||
return omp_get_max_threads();
|
||||
#else
|
||||
TORCH_WARN_ONCE(
|
||||
"vLLM CPU was built without OpenMP; running single-threaded.");
|
||||
return 1;
|
||||
#endif
|
||||
}
|
||||
} // namespace cpu_utils
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,987 @@
|
||||
#include <cmath>
|
||||
#include <type_traits>
|
||||
|
||||
#include <arm_neon.h>
|
||||
|
||||
#include "cpu/cpu_tanhf_neon.hpp"
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
|
||||
#if defined(__APPLE__)
|
||||
#include "omp.h"
|
||||
#endif
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
|
||||
#ifndef CPU_OP_GUARD
|
||||
#define CPU_KERNEL_GUARD_IN(NAME)
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME)
|
||||
#else
|
||||
#define CPU_KERNEL_GUARD_IN(NAME) \
|
||||
std::cout << #NAME << " invoked." << std::endl;
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME) \
|
||||
std::cout << #NAME << " exit." << std::endl;
|
||||
#endif
|
||||
|
||||
#define FORCE_INLINE __attribute__((always_inline)) inline
|
||||
// Number of elements in single ASIMD vector of given Datatype
|
||||
#define NUM_ELEMENTS_REG(vec) (sizeof(vec) / sizeof(vec[0]))
|
||||
|
||||
namespace {
|
||||
template <typename T, T... indexes, typename F>
|
||||
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
|
||||
(f(std::integral_constant<T, indexes>{}), ...);
|
||||
};
|
||||
}; // namespace
|
||||
|
||||
template <typename T, T count, typename F,
|
||||
typename = std::enable_if_t<std::is_invocable_v<F, T>>>
|
||||
inline constexpr void unroll_loop(F&& f) {
|
||||
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
|
||||
}
|
||||
|
||||
template <typename T, typename... Ts>
|
||||
struct is_one_of : std::bool_constant<(std::is_same_v<T, Ts> || ...)> {};
|
||||
|
||||
template <typename T, typename... Ts>
|
||||
inline constexpr bool is_one_of_v = is_one_of<T, Ts...>::value;
|
||||
|
||||
struct uninit_t {
|
||||
explicit constexpr uninit_t() = default;
|
||||
};
|
||||
inline constexpr uninit_t uninit{};
|
||||
|
||||
template <typename NxVectorizedTVecReg, typename T, int VEC_ELEM_NUM>
|
||||
union AliasReg {
|
||||
NxVectorizedTVecReg reg;
|
||||
T values[VEC_ELEM_NUM];
|
||||
};
|
||||
|
||||
// Template over at::vec::Vectorized<T> to support
|
||||
// multiple vectorised registers into 1 of length VEC_REG_NUM val
|
||||
template <int N, typename T>
|
||||
struct NxVectorizedTVecReg {
|
||||
using value_t = T;
|
||||
using VectorizedT = Vectorized<T>;
|
||||
|
||||
VectorizedT val[N];
|
||||
|
||||
NxVectorizedTVecReg() = default;
|
||||
NxVectorizedTVecReg(const NxVectorizedTVecReg&) = default;
|
||||
NxVectorizedTVecReg(NxVectorizedTVecReg&&) = default;
|
||||
NxVectorizedTVecReg& operator=(const NxVectorizedTVecReg&) = default;
|
||||
NxVectorizedTVecReg& operator=(NxVectorizedTVecReg&&) = default;
|
||||
|
||||
explicit NxVectorizedTVecReg(uninit_t) noexcept {};
|
||||
|
||||
FORCE_INLINE explicit NxVectorizedTVecReg(const VectorizedT& vec_t) {
|
||||
unroll_loop<int, N>([&](int i) { val[i] = vec_t; });
|
||||
};
|
||||
|
||||
FORCE_INLINE explicit NxVectorizedTVecReg(T v) noexcept {
|
||||
VectorizedT vv(v);
|
||||
unroll_loop<int, N>([&](int i) { val[i] = vv; });
|
||||
}
|
||||
|
||||
FORCE_INLINE explicit NxVectorizedTVecReg(const void* ptr) { load(ptr); }
|
||||
explicit NxVectorizedTVecReg(const void* ptr, const int elem_num) {
|
||||
load(ptr, elem_num);
|
||||
}
|
||||
|
||||
static constexpr int size() noexcept { return N * VectorizedT::size(); }
|
||||
|
||||
FORCE_INLINE void save(void* ptr) const {
|
||||
value_t* base = reinterpret_cast<value_t*>(ptr);
|
||||
unroll_loop<int, N>(
|
||||
[&](int i) { val[i].store(base + i * VectorizedT::size()); });
|
||||
}
|
||||
FORCE_INLINE void load(const void* ptr) {
|
||||
const value_t* base = reinterpret_cast<const value_t*>(ptr);
|
||||
unroll_loop<int, N>([&](int i) {
|
||||
val[i] = VectorizedT::loadu(base + i * VectorizedT::size());
|
||||
});
|
||||
}
|
||||
|
||||
FORCE_INLINE void save(void* ptr, const int elem_num) const {
|
||||
value_t* base = reinterpret_cast<value_t*>(ptr);
|
||||
save_partial(base, elem_num);
|
||||
}
|
||||
|
||||
FORCE_INLINE void load(const void* ptr, const int elem_num) {
|
||||
const value_t* base = reinterpret_cast<const value_t*>(ptr);
|
||||
load_partial(base, elem_num);
|
||||
}
|
||||
|
||||
FORCE_INLINE void save_partial(value_t* base, int elem_num) const {
|
||||
const int w = VectorizedT::size();
|
||||
int full = elem_num / w;
|
||||
int rem = elem_num % w;
|
||||
for (int i = 0; i < full; i++) val[i].store(base + i * w);
|
||||
if (rem) val[full].store(base + full * w, rem);
|
||||
}
|
||||
|
||||
FORCE_INLINE void load_partial(const value_t* base, int elem_num) {
|
||||
const int w = VectorizedT::size();
|
||||
int full = elem_num / w;
|
||||
int rem = elem_num % w;
|
||||
for (int i = 0; i < full; i++) val[i] = VectorizedT::loadu(base + i * w);
|
||||
if (rem) val[full] = VectorizedT::loadu(base + full * w, rem);
|
||||
}
|
||||
|
||||
template <VectorizedT (VectorizedT::*torch_vec_func)() const,
|
||||
value_t (*std_func)(value_t)>
|
||||
FORCE_INLINE NxVectorizedTVecReg opt_vec_func_impl() const {
|
||||
NxVectorizedTVecReg result;
|
||||
|
||||
if constexpr (torch_vec_func != nullptr) {
|
||||
unroll_loop<int, N>(
|
||||
[&](int i) { result.val[i] = (val[i].*torch_vec_func)(); });
|
||||
} else {
|
||||
for (int i = 0; i < N; i++) {
|
||||
alignas(64) value_t buf[VectorizedT::size()];
|
||||
val[i].store(buf);
|
||||
for (int j = 0; j < VectorizedT::size(); ++j) {
|
||||
buf[j] = std_func(buf[j]);
|
||||
}
|
||||
result.val[i] = VectorizedT::loadu(buf);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
template <typename DerivedClassT, int N, typename T>
|
||||
struct VectorizedRegWrapper {
|
||||
using ScalarT = T;
|
||||
using VectorizedT = Vectorized<T>;
|
||||
using NxVectorizedTArray = NxVectorizedTVecReg<N, T>;
|
||||
|
||||
constexpr static int VEC_REG_NUM = N;
|
||||
constexpr static int VEC_ELEM_NUM = VEC_REG_NUM * VectorizedT::size();
|
||||
constexpr static int get_elem_num() { return VEC_ELEM_NUM; };
|
||||
|
||||
NxVectorizedTArray reg;
|
||||
|
||||
VectorizedRegWrapper() noexcept = default;
|
||||
explicit VectorizedRegWrapper(uninit_t) noexcept : reg{uninit} {};
|
||||
explicit VectorizedRegWrapper(T v) : reg(v) {};
|
||||
explicit VectorizedRegWrapper(const void* ptr) : reg(ptr) {};
|
||||
explicit VectorizedRegWrapper(const void* ptr, const int elem_num)
|
||||
: reg(ptr, elem_num) {};
|
||||
explicit VectorizedRegWrapper(const VectorizedT& r) : reg(r) {};
|
||||
explicit VectorizedRegWrapper(const NxVectorizedTArray& r) : reg(r) {};
|
||||
|
||||
VectorizedRegWrapper(const VectorizedRegWrapper&) = default;
|
||||
VectorizedRegWrapper(VectorizedRegWrapper&&) = default;
|
||||
VectorizedRegWrapper& operator=(VectorizedRegWrapper&&) = default;
|
||||
VectorizedRegWrapper& operator=(const VectorizedRegWrapper&) = default;
|
||||
|
||||
FORCE_INLINE void save(void* ptr) const { reg.save(ptr); }
|
||||
void save(void* ptr, const int elem_num) const { reg.save(ptr, elem_num); }
|
||||
|
||||
// Define optimized functions using at::vec::Vectorized<T> where possible
|
||||
// Fallback to std:: functions when not available
|
||||
#define OPT_TORCH_IMPL(FUNC_NAME, STD_FUNC_NAME, TORCH_FUNC_NAME, ...) \
|
||||
FORCE_INLINE DerivedClassT FUNC_NAME() const { \
|
||||
if constexpr (is_one_of_v<T, __VA_ARGS__>) { \
|
||||
return DerivedClassT{ \
|
||||
reg.template opt_vec_func_impl<&VectorizedT::TORCH_FUNC_NAME, \
|
||||
std::STD_FUNC_NAME>()}; \
|
||||
} else { \
|
||||
return DerivedClassT{reg.template opt_vec_func_impl< \
|
||||
nullptr, static_cast<ScalarT (*)(ScalarT)>(&std::STD_FUNC_NAME)>()}; \
|
||||
} \
|
||||
}
|
||||
|
||||
// Define optimized functions for datatypes passed in __VA_ARGS__
|
||||
OPT_TORCH_IMPL(abs, abs, abs, c10::Half, float)
|
||||
OPT_TORCH_IMPL(er, erf, erf, float)
|
||||
OPT_TORCH_IMPL(exp, exp, fexp_u20, float)
|
||||
OPT_TORCH_IMPL(exp_u20, exp, exp_u20, float)
|
||||
OPT_TORCH_IMPL(sin, sin, sin, float)
|
||||
OPT_TORCH_IMPL(sinh, sinh, sinh, float)
|
||||
OPT_TORCH_IMPL(cos, cos, cos, float)
|
||||
OPT_TORCH_IMPL(cosh, cosh, cosh, float)
|
||||
OPT_TORCH_IMPL(log, log, log, float)
|
||||
OPT_TORCH_IMPL(log10, log10, log10, float)
|
||||
OPT_TORCH_IMPL(sqrt, sqrt, sqrt, c10::Half, float)
|
||||
OPT_TORCH_IMPL(tan, tan, tan, float)
|
||||
OPT_TORCH_IMPL(tanh, tanh, tanh, float)
|
||||
|
||||
#undef OPT_TORCH_IMPL
|
||||
};
|
||||
|
||||
// forward declare vectorised dtypes
|
||||
struct FP32Vec8;
|
||||
struct FP32Vec16;
|
||||
struct FP16Vec8;
|
||||
struct FP16Vec16;
|
||||
struct BF16Vec8;
|
||||
struct BF16Vec16;
|
||||
|
||||
struct INT8Vec16;
|
||||
struct INT32Vec16;
|
||||
|
||||
template <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using vec_t = typename VecType<T>::vec_type;
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using vec_type = FP32Vec8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using vec_type = FP16Vec8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using vec_type = BF16Vec8;
|
||||
};
|
||||
|
||||
struct FP16Vec8 : public VectorizedRegWrapper<FP16Vec8, 1, c10::Half> {
|
||||
using Base = VectorizedRegWrapper<FP16Vec8, 1, c10::Half>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
explicit FP16Vec8(const FP32Vec8&);
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public VectorizedRegWrapper<FP16Vec16, 2, c10::Half> {
|
||||
using Base = VectorizedRegWrapper<FP16Vec16, 2, c10::Half>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
// ASIMD does not support non-temporal loads
|
||||
explicit FP16Vec16(bool, const void* ptr) : Base(ptr) {}
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16& vec);
|
||||
};
|
||||
|
||||
struct BF16Vec8 : public VectorizedRegWrapper<BF16Vec8, 1, c10::BFloat16> {
|
||||
using Base = VectorizedRegWrapper<BF16Vec8, 1, c10::BFloat16>;
|
||||
using VectorizedT = typename Base::VectorizedT;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
explicit BF16Vec8(at_bfloat16x8_t data) : Base(VectorizedT(data)) {};
|
||||
|
||||
explicit BF16Vec8(float32x4x2_t v) {
|
||||
reg.val[0] = convert_float_bfloat16(v.val[0], v.val[1]);
|
||||
};
|
||||
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public VectorizedRegWrapper<BF16Vec16, 2, c10::BFloat16> {
|
||||
using Base = VectorizedRegWrapper<BF16Vec16, 2, c10::BFloat16>;
|
||||
using VectorizedT = typename Base::VectorizedT;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
// ASIMD does not support non-temporal loads
|
||||
explicit BF16Vec16(bool, const void* ptr) : Base(ptr) {}
|
||||
|
||||
explicit BF16Vec16(float32x4x4_t v) {
|
||||
reg.val[0] = convert_float_bfloat16(v.val[0], v.val[1]);
|
||||
reg.val[1] = convert_float_bfloat16(v.val[2], v.val[3]);
|
||||
};
|
||||
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
|
||||
using Base = VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& vec8_data) {
|
||||
reg.val[0] = vec8_data.reg.val[0];
|
||||
reg.val[1] = vec8_data.reg.val[0];
|
||||
reg.val[2] = vec8_data.reg.val[0];
|
||||
reg.val[3] = vec8_data.reg.val[0];
|
||||
};
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
|
||||
};
|
||||
|
||||
struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
|
||||
using Base = VectorizedRegWrapper<FP32Vec4, 1, float>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
using VectorizedT = typename Base::VectorizedT;
|
||||
using Vectorized1x4f = typename Base::NxVectorizedTArray;
|
||||
|
||||
FP32Vec4() : Base() {};
|
||||
explicit FP32Vec4(float v) : Base(v) {};
|
||||
|
||||
explicit FP32Vec4(float32x4_t data) : Base(VectorizedT(data)) {};
|
||||
|
||||
explicit FP32Vec4(const FP32Vec4& data) : Base(data) {};
|
||||
|
||||
FORCE_INLINE FP32Vec4 tanh() const {
|
||||
return FP32Vec4(fast_tanhf_f32x4(reg.val[0]));
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public VectorizedRegWrapper<FP32Vec8, 2, float> {
|
||||
using Base = VectorizedRegWrapper<FP32Vec8, 2, float>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
using Base::VEC_REG_NUM;
|
||||
|
||||
using VectorizedT = typename Base::VectorizedT;
|
||||
using Vectorized2x4f = typename Base::NxVectorizedTArray;
|
||||
|
||||
FP32Vec8() : Base() {};
|
||||
FP32Vec8(const FP32Vec8& data) : Base(data) {};
|
||||
|
||||
explicit FP32Vec8(float v) : Base(v) {};
|
||||
explicit FP32Vec8(const float* ptr)
|
||||
: Base(reinterpret_cast<const void*>(ptr)) {};
|
||||
explicit FP32Vec8(const float* ptr, const int elem_num)
|
||||
: Base(reinterpret_cast<const void*>(ptr), elem_num) {};
|
||||
|
||||
explicit FP32Vec8(const Vectorized2x4f& data) {
|
||||
reg.val[0] = data.val[0];
|
||||
reg.val[1] = data.val[1];
|
||||
};
|
||||
|
||||
explicit FP32Vec8(const BF16Vec8& v) {
|
||||
std::tie(reg.val[0], reg.val[1]) = convert_bfloat16_float(v.reg.val[0]);
|
||||
};
|
||||
explicit FP32Vec8(const FP16Vec8& v) {
|
||||
reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
|
||||
reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
|
||||
};
|
||||
explicit FP32Vec8(float16x8_t v) {
|
||||
reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v)));
|
||||
reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v)));
|
||||
};
|
||||
explicit FP32Vec8(at_bfloat16x8_t v) {
|
||||
std::tie(reg.val[0], reg.val[1]) =
|
||||
convert_bfloat16_float(Vectorized<c10::BFloat16>(v));
|
||||
};
|
||||
explicit FP32Vec8(float32x4x2_t data) {
|
||||
reg.val[0] = Vectorized<float>(data.val[0]);
|
||||
reg.val[1] = Vectorized<float>(data.val[1]);
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 tanh() const {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
|
||||
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE float reduce_sum() const noexcept {
|
||||
float answer = 0;
|
||||
std::plus<VectorizedT> add;
|
||||
|
||||
unroll_loop<int, VEC_REG_NUM>([&](int i) {
|
||||
answer += at::vec::vec_reduce_all<float, std::plus<VectorizedT>>(
|
||||
add, reg.val[i]);
|
||||
});
|
||||
return answer;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 operator+(const FP32Vec8& b) const noexcept {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] + b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] + b.reg.val[1];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 operator-(const FP32Vec8& b) const noexcept {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] - b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] - b.reg.val[1];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 operator*(const FP32Vec8& b) const noexcept {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] * b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] * b.reg.val[1];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec8 operator/(const FP32Vec8& b) const noexcept {
|
||||
FP32Vec8 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] / b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] / b.reg.val[1];
|
||||
return r;
|
||||
}
|
||||
};
|
||||
|
||||
struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
|
||||
using Base = VectorizedRegWrapper<FP32Vec16, 4, float>;
|
||||
using Base::Base;
|
||||
using Base::get_elem_num;
|
||||
using Base::VEC_ELEM_NUM;
|
||||
|
||||
using ScalarT = typename Base::ScalarT;
|
||||
using VectorizedT = typename Base::VectorizedT;
|
||||
using Vectorized4x4f = typename Base::NxVectorizedTArray;
|
||||
|
||||
FP32Vec16() : Base() {};
|
||||
FP32Vec16(const FP32Vec16& data) : Base(data) {};
|
||||
explicit FP32Vec16(float v) : Base(v) {};
|
||||
explicit FP32Vec16(const float* ptr)
|
||||
: Base(reinterpret_cast<const void*>(ptr)) {};
|
||||
explicit FP32Vec16(const float* ptr, const int elem_num)
|
||||
: Base(reinterpret_cast<const void*>(ptr), elem_num) {};
|
||||
explicit FP32Vec16(const Vectorized4x4f& data) {
|
||||
reg.val[0] = data.val[0];
|
||||
reg.val[1] = data.val[1];
|
||||
reg.val[2] = data.val[2];
|
||||
reg.val[3] = data.val[3];
|
||||
};
|
||||
|
||||
// ASIMD does not support non-temporal loads
|
||||
explicit FP32Vec16(bool, const float* ptr) : Base(ptr) {}
|
||||
|
||||
explicit FP32Vec16(float32x4x4_t data) {
|
||||
reg.val[0] = data.val[0];
|
||||
reg.val[1] = data.val[1];
|
||||
reg.val[2] = data.val[2];
|
||||
reg.val[3] = data.val[3];
|
||||
};
|
||||
|
||||
explicit FP32Vec16(const FP32Vec4& data) {
|
||||
reg.val[0] = data.reg.val[0];
|
||||
reg.val[1] = data.reg.val[0];
|
||||
reg.val[2] = data.reg.val[0];
|
||||
reg.val[3] = data.reg.val[0];
|
||||
};
|
||||
|
||||
explicit FP32Vec16(const FP32Vec8& data) {
|
||||
reg.val[0] = data.reg.val[0];
|
||||
reg.val[1] = data.reg.val[1];
|
||||
reg.val[2] = data.reg.val[0];
|
||||
reg.val[3] = data.reg.val[1];
|
||||
};
|
||||
|
||||
explicit FP32Vec16(const BF16Vec16& v) {
|
||||
std::tie(reg.val[0], reg.val[1]) = convert_bfloat16_float(v.reg.val[0]);
|
||||
std::tie(reg.val[2], reg.val[3]) = convert_bfloat16_float(v.reg.val[1]);
|
||||
};
|
||||
|
||||
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
|
||||
|
||||
// FP8 stub: dead code on ARM (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : Base() {}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v) {
|
||||
reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
|
||||
reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
|
||||
reg.val[2] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[1])));
|
||||
reg.val[3] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[1])));
|
||||
};
|
||||
|
||||
FORCE_INLINE FP32Vec16 tanh() const {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = Vectorized<float>(fast_tanhf_f32x4(reg.val[0]));
|
||||
r.reg.val[1] = Vectorized<float>(fast_tanhf_f32x4(reg.val[1]));
|
||||
r.reg.val[2] = Vectorized<float>(fast_tanhf_f32x4(reg.val[2]));
|
||||
r.reg.val[3] = Vectorized<float>(fast_tanhf_f32x4(reg.val[3]));
|
||||
return r;
|
||||
}
|
||||
|
||||
static FORCE_INLINE void load_even_odd(const float* ptr, FP32Vec16& even,
|
||||
FP32Vec16& odd) noexcept {
|
||||
const float32x4x2_t x01 = vuzpq_f32(vld1q_f32(ptr), vld1q_f32(ptr + 4));
|
||||
const float32x4x2_t x23 =
|
||||
vuzpq_f32(vld1q_f32(ptr + 8), vld1q_f32(ptr + 12));
|
||||
const float32x4x2_t x45 =
|
||||
vuzpq_f32(vld1q_f32(ptr + 16), vld1q_f32(ptr + 20));
|
||||
const float32x4x2_t x67 =
|
||||
vuzpq_f32(vld1q_f32(ptr + 24), vld1q_f32(ptr + 28));
|
||||
|
||||
even.reg.val[0] = VectorizedT(x01.val[0]);
|
||||
even.reg.val[1] = VectorizedT(x23.val[0]);
|
||||
even.reg.val[2] = VectorizedT(x45.val[0]);
|
||||
even.reg.val[3] = VectorizedT(x67.val[0]);
|
||||
odd.reg.val[0] = VectorizedT(x01.val[1]);
|
||||
odd.reg.val[1] = VectorizedT(x23.val[1]);
|
||||
odd.reg.val[2] = VectorizedT(x45.val[1]);
|
||||
odd.reg.val[3] = VectorizedT(x67.val[1]);
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 operator+(const FP32Vec16& b) const noexcept {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] + b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] + b.reg.val[1];
|
||||
r.reg.val[2] = reg.val[2] + b.reg.val[2];
|
||||
r.reg.val[3] = reg.val[3] + b.reg.val[3];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 operator-(const FP32Vec16& b) const noexcept {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] - b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] - b.reg.val[1];
|
||||
r.reg.val[2] = reg.val[2] - b.reg.val[2];
|
||||
r.reg.val[3] = reg.val[3] - b.reg.val[3];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 operator-() const noexcept {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = reg.val[0].neg();
|
||||
r.reg.val[1] = reg.val[1].neg();
|
||||
r.reg.val[2] = reg.val[2].neg();
|
||||
r.reg.val[3] = reg.val[3].neg();
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 operator*(const FP32Vec16& b) const noexcept {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] * b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] * b.reg.val[1];
|
||||
r.reg.val[2] = reg.val[2] * b.reg.val[2];
|
||||
r.reg.val[3] = reg.val[3] * b.reg.val[3];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 operator/(const FP32Vec16& b) const noexcept {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = reg.val[0] / b.reg.val[0];
|
||||
r.reg.val[1] = reg.val[1] / b.reg.val[1];
|
||||
r.reg.val[2] = reg.val[2] / b.reg.val[2];
|
||||
r.reg.val[3] = reg.val[3] / b.reg.val[3];
|
||||
return r;
|
||||
}
|
||||
|
||||
FORCE_INLINE FP32Vec16 clamp(const FP32Vec16& min,
|
||||
const FP32Vec16& max) const {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = at::vec::clamp(reg.val[0], min.reg.val[0], max.reg.val[0]);
|
||||
r.reg.val[1] = at::vec::clamp(reg.val[1], min.reg.val[1], max.reg.val[1]);
|
||||
r.reg.val[2] = at::vec::clamp(reg.val[2], min.reg.val[2], max.reg.val[2]);
|
||||
r.reg.val[3] = at::vec::clamp(reg.val[3], min.reg.val[3], max.reg.val[3]);
|
||||
return r;
|
||||
};
|
||||
|
||||
FORCE_INLINE FP32Vec16 min(const FP32Vec16& b) const {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = minimum(b.reg.val[0], reg.val[0]),
|
||||
r.reg.val[1] = minimum(b.reg.val[1], reg.val[1]);
|
||||
r.reg.val[2] = minimum(b.reg.val[2], reg.val[2]);
|
||||
r.reg.val[3] = minimum(b.reg.val[3], reg.val[3]);
|
||||
return r;
|
||||
};
|
||||
|
||||
FORCE_INLINE FP32Vec16 max(const FP32Vec16& b) const {
|
||||
FP32Vec16 r(uninit);
|
||||
r.reg.val[0] = maximum(b.reg.val[0], reg.val[0]);
|
||||
r.reg.val[1] = maximum(b.reg.val[1], reg.val[1]);
|
||||
r.reg.val[2] = maximum(b.reg.val[2], reg.val[2]);
|
||||
r.reg.val[3] = maximum(b.reg.val[3], reg.val[3]);
|
||||
return r;
|
||||
};
|
||||
|
||||
FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
|
||||
size_t num_elements = reg.val[0].size();
|
||||
|
||||
if (elem_num == VEC_ELEM_NUM) {
|
||||
return FP32Vec16::min(b);
|
||||
}
|
||||
|
||||
int full_blocks = elem_num / num_elements;
|
||||
const int remainder = elem_num % num_elements;
|
||||
|
||||
FP32Vec16 res(uninit);
|
||||
for (int i = 0; i < full_blocks; i++)
|
||||
res.reg.val[i] = minimum(b.reg.val[i], reg.val[i]);
|
||||
|
||||
if (remainder > 0) {
|
||||
float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 0),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 0));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(min_v, res.reg.val[full_blocks], 0);
|
||||
}
|
||||
if (remainder > 1) {
|
||||
float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 1),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 1));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(min_v, res.reg.val[full_blocks], 1);
|
||||
}
|
||||
if (remainder > 2) {
|
||||
float min_v = std::min(vgetq_lane_f32(reg.val[full_blocks], 2),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 2));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(min_v, res.reg.val[full_blocks], 2);
|
||||
}
|
||||
|
||||
return res;
|
||||
};
|
||||
|
||||
FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
|
||||
size_t num_elements = reg.val[0].size();
|
||||
|
||||
if (elem_num == VEC_ELEM_NUM) {
|
||||
return FP32Vec16::max(b);
|
||||
}
|
||||
|
||||
int full_blocks = elem_num / num_elements;
|
||||
int remainder = elem_num % num_elements;
|
||||
|
||||
FP32Vec16 res(uninit);
|
||||
|
||||
for (int i = 0; i < full_blocks; i++)
|
||||
res.reg.val[i] = maximum(b.reg.val[i], reg.val[i]);
|
||||
|
||||
if (remainder > 0) {
|
||||
float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 0),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 0));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(max_v, res.reg.val[full_blocks], 0);
|
||||
}
|
||||
if (remainder > 1) {
|
||||
float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 1),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 1));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(max_v, res.reg.val[full_blocks], 1);
|
||||
}
|
||||
if (remainder > 2) {
|
||||
float max_v = std::max(vgetq_lane_f32(reg.val[full_blocks], 2),
|
||||
vgetq_lane_f32(b.reg.val[full_blocks], 2));
|
||||
res.reg.val[full_blocks] =
|
||||
vsetq_lane_f32(max_v, res.reg.val[full_blocks], 2);
|
||||
}
|
||||
return res;
|
||||
};
|
||||
|
||||
float reduce_max() const {
|
||||
VectorizedT max_vec = reg.val[0];
|
||||
unroll_loop<int, VEC_REG_NUM>([&](int i) {
|
||||
if (i > 0) max_vec = maximum(max_vec, reg.val[i]);
|
||||
});
|
||||
|
||||
return vmaxvq_f32(max_vec);
|
||||
}
|
||||
|
||||
float reduce_min() const {
|
||||
VectorizedT min_vec = reg.val[0];
|
||||
unroll_loop<int, VEC_REG_NUM>([&](int i) {
|
||||
if (i > 0) min_vec = minimum(min_vec, reg.val[i]);
|
||||
});
|
||||
|
||||
return vminvq_f32(min_vec);
|
||||
}
|
||||
|
||||
template <int group_size>
|
||||
float reduce_sub_sum(int idx) {
|
||||
static_assert(VEC_ELEM_NUM % group_size == 0);
|
||||
|
||||
AliasReg<NxVectorizedTArray, ScalarT, VEC_ELEM_NUM> ar{reg};
|
||||
float answer = 0;
|
||||
const int start = idx * group_size;
|
||||
unroll_loop<int, group_size>(
|
||||
[&](int i) { answer += ar.values[start + i]; });
|
||||
|
||||
return answer;
|
||||
};
|
||||
|
||||
float reduce_sum() const {
|
||||
float answer = 0;
|
||||
std::plus<VectorizedT> add;
|
||||
unroll_loop<int, VEC_REG_NUM>([&](int i) {
|
||||
answer += at::vec::vec_reduce_all<float>(add, reg.val[i]);
|
||||
});
|
||||
|
||||
return answer;
|
||||
}
|
||||
};
|
||||
|
||||
// Only used for int types for now could be replaced when
|
||||
// int8/32 vectorised ops are added in ATen
|
||||
template <typename T>
|
||||
struct Vec {
|
||||
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; };
|
||||
};
|
||||
|
||||
struct INT8Vec16 : public Vec<INT8Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
union AliasReg {
|
||||
int8x16_t reg;
|
||||
int8_t values[VEC_ELEM_NUM];
|
||||
};
|
||||
int8x16_t reg;
|
||||
|
||||
explicit INT8Vec16(const FP32Vec16& vec) {
|
||||
// Convert each 128-bit float32 vector to int32
|
||||
int32x4_t part0 =
|
||||
vcvtq_s32_f32(vec.reg.val[0]); // Convert first 128-bit block
|
||||
int32x4_t part1 =
|
||||
vcvtq_s32_f32(vec.reg.val[1]); // Convert second 128-bit block
|
||||
int32x4_t part2 =
|
||||
vcvtq_s32_f32(vec.reg.val[2]); // Convert third 128-bit block
|
||||
int32x4_t part3 =
|
||||
vcvtq_s32_f32(vec.reg.val[3]); // Convert fourth 128-bit block
|
||||
|
||||
// Narrow each 32-bit vector to 8 bits and combine
|
||||
int8x8_t lower =
|
||||
vqmovn_s16(vcombine_s16(vqmovn_s32(part0), vqmovn_s32(part1)));
|
||||
int8x8_t upper =
|
||||
vqmovn_s16(vcombine_s16(vqmovn_s32(part2), vqmovn_s32(part3)));
|
||||
reg = vcombine_s8(lower, upper); // Combine to form a single 128-bit vector
|
||||
}
|
||||
|
||||
void save(int8_t* ptr) const { vst1q_s8(ptr, reg); };
|
||||
|
||||
void save(int8_t* ptr, const int elem_num) const {
|
||||
int full_blocks = elem_num / NUM_ELEMENTS_REG(reg);
|
||||
int remainder = elem_num % NUM_ELEMENTS_REG(reg);
|
||||
|
||||
for (int i = 0; i < full_blocks; i++)
|
||||
vst1q_s8(reinterpret_cast<int8_t*>(ptr) + NUM_ELEMENTS_REG(reg) * i, reg);
|
||||
if (remainder > 0) {
|
||||
int8x16_t temp = reg;
|
||||
int8_t* base =
|
||||
reinterpret_cast<int8_t*>(ptr) + full_blocks * NUM_ELEMENTS_REG(reg);
|
||||
if (remainder > 0) base[0] = vgetq_lane_s8(temp, 0);
|
||||
if (remainder > 1) base[1] = vgetq_lane_s8(temp, 1);
|
||||
if (remainder > 2) base[2] = vgetq_lane_s8(temp, 2);
|
||||
if (remainder > 3) base[3] = vgetq_lane_s8(temp, 3);
|
||||
if (remainder > 4) base[4] = vgetq_lane_s8(temp, 4);
|
||||
if (remainder > 5) base[5] = vgetq_lane_s8(temp, 5);
|
||||
if (remainder > 6) base[6] = vgetq_lane_s8(temp, 6);
|
||||
if (remainder > 7) base[7] = vgetq_lane_s8(temp, 7);
|
||||
if (remainder > 8) base[8] = vgetq_lane_s8(temp, 8);
|
||||
if (remainder > 9) base[9] = vgetq_lane_s8(temp, 9);
|
||||
if (remainder > 10) base[10] = vgetq_lane_s8(temp, 10);
|
||||
if (remainder > 11) base[11] = vgetq_lane_s8(temp, 11);
|
||||
if (remainder > 12) base[12] = vgetq_lane_s8(temp, 12);
|
||||
if (remainder > 13) base[13] = vgetq_lane_s8(temp, 13);
|
||||
if (remainder > 14) base[14] = vgetq_lane_s8(temp, 14);
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
struct INT8Vec64 : public Vec<INT8Vec64> {
|
||||
constexpr static int VEC_ELEM_NUM = 64;
|
||||
union AliasReg {
|
||||
int8x16x4_t reg;
|
||||
int8_t values[VEC_ELEM_NUM];
|
||||
};
|
||||
int8x16x4_t reg;
|
||||
|
||||
explicit INT8Vec64(const int8_t* ptr) { reg = vld1q_s8_x4(ptr); }
|
||||
|
||||
// ASIMD does not support non-temporal loads
|
||||
explicit INT8Vec64(bool, const int8_t* ptr) : INT8Vec64(ptr) {}
|
||||
|
||||
void save(int8_t* ptr) const { vst1q_s8_x4(ptr, reg); }
|
||||
|
||||
// masked store
|
||||
void save(int8_t* p, int elem_num) const {
|
||||
TORCH_CHECK(elem_num <= VEC_ELEM_NUM && elem_num > 0);
|
||||
|
||||
if (elem_num == VEC_ELEM_NUM) {
|
||||
vst1q_s8_x4(p, reg);
|
||||
return;
|
||||
}
|
||||
|
||||
const int full_quadwords = elem_num / 16;
|
||||
const int remaining_bytes = elem_num % 16;
|
||||
|
||||
for (int i = 0; i < full_quadwords; ++i) {
|
||||
vst1q_s8(p + 16 * i, reg.val[i]);
|
||||
}
|
||||
|
||||
if (remaining_bytes) {
|
||||
const int8x16_t v = reg.val[full_quadwords];
|
||||
int8_t* tail = p + 16 * full_quadwords;
|
||||
switch (remaining_bytes) {
|
||||
case 15:
|
||||
tail[14] = vgetq_lane_s8(v, 14);
|
||||
[[fallthrough]];
|
||||
case 14:
|
||||
tail[13] = vgetq_lane_s8(v, 13);
|
||||
[[fallthrough]];
|
||||
case 13:
|
||||
tail[12] = vgetq_lane_s8(v, 12);
|
||||
[[fallthrough]];
|
||||
case 12:
|
||||
tail[11] = vgetq_lane_s8(v, 11);
|
||||
[[fallthrough]];
|
||||
case 11:
|
||||
tail[10] = vgetq_lane_s8(v, 10);
|
||||
[[fallthrough]];
|
||||
case 10:
|
||||
tail[9] = vgetq_lane_s8(v, 9);
|
||||
[[fallthrough]];
|
||||
case 9:
|
||||
tail[8] = vgetq_lane_s8(v, 8);
|
||||
[[fallthrough]];
|
||||
case 8:
|
||||
tail[7] = vgetq_lane_s8(v, 7);
|
||||
[[fallthrough]];
|
||||
case 7:
|
||||
tail[6] = vgetq_lane_s8(v, 6);
|
||||
[[fallthrough]];
|
||||
case 6:
|
||||
tail[5] = vgetq_lane_s8(v, 5);
|
||||
[[fallthrough]];
|
||||
case 5:
|
||||
tail[4] = vgetq_lane_s8(v, 4);
|
||||
[[fallthrough]];
|
||||
case 4:
|
||||
tail[3] = vgetq_lane_s8(v, 3);
|
||||
[[fallthrough]];
|
||||
case 3:
|
||||
tail[2] = vgetq_lane_s8(v, 2);
|
||||
[[fallthrough]];
|
||||
case 2:
|
||||
tail[1] = vgetq_lane_s8(v, 1);
|
||||
[[fallthrough]];
|
||||
case 1:
|
||||
tail[0] = vgetq_lane_s8(v, 0);
|
||||
break;
|
||||
default:
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ASIMD does not support non-temporal stores
|
||||
void nt_save(int8_t* ptr) const { save(ptr); }
|
||||
}; // INT8Vec64
|
||||
|
||||
struct INT32Vec16 : public Vec<INT32Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
union AliasReg {
|
||||
int32x4x4_t reg;
|
||||
int32_t values[VEC_ELEM_NUM];
|
||||
};
|
||||
int32x4x4_t reg;
|
||||
|
||||
explicit INT32Vec16(const void* ptr) {
|
||||
reg.val[0] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr));
|
||||
reg.val[1] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 4);
|
||||
reg.val[2] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 8);
|
||||
reg.val[3] = vld1q_s32(reinterpret_cast<const int32_t*>(ptr) + 12);
|
||||
}
|
||||
|
||||
void save(int32_t* ptr) const {
|
||||
vst1q_s32(ptr, reg.val[0]);
|
||||
vst1q_s32(ptr + 4, reg.val[1]);
|
||||
vst1q_s32(ptr + 8, reg.val[2]);
|
||||
vst1q_s32(ptr + 12, reg.val[3]);
|
||||
};
|
||||
|
||||
void save(int32_t* ptr, const int elem_num) const {
|
||||
int full_blocks = elem_num / NUM_ELEMENTS_REG(reg.val[0]);
|
||||
int remainder = elem_num % NUM_ELEMENTS_REG(reg.val[0]);
|
||||
|
||||
for (int i = 0; i < full_blocks; i++)
|
||||
vst1q_s32(
|
||||
reinterpret_cast<__int32_t*>(ptr) + NUM_ELEMENTS_REG(reg.val[0]) * i,
|
||||
reg.val[i]);
|
||||
|
||||
if (remainder > 0) {
|
||||
int32x4_t temp = reg.val[full_blocks];
|
||||
int32_t* base = reinterpret_cast<int32_t*>(ptr) + full_blocks * 4;
|
||||
if (remainder > 0) base[0] = vgetq_lane_s32(temp, 0);
|
||||
if (remainder > 1) base[1] = vgetq_lane_s32(temp, 1);
|
||||
if (remainder > 2) base[2] = vgetq_lane_s32(temp, 2);
|
||||
if (remainder > 3) base[3] = vgetq_lane_s32(temp, 3);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void storeFP32(float v, T* ptr) {
|
||||
*ptr = v;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
|
||||
*reinterpret_cast<__fp16*>(ptr) = v;
|
||||
}
|
||||
|
||||
inline FP16Vec8::FP16Vec8(const FP32Vec8& v) {
|
||||
reg.val[0] = convert_float_half(v.reg.val[0], v.reg.val[1]);
|
||||
};
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
reg.val[0] = convert_float_half(v.reg.val[0], v.reg.val[1]);
|
||||
reg.val[1] = convert_float_half(v.reg.val[2], v.reg.val[3]);
|
||||
};
|
||||
|
||||
inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
|
||||
fmadd(acc.reg.val[0], a.reg.val[0], b.reg.val[0]);
|
||||
fmadd(acc.reg.val[1], a.reg.val[1], b.reg.val[1]);
|
||||
fmadd(acc.reg.val[2], a.reg.val[2], b.reg.val[2]);
|
||||
fmadd(acc.reg.val[3], a.reg.val[3], b.reg.val[3]);
|
||||
};
|
||||
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
|
||||
reg.val[0] = convert_float_bfloat16(v.reg.val[0], v.reg.val[1]);
|
||||
};
|
||||
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
|
||||
reg.val[0] = convert_float_bfloat16(v.reg.val[0], v.reg.val[1]);
|
||||
reg.val[1] = convert_float_bfloat16(v.reg.val[2], v.reg.val[3]);
|
||||
};
|
||||
|
||||
inline void fma(FP32Vec16& acc, BF16Vec32& a, BF16Vec32& b) {
|
||||
Vectorized<float> a0_low, a0_high, a1_low, a1_high, b0_low, b0_high, b1_low,
|
||||
b1_high;
|
||||
|
||||
std::tie(a0_low, a0_high) = convert_bfloat16_float(a.reg.val[0]);
|
||||
std::tie(a1_low, a1_high) = convert_bfloat16_float(a.reg.val[1]);
|
||||
std::tie(b0_low, b0_high) = convert_bfloat16_float(b.reg.val[0]);
|
||||
std::tie(b1_low, b1_high) = convert_bfloat16_float(b.reg.val[1]);
|
||||
|
||||
fmadd(acc.reg.val[0], a0_low, b0_low);
|
||||
fmadd(acc.reg.val[1], a0_high, b0_high);
|
||||
fmadd(acc.reg.val[2], a1_low, b1_low);
|
||||
fmadd(acc.reg.val[3], a1_high, b1_high);
|
||||
};
|
||||
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
#ifdef ARM_BF16_SUPPORT
|
||||
*reinterpret_cast<__bf16*>(ptr) = vcvth_bf16_f32(v);
|
||||
#else
|
||||
*ptr = static_cast<c10::BFloat16>(v);
|
||||
#endif
|
||||
};
|
||||
|
||||
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 1); };
|
||||
|
||||
}; // namespace vec_op
|
||||
@@ -0,0 +1,25 @@
|
||||
#ifndef CPU_TYPES_RISCV_HPP
|
||||
#define CPU_TYPES_RISCV_HPP
|
||||
|
||||
// RISC-V Vector (RVV) CPU type definitions for vLLM.
|
||||
//
|
||||
// Supports multiple VLENs via compile-time dispatch. The compiler defines
|
||||
// __riscv_v_min_vlen from the zvl<N>b extension in -march. The defs header
|
||||
// maps VLEN to the correct LMUL suffixes, and the impl header provides
|
||||
// VLEN-independent class implementations.
|
||||
//
|
||||
// To add support for a new VLEN, add the LMUL mapping in
|
||||
// cpu_types_riscv_defs.hpp (the impl header needs no changes).
|
||||
|
||||
#ifndef __riscv_vector
|
||||
#error "cpu_types_riscv.hpp included in a non-RVV translation unit"
|
||||
#endif
|
||||
|
||||
#ifndef __riscv_v_min_vlen
|
||||
#error "compiler did not define __riscv_v_min_vlen; pass -march=...zvl<N>b"
|
||||
#endif
|
||||
|
||||
#include "cpu_types_riscv_defs.hpp"
|
||||
#include "cpu_types_riscv_impl.hpp"
|
||||
|
||||
#endif // CPU_TYPES_RISCV_HPP
|
||||
@@ -0,0 +1,121 @@
|
||||
#ifndef CPU_TYPES_RISCV_DEFS_HPP
|
||||
#define CPU_TYPES_RISCV_DEFS_HPP
|
||||
|
||||
// VLEN-to-LMUL mapping for RISC-V Vector extension.
|
||||
//
|
||||
// LMUL_<N> expands to the LMUL suffix giving N total bits of vector data.
|
||||
// LMUL_64 is used by 8-lane int8/uint8 vectors.
|
||||
// VLEN=128:
|
||||
// LMUL_64=mf2, LMUL_128=m1, LMUL_256=m2, LMUL_512=m4, LMUL_1024=m8
|
||||
// VLEN=256:
|
||||
// LMUL_64=mf4, LMUL_128=mf2, LMUL_256=m1, LMUL_512=m2, LMUL_1024=m4
|
||||
|
||||
#include <riscv_vector.h>
|
||||
|
||||
#if __riscv_v_min_vlen == 128
|
||||
#define LMUL_64 mf2
|
||||
#define LMUL_128 m1
|
||||
#define LMUL_256 m2
|
||||
#define LMUL_512 m4
|
||||
#define LMUL_1024 m8
|
||||
#define BOOL_256 b16
|
||||
#define BOOL_512 b8
|
||||
#elif __riscv_v_min_vlen == 256
|
||||
#define LMUL_64 mf4
|
||||
#define LMUL_128 mf2
|
||||
#define LMUL_256 m1
|
||||
#define LMUL_512 m2
|
||||
#define LMUL_1024 m4
|
||||
#define BOOL_256 b32
|
||||
#define BOOL_512 b16
|
||||
#else
|
||||
#error "cpu_types_riscv_defs.hpp: unsupported __riscv_v_min_vlen"
|
||||
#endif
|
||||
|
||||
// Token-paste helpers.
|
||||
#define _RVV_P2(a, b) a##b
|
||||
#define _RVV_P3(a, b, c) a##b##c
|
||||
#define _RVV_P4(a, b, c, d) a##b##c##d
|
||||
#define RVVTYPE(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
|
||||
#define RVVI(base, lmul) _RVV_P2(base, lmul)
|
||||
#define RVVI3(base, lmul, suffix) _RVV_P3(base, lmul, suffix)
|
||||
#define RVVI4(a, b, c, d) _RVV_P4(a, b, c, d)
|
||||
// For mask intrinsics: RVVIB(base, LMUL_256, BOOL_256) → base##m2##_##b16
|
||||
#define _RVV_PB(base, lmul, btype) base##lmul##_##btype
|
||||
#define RVVIB(base, lmul, btype) _RVV_PB(base, lmul, btype)
|
||||
|
||||
// ---- Semantic fixed-vector typedefs (named by element count) ----
|
||||
|
||||
// uint8 / int8
|
||||
typedef RVVTYPE(vuint8, LMUL_64, _t) fixed_u8x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(64)));
|
||||
typedef RVVTYPE(vint8, LMUL_64, _t) fixed_i8x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(64)));
|
||||
|
||||
// int16
|
||||
typedef RVVTYPE(vint16, LMUL_128, _t) fixed_i16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
|
||||
// float16
|
||||
typedef RVVTYPE(vfloat16, LMUL_128, _t) fixed_fp16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vfloat16, LMUL_256, _t) fixed_fp16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
|
||||
// float32
|
||||
typedef RVVTYPE(vfloat32, LMUL_128, _t) fixed_fp32x4_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_256, _t) fixed_fp32x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_512, _t) fixed_fp32x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
typedef RVVTYPE(vfloat32, LMUL_1024, _t) fixed_fp32x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(1024)));
|
||||
|
||||
// int8
|
||||
typedef RVVTYPE(vint8, LMUL_128, _t) fixed_i8x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
|
||||
// int32
|
||||
typedef RVVTYPE(vint32, LMUL_256, _t) fixed_i32x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vint32, LMUL_512, _t) fixed_i32x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
// uint16
|
||||
typedef RVVTYPE(vuint16, LMUL_128, _t) fixed_u16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vuint16, LMUL_256, _t) fixed_u16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vuint16, LMUL_512, _t) fixed_u16x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
|
||||
// uint32
|
||||
typedef RVVTYPE(vuint32, LMUL_256, _t) fixed_u32x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
|
||||
// bfloat16
|
||||
#ifdef __riscv_zvfbfmin
|
||||
typedef RVVTYPE(vbfloat16, LMUL_128, _t) fixed_bf16x8_t
|
||||
__attribute__((riscv_rvv_vector_bits(128)));
|
||||
typedef RVVTYPE(vbfloat16, LMUL_256, _t) fixed_bf16x16_t
|
||||
__attribute__((riscv_rvv_vector_bits(256)));
|
||||
typedef RVVTYPE(vbfloat16, LMUL_512, _t) fixed_bf16x32_t
|
||||
__attribute__((riscv_rvv_vector_bits(512)));
|
||||
#endif
|
||||
|
||||
// ---- Reduction accumulator type (always m1 = one register of f32) ----
|
||||
// Used for scalar reductions; only element [0] is meaningful.
|
||||
typedef vfloat32m1_t rvv_f32_accum_t
|
||||
__attribute__((riscv_rvv_vector_bits(__riscv_v_min_vlen)));
|
||||
|
||||
// ---- Mask types for f32 elements ----
|
||||
#if __riscv_v_min_vlen == 128
|
||||
typedef vbool16_t rvv_mask_f32x8_t;
|
||||
typedef vbool8_t rvv_mask_f32x16_t;
|
||||
#elif __riscv_v_min_vlen == 256
|
||||
typedef vbool32_t rvv_mask_f32x8_t;
|
||||
typedef vbool16_t rvv_mask_f32x16_t;
|
||||
#endif
|
||||
|
||||
#endif // CPU_TYPES_RISCV_DEFS_HPP
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,482 @@
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <torch/all.h>
|
||||
#include "float_convert.hpp"
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
|
||||
#ifndef CPU_OP_GUARD
|
||||
#define CPU_KERNEL_GUARD_IN(NAME)
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME)
|
||||
#else
|
||||
#define CPU_KERNEL_GUARD_IN(NAME) \
|
||||
std::cout << #NAME << " invoked." << std::endl;
|
||||
#define CPU_KERNEL_GUARD_OUT(NAME) \
|
||||
std::cout << #NAME << " exit." << std::endl;
|
||||
#endif
|
||||
|
||||
#define FORCE_INLINE __attribute__((always_inline)) inline
|
||||
|
||||
typedef struct f16x8_t {
|
||||
uint16_t val[8];
|
||||
} f16x8_t;
|
||||
|
||||
typedef struct f16x16_t {
|
||||
uint16_t val[16];
|
||||
} f16x16_t;
|
||||
|
||||
typedef struct f16x32_t {
|
||||
uint16_t val[32];
|
||||
} f16x32_t;
|
||||
|
||||
typedef struct f32x4_t {
|
||||
float val[4];
|
||||
} f32x4_t;
|
||||
|
||||
typedef struct f32x8_t {
|
||||
float val[8];
|
||||
} f32x8_t;
|
||||
|
||||
typedef struct f32x16_t {
|
||||
float val[16];
|
||||
} f32x16_t;
|
||||
|
||||
namespace {
|
||||
template <typename T, T... indexes, typename F>
|
||||
constexpr void unroll_loop_item(std::integer_sequence<T, indexes...>, F&& f) {
|
||||
(f(std::integral_constant<T, indexes>{}), ...);
|
||||
};
|
||||
}; // namespace
|
||||
|
||||
template <typename T, T count, typename F,
|
||||
typename = std::enable_if_t<std::is_invocable_v<F, T> > >
|
||||
constexpr void unroll_loop(F&& f) {
|
||||
unroll_loop_item(std::make_integer_sequence<T, count>{}, std::forward<F>(f));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct Vec {
|
||||
constexpr static int get_elem_num() { return T::VEC_ELEM_NUM; }
|
||||
};
|
||||
|
||||
struct FP32Vec8;
|
||||
struct FP32Vec16;
|
||||
|
||||
struct FP16Vec8 : public Vec<FP16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
f16x8_t reg;
|
||||
|
||||
explicit FP16Vec8(const void* ptr)
|
||||
: reg(*reinterpret_cast<const f16x8_t*>(ptr)) {};
|
||||
|
||||
explicit FP16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x8_t*>(ptr) = reg; }
|
||||
};
|
||||
|
||||
struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
f16x16_t reg;
|
||||
|
||||
explicit FP16Vec16(const void* ptr)
|
||||
: reg(*reinterpret_cast<const f16x16_t*>(ptr)) {};
|
||||
|
||||
explicit FP16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
|
||||
|
||||
void save(void* ptr, const int elem_num) const {
|
||||
int num = std::min(elem_num, VEC_ELEM_NUM);
|
||||
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec8 : public Vec<BF16Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
f16x8_t reg;
|
||||
|
||||
explicit BF16Vec8(const void* ptr)
|
||||
: reg(*reinterpret_cast<const f16x8_t*>(ptr)) {};
|
||||
|
||||
explicit BF16Vec8(const FP32Vec8&);
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x8_t*>(ptr) = reg; }
|
||||
};
|
||||
|
||||
struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
f16x16_t reg;
|
||||
|
||||
explicit BF16Vec16(const void* ptr)
|
||||
: reg(*reinterpret_cast<const f16x16_t*>(ptr)) {};
|
||||
|
||||
explicit BF16Vec16(const FP32Vec16&);
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x16_t*>(ptr) = reg; }
|
||||
|
||||
void save(void* ptr, const int elem_num) const {
|
||||
int num = std::min(elem_num, VEC_ELEM_NUM);
|
||||
std::memcpy(ptr, &(reg.val[0]), num * sizeof(uint16_t));
|
||||
}
|
||||
};
|
||||
|
||||
struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
constexpr static int VEC_ELEM_NUM = 32;
|
||||
f16x32_t reg;
|
||||
|
||||
explicit BF16Vec32(const void* ptr)
|
||||
: reg(*reinterpret_cast<const f16x32_t*>(ptr)) {};
|
||||
|
||||
explicit BF16Vec32(f16x32_t data) : reg(data) {};
|
||||
|
||||
explicit BF16Vec32(BF16Vec8& vec8_data) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&vec8_data, this](int i) {
|
||||
reg.val[i] = vec8_data.reg.val[i % BF16Vec8::VEC_ELEM_NUM];
|
||||
});
|
||||
}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
|
||||
};
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
constexpr static int VEC_ELEM_NUM = 4;
|
||||
|
||||
f32x4_t reg;
|
||||
|
||||
explicit FP32Vec4(float v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
|
||||
}
|
||||
|
||||
explicit FP32Vec4() {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
|
||||
}
|
||||
|
||||
explicit FP32Vec4(const float* ptr)
|
||||
: reg(*reinterpret_cast<const f32x4_t*>(ptr)) {};
|
||||
|
||||
explicit FP32Vec4(f32x4_t data) : reg(data) {};
|
||||
|
||||
explicit FP32Vec4(const FP32Vec4& data) : reg(data.reg) {};
|
||||
};
|
||||
|
||||
struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
constexpr static int VEC_ELEM_NUM = 8;
|
||||
|
||||
f32x8_t reg;
|
||||
|
||||
explicit FP32Vec8(float v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
|
||||
}
|
||||
|
||||
explicit FP32Vec8() {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
|
||||
}
|
||||
|
||||
explicit FP32Vec8(const float* ptr)
|
||||
: reg(*reinterpret_cast<const f32x8_t*>(ptr)) {};
|
||||
|
||||
explicit FP32Vec8(f32x8_t data) : reg(data) {};
|
||||
|
||||
explicit FP32Vec8(const FP32Vec8& data) : reg(data.reg) {};
|
||||
|
||||
explicit FP32Vec8(const FP16Vec8& v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
FP32Vec8(const BF16Vec8& v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
float reduce_sum() const {
|
||||
float result = 0;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, this](int i) { result += reg.val[i]; });
|
||||
return result;
|
||||
}
|
||||
|
||||
FP32Vec8 exp() const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = expf(reg.val[i]); });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 tanh() const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = tanhf(reg.val[i]); });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 er() const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = erf(reg.val[i]); });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 operator*(const FP32Vec8& b) const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 operator+(const FP32Vec8& b) const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 operator-(const FP32Vec8& b) const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 operator/(const FP32Vec8& b) const {
|
||||
f32x8_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f32x8_t*>(ptr) = reg; }
|
||||
};
|
||||
|
||||
struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
f32x16_t reg;
|
||||
|
||||
explicit FP32Vec16(float v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&v, this](int i) { reg.val[i] = v; });
|
||||
}
|
||||
|
||||
explicit FP32Vec16() {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([this](int i) { reg.val[i] = 0.0f; });
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const float* ptr)
|
||||
: reg(*reinterpret_cast<const f32x16_t*>(ptr)) {};
|
||||
|
||||
explicit FP32Vec16(f32x16_t data) : reg(data) {};
|
||||
|
||||
FP32Vec16(const FP32Vec4& data) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
|
||||
reg.val[i] = data.reg.val[i % FP32Vec4::VEC_ELEM_NUM];
|
||||
});
|
||||
}
|
||||
|
||||
FP32Vec16(const FP32Vec8& data) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&data, this](int i) {
|
||||
reg.val[i] = data.reg.val[i % FP32Vec8::VEC_ELEM_NUM];
|
||||
});
|
||||
}
|
||||
|
||||
FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = fp16_to_float(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const BF16Vec16& v) {
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = bf16_to_float(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
|
||||
|
||||
FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
|
||||
|
||||
// FP8 stub: dead code on scalar path (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
|
||||
|
||||
FP32Vec16 operator*(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] * b.reg.val[i]; });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 operator+(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] + b.reg.val[i]; });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 operator-(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] - b.reg.val[i]; });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 operator/(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, &b, this](int i) { ret.val[i] = reg.val[i] / b.reg.val[i]; });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 max(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
|
||||
ret.val[i] = std::max(reg.val[i], b.reg.val[i]);
|
||||
});
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 min(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>([&ret, &b, this](int i) {
|
||||
ret.val[i] = std::min(reg.val[i], b.reg.val[i]);
|
||||
});
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 abs() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = std::abs(reg.val[i]); });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
FP32Vec16 tanh() const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&ret, this](int i) { ret.val[i] = std::tanh(reg.val[i]); });
|
||||
return FP32Vec16(ret);
|
||||
}
|
||||
|
||||
float reduce_sum() const {
|
||||
float result = 0.0f;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, this](int i) { result += reg.val[i]; });
|
||||
return result;
|
||||
}
|
||||
|
||||
float reduce_max() const {
|
||||
float result = std::numeric_limits<float>::lowest();
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, this](int i) { result = std::max(reg.val[i], result); });
|
||||
return result;
|
||||
}
|
||||
|
||||
float reduce_min() const {
|
||||
float result = std::numeric_limits<float>::max();
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, this](int i) { result = std::min(reg.val[i], result); });
|
||||
return result;
|
||||
}
|
||||
|
||||
template <int group_size>
|
||||
float reduce_sub_sum(int idx) {
|
||||
static_assert(VEC_ELEM_NUM % group_size == 0);
|
||||
float sum = 0.0;
|
||||
const int start = idx * group_size;
|
||||
unroll_loop<int, group_size>(
|
||||
[&sum, &start, this](int i) { sum += reg.val[start + i]; });
|
||||
return sum;
|
||||
}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f32x16_t*>(ptr) = reg; }
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct VecType {
|
||||
using vec_type = void;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using vec_t = typename VecType<T>::vec_type;
|
||||
|
||||
template <>
|
||||
struct VecType<float> {
|
||||
using vec_type = FP32Vec8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::Half> {
|
||||
using vec_type = FP16Vec8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecType<c10::BFloat16> {
|
||||
using vec_type = BF16Vec8;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
void storeFP32(float v, T* ptr) {
|
||||
*ptr = v;
|
||||
}
|
||||
|
||||
/*
|
||||
template <> inline void storeFP32<c10::Half>(float v, c10::Half *ptr) {
|
||||
c10::Half __attribute__((__may_alias__)) *v_ptr =
|
||||
reinterpret_cast<c10::Half *>(&v);
|
||||
*ptr = *(v_ptr + 1);
|
||||
}
|
||||
*/
|
||||
|
||||
template <>
|
||||
inline void storeFP32<c10::Half>(float v, c10::Half* ptr) {
|
||||
uint16_t fp16 = float_to_fp16(v);
|
||||
*reinterpret_cast<uint16_t*>(ptr) = fp16;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline void storeFP32<c10::BFloat16>(float v, c10::BFloat16* ptr) {
|
||||
c10::BFloat16 __attribute__((__may_alias__))* v_ptr =
|
||||
reinterpret_cast<c10::BFloat16*>(&v);
|
||||
*ptr = *(v_ptr + 1);
|
||||
}
|
||||
|
||||
inline FP16Vec16::FP16Vec16(const FP32Vec16& v) {
|
||||
unroll_loop<int, FP16Vec16::VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
inline FP16Vec8 ::FP16Vec8(const FP32Vec8& v) {
|
||||
unroll_loop<int, FP16Vec8::VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = float_to_fp16(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
inline void fma(FP32Vec16& acc, FP32Vec16& a, FP32Vec16& b) {
|
||||
acc = acc + a * b;
|
||||
}
|
||||
|
||||
inline BF16Vec8::BF16Vec8(const FP32Vec8& v) {
|
||||
unroll_loop<int, BF16Vec8::VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
inline BF16Vec16::BF16Vec16(const FP32Vec16& v) {
|
||||
unroll_loop<int, BF16Vec16::VEC_ELEM_NUM>(
|
||||
[&v, this](int i) { reg.val[i] = float_to_bf16(v.reg.val[i]); });
|
||||
}
|
||||
|
||||
inline void prefetch(const void* addr) { __builtin_prefetch(addr, 0, 3); }
|
||||
|
||||
}; // namespace vec_op
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,441 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_amx.hpp"
|
||||
#endif
|
||||
#if defined(__riscv_v)
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_rvv.hpp"
|
||||
#endif
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
|
||||
|
||||
#define VLLM_DISPATCH_CASE_16B_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
|
||||
|
||||
#define VLLM_DISPATCH_16B_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_16B_TYPES(__VA_ARGS__))
|
||||
|
||||
template <typename T>
|
||||
void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
|
||||
int32_t stride) {
|
||||
std::stringstream ss;
|
||||
ss << std::fixed << std::setprecision(5) << name << ": [\n";
|
||||
auto* curr_logits_buffer = ptr;
|
||||
for (int32_t m = 0; m < row; ++m) {
|
||||
for (int32_t n = 0; n < col; ++n) {
|
||||
ss << curr_logits_buffer[n] << ", ";
|
||||
}
|
||||
ss << "\n";
|
||||
curr_logits_buffer += stride;
|
||||
}
|
||||
ss << "]\n";
|
||||
std::printf("%s", ss.str().c_str());
|
||||
}
|
||||
|
||||
namespace {
|
||||
using cpu_utils::ISA;
|
||||
using cpu_utils::VecTypeTrait;
|
||||
|
||||
template <typename scalar_t, ISA isa, bool has_zp, bool use_desc_act>
|
||||
class Dequantizer4b {
|
||||
public:
|
||||
constexpr static int32_t pack_num = 32 / 4;
|
||||
using scalar_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
public:
|
||||
static void dequant(int32_t* __restrict__ q_weight,
|
||||
scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ scales,
|
||||
int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
|
||||
const int64_t scales_stride, const int64_t zeros_stride,
|
||||
const int32_t k_size, const int32_t group_size) {
|
||||
vec_op::FP32Vec16 lut;
|
||||
if constexpr (has_zp) {
|
||||
// AWQ
|
||||
alignas(64) static const float LUT[16] = {
|
||||
0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f,
|
||||
8.0f, 9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f};
|
||||
lut = vec_op::FP32Vec16(LUT);
|
||||
} else {
|
||||
// GPTQ
|
||||
alignas(64) static const float LUT[16] = {
|
||||
-8.0f, -7.0f, -6.0f, -5.0f, -4.0f, -3.0f, -2.0f, -1.0f,
|
||||
0.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f};
|
||||
lut = vec_op::FP32Vec16(LUT);
|
||||
}
|
||||
|
||||
// per 64-bits elem contains 16 output channels
|
||||
int64_t* __restrict__ curr_q_weight = reinterpret_cast<int64_t*>(q_weight);
|
||||
int64_t* __restrict__ curr_zeros = reinterpret_cast<int64_t*>(zeros);
|
||||
scalar_t* __restrict__ curr_weight = weight;
|
||||
scalar_t* __restrict__ curr_scale = scales;
|
||||
vec_op::FP32Vec16 scale_0;
|
||||
vec_op::FP32Vec16 scale_1;
|
||||
vec_op::FP32Vec16 zero_0;
|
||||
vec_op::FP32Vec16 zero_1;
|
||||
int32_t group_counter = 0;
|
||||
for (int32_t k_idx = 0; k_idx < k_size; k_idx += 2) {
|
||||
int64_t qwb_0 = *curr_q_weight;
|
||||
int64_t qwb_1 = *(curr_q_weight + 1);
|
||||
vec_op::FP32Vec16 wb_0(qwb_0, lut);
|
||||
vec_op::FP32Vec16 wb_1(qwb_1, lut);
|
||||
|
||||
if constexpr (!use_desc_act) {
|
||||
if (group_counter == 0) {
|
||||
scale_0 = vec_op::FP32Vec16(scalar_vec_t(curr_scale));
|
||||
scale_1 = vec_op::FP32Vec16(scale_0);
|
||||
curr_scale += scales_stride;
|
||||
|
||||
if constexpr (has_zp) {
|
||||
zero_0 = vec_op::FP32Vec16(*curr_zeros, lut);
|
||||
zero_1 = vec_op::FP32Vec16(zero_0);
|
||||
curr_zeros += zeros_stride / 2;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
int32_t g_idx_0 = g_idx[k_idx];
|
||||
int32_t g_idx_1 = g_idx[k_idx + 1];
|
||||
scale_0 = vec_op::FP32Vec16(
|
||||
scalar_vec_t(curr_scale + g_idx_0 * scales_stride));
|
||||
scale_1 = vec_op::FP32Vec16(
|
||||
scalar_vec_t(curr_scale + g_idx_1 * scales_stride));
|
||||
if constexpr (has_zp) {
|
||||
zero_0 = vec_op::FP32Vec16(*(curr_zeros + g_idx_0 * zeros_stride / 2),
|
||||
lut);
|
||||
zero_1 = vec_op::FP32Vec16(*(curr_zeros + g_idx_1 * zeros_stride / 2),
|
||||
lut);
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (has_zp) {
|
||||
wb_0 = wb_0 - zero_0;
|
||||
wb_1 = wb_1 - zero_1;
|
||||
}
|
||||
|
||||
wb_0 = wb_0 * scale_0;
|
||||
wb_1 = wb_1 * scale_1;
|
||||
|
||||
scalar_vec_t output_vec_0(wb_0);
|
||||
scalar_vec_t output_vec_1(wb_1);
|
||||
|
||||
// AMX needs to interleave K elements to pack as 32 bits
|
||||
if constexpr (isa == ISA::AMX) {
|
||||
vec_op::interleave_save(output_vec_0, output_vec_1, curr_weight);
|
||||
} else {
|
||||
output_vec_0.save(curr_weight);
|
||||
output_vec_1.save(curr_weight + 16);
|
||||
}
|
||||
|
||||
// update
|
||||
curr_q_weight += 2;
|
||||
curr_weight += 32;
|
||||
if constexpr (!use_desc_act) {
|
||||
group_counter += 2;
|
||||
if (group_counter == group_size) {
|
||||
group_counter = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
}; // namespace
|
||||
|
||||
template <typename scalar_t, typename dequantizer_t, typename gemm_t>
|
||||
void cpu_gemm_wna16_impl(
|
||||
scalar_t* __restrict__ input, int32_t* __restrict__ q_weight,
|
||||
scalar_t* __restrict__ output, scalar_t* __restrict__ scales,
|
||||
int32_t* __restrict__ zeros, int32_t* __restrict__ g_idx,
|
||||
scalar_t* __restrict__ bias, const int32_t m_size, const int32_t n_size,
|
||||
const int32_t k_size, const int64_t input_stride,
|
||||
const int64_t output_stride, const int64_t scales_group_stride,
|
||||
const int64_t zeros_group_stride, const int32_t group_num,
|
||||
const int32_t group_size, const int64_t pack_factor) {
|
||||
constexpr int32_t gemm_n_tile_size = gemm_t::NSize;
|
||||
constexpr int32_t gemm_m_tile_size = gemm_t::MaxMSize;
|
||||
constexpr int32_t n_block_size = 16;
|
||||
static_assert(gemm_n_tile_size % n_block_size == 0);
|
||||
const int32_t thread_num = cpu_utils::get_max_threads();
|
||||
|
||||
// a simple schedule policy, just to hold more B tiles in L2 and make sure
|
||||
// each thread has tasks
|
||||
const int32_t n_partition_size = [&]() {
|
||||
const int64_t cache_size = cpu_utils::get_available_l2_size();
|
||||
int64_t ps_cache_limit = cache_size / (k_size * sizeof(scalar_t));
|
||||
int64_t ps_thread_limit = n_size / thread_num;
|
||||
ps_cache_limit =
|
||||
std::max((ps_cache_limit / gemm_n_tile_size) * gemm_n_tile_size,
|
||||
(int64_t)gemm_n_tile_size);
|
||||
ps_thread_limit =
|
||||
std::max((ps_thread_limit / gemm_n_tile_size) * gemm_n_tile_size,
|
||||
(int64_t)gemm_n_tile_size);
|
||||
return std::min(ps_cache_limit, ps_thread_limit);
|
||||
}();
|
||||
const int32_t task_num = (n_size + n_partition_size - 1) / n_partition_size;
|
||||
|
||||
// get buffer size
|
||||
const int64_t b_buffer_size =
|
||||
(((n_partition_size * k_size * sizeof(scalar_t) + 63) / 64) * 64);
|
||||
const int64_t c_buffer_size =
|
||||
(((gemm_m_tile_size * gemm_n_tile_size * sizeof(float) + 63) / 64) * 64);
|
||||
const int64_t b_buffer_offset = 0;
|
||||
const int64_t c_buffer_offset = b_buffer_size;
|
||||
const int64_t buffer_size = b_buffer_size + c_buffer_size;
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size *
|
||||
thread_num);
|
||||
|
||||
alignas(64) cpu_utils::Counter counter;
|
||||
cpu_utils::Counter* counter_ptr = &counter;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int32_t thread_id = 0; thread_id < thread_num; ++thread_id) {
|
||||
scalar_t* __restrict__ b_buffer = nullptr;
|
||||
float* __restrict__ c_buffer = nullptr;
|
||||
{
|
||||
uint8_t* buffer_ptr =
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()
|
||||
->get_data<uint8_t>() +
|
||||
thread_id * buffer_size;
|
||||
b_buffer = reinterpret_cast<scalar_t*>(buffer_ptr + b_buffer_offset);
|
||||
c_buffer = reinterpret_cast<float*>(buffer_ptr + c_buffer_offset);
|
||||
}
|
||||
|
||||
const int64_t q_weight_block_stride = n_block_size / pack_factor * k_size;
|
||||
const int64_t b_buffer_block_stride = n_block_size * k_size;
|
||||
const int32_t zeros_block_stride = n_block_size / pack_factor;
|
||||
|
||||
gemm_t gemm;
|
||||
|
||||
for (;;) {
|
||||
int32_t task_id = counter_ptr->acquire_counter();
|
||||
|
||||
if (task_id >= task_num) {
|
||||
break;
|
||||
}
|
||||
|
||||
const int32_t n_start_idx = task_id * n_partition_size;
|
||||
const int32_t n_block_start_idx = n_start_idx / n_block_size;
|
||||
const int32_t n_num = std::min(n_partition_size, n_size - n_start_idx);
|
||||
const int32_t n_block_num = n_num / n_block_size;
|
||||
// std::printf("thread_id: %d, task_id: %d, n_start_idx: %d, n_num: %d\n",
|
||||
// thread_id, task_id, n_start_idx, n_num);
|
||||
|
||||
// dequant weight
|
||||
{
|
||||
int32_t* __restrict__ curr_q_weight =
|
||||
q_weight + n_block_start_idx * q_weight_block_stride;
|
||||
scalar_t* __restrict__ curr_b_buffer = b_buffer;
|
||||
scalar_t* __restrict__ curr_scales = scales + n_start_idx;
|
||||
int32_t* __restrict__ curr_zeros = zeros + n_start_idx / pack_factor;
|
||||
for (int32_t block_idx = 0; block_idx < n_block_num; ++block_idx) {
|
||||
dequantizer_t::dequant(curr_q_weight, curr_b_buffer, curr_scales,
|
||||
curr_zeros, g_idx, scales_group_stride,
|
||||
zeros_group_stride, k_size, group_size);
|
||||
|
||||
// if (block_idx == 0 && n_start_idx == 0) {
|
||||
// print_logits("depacked weight", curr_b_buffer, k_size,
|
||||
// n_block_size, n_block_size);
|
||||
// }
|
||||
|
||||
// update
|
||||
curr_q_weight += q_weight_block_stride;
|
||||
curr_b_buffer += b_buffer_block_stride;
|
||||
curr_scales += n_block_size;
|
||||
curr_zeros += zeros_block_stride;
|
||||
}
|
||||
}
|
||||
|
||||
// compute loop
|
||||
{
|
||||
const int32_t n_tile_num = n_num / gemm_n_tile_size;
|
||||
scalar_t* __restrict__ curr_input = input;
|
||||
scalar_t* __restrict__ init_bias = bias;
|
||||
if (bias != nullptr) {
|
||||
init_bias += n_start_idx;
|
||||
}
|
||||
scalar_t* __restrict__ init_output = output + n_start_idx;
|
||||
for (int32_t m_idx = 0; m_idx < m_size; m_idx += gemm_m_tile_size) {
|
||||
const int32_t curr_m_size =
|
||||
std::min(gemm_m_tile_size, m_size - m_idx);
|
||||
scalar_t* __restrict__ curr_b_buffer = b_buffer;
|
||||
scalar_t* __restrict__ curr_bias = init_bias;
|
||||
scalar_t* __restrict__ curr_output = init_output;
|
||||
for (int32_t n_tile_idx = 0; n_tile_idx < n_tile_num; ++n_tile_idx) {
|
||||
gemm.gemm(curr_input, curr_b_buffer, c_buffer, curr_m_size, k_size,
|
||||
input_stride, b_buffer_block_stride, gemm_n_tile_size,
|
||||
false);
|
||||
|
||||
if (bias != nullptr) {
|
||||
cpu_micro_gemm::bias_epilogue<gemm_n_tile_size>(
|
||||
c_buffer, curr_output, curr_bias, curr_m_size,
|
||||
gemm_n_tile_size, output_stride);
|
||||
curr_bias += gemm_n_tile_size;
|
||||
} else {
|
||||
cpu_micro_gemm::default_epilogue<gemm_n_tile_size>(
|
||||
c_buffer, curr_output, curr_m_size, gemm_n_tile_size,
|
||||
output_stride);
|
||||
}
|
||||
|
||||
curr_b_buffer +=
|
||||
b_buffer_block_stride * (gemm_n_tile_size / n_block_size);
|
||||
curr_output += gemm_n_tile_size;
|
||||
}
|
||||
curr_input += gemm_m_tile_size * input_stride;
|
||||
init_output += gemm_m_tile_size * output_stride;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void cpu_gemm_wna16(
|
||||
const torch::Tensor& input, // [M, K]
|
||||
const torch::Tensor&
|
||||
q_weight, // [N / 16, K * 16 / pack_factor], packed as int32
|
||||
torch::Tensor& output, // [M, N]
|
||||
const torch::Tensor& scales, // [group_num, N]
|
||||
const std::optional<torch::Tensor>&
|
||||
zeros, // [group_num, N / pack_factor], packed as int32
|
||||
const std::optional<torch::Tensor>& g_idx, // [K]
|
||||
const std::optional<torch::Tensor>& bias, // [N]
|
||||
const int64_t pack_factor, const std::string& isa_hint) {
|
||||
using cpu_utils::ISA;
|
||||
TORCH_CHECK_EQ(pack_factor, 8); // only supports 4bits
|
||||
const int32_t a_m_size = input.size(0);
|
||||
const int32_t a_k_size = input.size(1);
|
||||
const int64_t a_m_stride = input.stride(0);
|
||||
const int32_t b_n_size = q_weight.size(0) * 16;
|
||||
TORCH_CHECK_EQ(a_k_size % 32, 0);
|
||||
TORCH_CHECK_EQ(b_n_size % 32, 0);
|
||||
const int32_t group_num = scales.size(0);
|
||||
const int32_t group_size = a_k_size / group_num;
|
||||
TORCH_CHECK_EQ(group_size % 2, 0);
|
||||
const int64_t scales_group_stride = scales.stride(0);
|
||||
const int64_t output_m_stride = output.stride(0);
|
||||
|
||||
bool has_zp = zeros.has_value();
|
||||
bool use_desc_act = g_idx.has_value();
|
||||
TORCH_CHECK(!(has_zp && use_desc_act));
|
||||
|
||||
ISA isa = [&]() {
|
||||
if (isa_hint == "amx") {
|
||||
return ISA::AMX;
|
||||
} else if (isa_hint == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa_hint == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else {
|
||||
TORCH_CHECK(false, "unsupported isa hint: " + isa_hint);
|
||||
}
|
||||
}();
|
||||
|
||||
int32_t* zeros_ptr = has_zp ? zeros->data_ptr<int32_t>() : nullptr;
|
||||
const int64_t zeros_group_stride = has_zp ? zeros->stride(0) : 0;
|
||||
int32_t* g_idx_ptr = use_desc_act ? g_idx->data_ptr<int32_t>() : nullptr;
|
||||
|
||||
VLLM_DISPATCH_16B_TYPES(input.scalar_type(), "cpu_gemm_wna16", [&]() {
|
||||
if (isa == ISA::AMX) {
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::AMX, scalar_t>;
|
||||
if (has_zp) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, true, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
if (use_desc_act) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, true>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
} else {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::AMX, false, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
} else if (isa == ISA::VEC) {
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::VEC, scalar_t>;
|
||||
if (has_zp) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, true, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
if (use_desc_act) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, true>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
} else {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::VEC, false, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
} else if (isa == ISA::RVV) {
|
||||
using gemm_t = cpu_micro_gemm::MicroGemm<ISA::RVV, scalar_t>;
|
||||
if (has_zp) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, true, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
if (use_desc_act) {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, true>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
} else {
|
||||
using dequantizer_t = Dequantizer4b<scalar_t, ISA::RVV, false, false>;
|
||||
cpu_gemm_wna16_impl<scalar_t, dequantizer_t, gemm_t>(
|
||||
input.data_ptr<scalar_t>(), q_weight.data_ptr<int32_t>(),
|
||||
output.data_ptr<scalar_t>(), scales.data_ptr<scalar_t>(), zeros_ptr,
|
||||
g_idx_ptr, bias.has_value() ? bias->data_ptr<scalar_t>() : nullptr,
|
||||
a_m_size, b_n_size, a_k_size, a_m_stride, output_m_stride,
|
||||
scales_group_stride, zeros_group_stride, group_num, group_size,
|
||||
pack_factor);
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,572 @@
|
||||
#include <list>
|
||||
#include <optional>
|
||||
|
||||
#include "common/memory_desc.hpp"
|
||||
#include "common/memory.hpp"
|
||||
|
||||
#include "cpu/utils.hpp"
|
||||
#include "cpu/dnnl_helper.h"
|
||||
|
||||
static dnnl::engine& default_engine() {
|
||||
static dnnl::engine engine(dnnl::engine::kind::cpu, 0);
|
||||
return engine;
|
||||
}
|
||||
|
||||
static dnnl::stream& default_stream() {
|
||||
static dnnl::stream stream(default_engine());
|
||||
return stream;
|
||||
}
|
||||
|
||||
void release_dnnl_matmul_handler(int64_t handler) {
|
||||
DNNLMatMulPrimitiveHandler* ptr =
|
||||
reinterpret_cast<DNNLMatMulPrimitiveHandler*>(handler);
|
||||
delete ptr;
|
||||
}
|
||||
|
||||
template <typename KT, typename VT>
|
||||
class DNNLPrimitiveCache {
|
||||
public:
|
||||
using cache_value_t = std::pair<KT, VT>;
|
||||
using result_value_t = VT;
|
||||
using container_t = std::list<cache_value_t>;
|
||||
using value_iterator_t = typename container_t::iterator;
|
||||
using map_t = std::unordered_map<KT, value_iterator_t>;
|
||||
using creator_t = VT (*)();
|
||||
|
||||
public:
|
||||
DNNLPrimitiveCache(size_t capacity)
|
||||
: capacity_(capacity),
|
||||
values_(),
|
||||
key_to_value_(std::min(256lu, capacity)) {
|
||||
assert(capacity > 0);
|
||||
}
|
||||
|
||||
template <typename F>
|
||||
result_value_t get_or_create(const KT& key, F&& creator) {
|
||||
std::optional<value_iterator_t> value = get_value(key);
|
||||
if (value.has_value()) {
|
||||
return value.value()->second;
|
||||
} else {
|
||||
return add_value({key, creator()})->second;
|
||||
}
|
||||
}
|
||||
|
||||
size_t size() const { return values_.size(); }
|
||||
|
||||
private:
|
||||
void dump_data() {
|
||||
std::stringstream ss;
|
||||
ss << "table_id: " << std::hex << reinterpret_cast<size_t>(this) << std::dec
|
||||
<< "\n";
|
||||
ss << "container: [";
|
||||
for (auto&& iter : values_) {
|
||||
ss << "(" << iter.first << ", " << std::hex
|
||||
<< reinterpret_cast<size_t>(iter.second.get()) << "), " << std::dec;
|
||||
}
|
||||
ss << "]\n";
|
||||
|
||||
ss << "map: [";
|
||||
for (auto&& iter : key_to_value_) {
|
||||
ss << "(" << iter.first << ", " << iter.second->first << ", " << std::hex
|
||||
<< reinterpret_cast<size_t>(iter.second->second.get()) << std::dec
|
||||
<< "), ";
|
||||
}
|
||||
ss << "]\n";
|
||||
std::printf("%s\n", ss.str().c_str());
|
||||
}
|
||||
|
||||
value_iterator_t add_value(cache_value_t&& new_value) {
|
||||
if (size() == capacity_) {
|
||||
cache_value_t& last_item = values_.back();
|
||||
key_to_value_.erase(last_item.first);
|
||||
values_.pop_back();
|
||||
}
|
||||
|
||||
auto& added_value_ = values_.emplace_front(std::move(new_value));
|
||||
key_to_value_.emplace(added_value_.first, values_.begin());
|
||||
return values_.begin();
|
||||
}
|
||||
|
||||
std::optional<value_iterator_t> get_value(const KT& key) {
|
||||
if (key_to_value_.size() > 0 && key == values_.begin()->first) {
|
||||
return values_.begin();
|
||||
}
|
||||
|
||||
auto value_map_iterator = key_to_value_.find(key);
|
||||
if (value_map_iterator != key_to_value_.end()) {
|
||||
values_.splice(values_.begin(), values_, value_map_iterator->second);
|
||||
return value_map_iterator->second;
|
||||
} else {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
const size_t capacity_;
|
||||
container_t values_;
|
||||
map_t key_to_value_;
|
||||
};
|
||||
|
||||
DNNLMatMulPrimitiveHandler::DNNLMatMulPrimitiveHandler(
|
||||
const Args& args, dnnl::memory::data_type b_type)
|
||||
: b_n_size_(args.b_n_size),
|
||||
b_n_stride_(args.b_n_stride),
|
||||
b_k_size_(args.b_k_size),
|
||||
b_k_stride_(args.b_k_stride),
|
||||
b_type_(b_type),
|
||||
c_type_(args.c_type),
|
||||
runtime_memory_ptrs_(8),
|
||||
primitive_cache_size_(args.primitive_cache_size) {
|
||||
assert(primitive_cache_size_ > 0);
|
||||
}
|
||||
|
||||
void DNNLMatMulPrimitiveHandler::prepack_weight(
|
||||
void* original_b_ptr, dnnl::memory::desc original_b_md,
|
||||
dnnl::memory::desc b_target_mem_desc) {
|
||||
dnnl::memory original_weight(original_b_md, default_engine(), original_b_ptr);
|
||||
dnnl::memory packed_weight(b_target_mem_desc, default_engine());
|
||||
{
|
||||
dnnl::reorder(original_weight, packed_weight)
|
||||
.execute(default_stream(), original_weight, packed_weight);
|
||||
default_stream().wait();
|
||||
}
|
||||
memory_cache_[DNNL_ARG_WEIGHTS] = packed_weight;
|
||||
b_target_mem_desc_ = b_target_mem_desc;
|
||||
}
|
||||
|
||||
void DNNLMatMulPrimitiveHandler::set_runtime_memory_ptr(
|
||||
size_t index, dnnl_memory* memory_ptr) {
|
||||
dnnl::impl::memory_storage_t* mem_storage_ptr = memory_ptr->memory_storage();
|
||||
dnnl_memory_desc* mem_desc = const_cast<dnnl_memory_desc*>(memory_ptr->md());
|
||||
runtime_memory_ptrs_[index] = {mem_storage_ptr, mem_desc};
|
||||
}
|
||||
|
||||
std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>
|
||||
DNNLMatMulPrimitiveHandler::get_runtime_memory_ptr(size_t index) {
|
||||
return runtime_memory_ptrs_[index];
|
||||
}
|
||||
|
||||
namespace std {
|
||||
template <>
|
||||
struct hash<W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey> {
|
||||
size_t operator()(
|
||||
const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
|
||||
return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
|
||||
hash<int>()(static_cast<int>(val.a_qs)) ^
|
||||
hash<int>()(static_cast<int>(val.b_qs)) ^ hash<bool>()(val.use_azp) ^
|
||||
hash<int>()(static_cast<int>(val.c_type));
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct hash<W8A8MatMulPrimitiveHandler::MSizeCacheKey> {
|
||||
size_t operator()(
|
||||
const W8A8MatMulPrimitiveHandler::MSizeCacheKey& val) const {
|
||||
return hash<dnnl_dim_t>()(val.a_m_size) ^ hash<bool>()(val.use_bias) ^
|
||||
hash<int>()(static_cast<int>(val.bias_type));
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct hash<MatMulPrimitiveHandler::ClassMatmulCacheKey> {
|
||||
size_t operator()(
|
||||
const MatMulPrimitiveHandler::ClassMatmulCacheKey& val) const {
|
||||
return hash<dnnl_dim_t>()(val.b_n_size) ^ hash<dnnl_dim_t>()(val.b_k_size) ^
|
||||
hash<int>()(static_cast<int>(val.b_type));
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct hash<MatMulPrimitiveHandler::MSizeCacheKey> {
|
||||
size_t operator()(const MatMulPrimitiveHandler::MSizeCacheKey& val) const {
|
||||
return hash<dnnl_dim_t>()(val.a_m_size) ^
|
||||
hash<dnnl_dim_t>()(val.a_m_stride) ^ hash<bool>()(val.use_bias) ^
|
||||
hash<int>()(static_cast<int>(val.bias_type));
|
||||
}
|
||||
};
|
||||
} // namespace std
|
||||
|
||||
bool operator==(const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
|
||||
const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
|
||||
return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
|
||||
l.a_qs == r.a_qs && l.b_qs == r.b_qs && l.use_azp == r.use_azp &&
|
||||
l.c_type == r.c_type;
|
||||
}
|
||||
|
||||
bool operator==(const W8A8MatMulPrimitiveHandler::MSizeCacheKey& l,
|
||||
const W8A8MatMulPrimitiveHandler::MSizeCacheKey& r) {
|
||||
return l.use_bias == r.use_bias && l.a_m_size == r.a_m_size &&
|
||||
l.bias_type == r.bias_type;
|
||||
}
|
||||
|
||||
bool operator==(const MatMulPrimitiveHandler::ClassMatmulCacheKey& l,
|
||||
const MatMulPrimitiveHandler::ClassMatmulCacheKey& r) {
|
||||
return l.b_n_size == r.b_n_size && l.b_k_size == r.b_k_size &&
|
||||
l.b_type == r.b_type;
|
||||
}
|
||||
|
||||
bool operator==(const MatMulPrimitiveHandler::MSizeCacheKey& l,
|
||||
const MatMulPrimitiveHandler::MSizeCacheKey& r) {
|
||||
return l.a_m_size == r.a_m_size && l.a_m_stride == r.a_m_stride &&
|
||||
l.use_bias == r.use_bias && l.bias_type == r.bias_type;
|
||||
}
|
||||
|
||||
static std::shared_ptr<W8A8MatMulPrimitiveHandler::MSizeCache>
|
||||
get_w8a8_class_primitive_cache(
|
||||
const W8A8MatMulPrimitiveHandler::ClassMatmulCacheKey& key,
|
||||
int64_t cache_size) {
|
||||
static W8A8MatMulPrimitiveHandler::ClassMatmulCache cache(128);
|
||||
assert(cache_size > 0);
|
||||
return cache.get_or_create(key, [&]() {
|
||||
return std::make_shared<W8A8MatMulPrimitiveHandler::MSizeCache>(cache_size);
|
||||
});
|
||||
}
|
||||
|
||||
W8A8MatMulPrimitiveHandler::W8A8MatMulPrimitiveHandler(const Args& args)
|
||||
: DNNLMatMulPrimitiveHandler(
|
||||
static_cast<const DNNLMatMulPrimitiveHandler::Args&>(args),
|
||||
dnnl::memory::data_type::s8),
|
||||
use_azp_(args.use_a_zero_point),
|
||||
a_qs_(args.a_quantization_strategy),
|
||||
b_qs_(args.b_quantization_strategy),
|
||||
m_size_cache_(nullptr) {
|
||||
assert(a_qs_ != QuantizationStrategy::PER_OUTPUT_CHANNEL);
|
||||
assert(b_qs_ != QuantizationStrategy::PER_TOKEN);
|
||||
if (a_qs_ == QuantizationStrategy::PER_TOKEN) {
|
||||
assert(!use_azp_);
|
||||
};
|
||||
dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
|
||||
{b_k_stride_, b_n_stride_});
|
||||
|
||||
// dummy M size for prepacking weights
|
||||
// Prepacking weights improves performance and avoid runtime reorders
|
||||
constexpr dnnl_dim_t kProbeM = 128;
|
||||
|
||||
prepack_weight(args.b_ptr, original_b_md,
|
||||
create_primitive_desc(
|
||||
MSizeCacheKey{.a_m_size = kProbeM,
|
||||
.use_bias = false,
|
||||
.bias_type = dnnl::memory::data_type::undef},
|
||||
/*first_time=*/true)
|
||||
.weights_desc());
|
||||
init_runtime_memory_cache(args);
|
||||
}
|
||||
|
||||
void W8A8MatMulPrimitiveHandler::execute(ExecArgs& args) {
|
||||
auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0);
|
||||
auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1);
|
||||
a_storage->set_data_handle((void*)args.a_ptr);
|
||||
a_mem_desc->dims[0] = args.a_m_size;
|
||||
c_storage->set_data_handle((void*)args.c_ptr);
|
||||
c_mem_desc->dims[0] = args.a_m_size;
|
||||
|
||||
if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
|
||||
auto&& [a_scale_storage, a_scale_mem_desc] = get_runtime_memory_ptr(2);
|
||||
a_scale_storage->set_data_handle((void*)args.a_scales_ptr);
|
||||
}
|
||||
if (use_azp_) {
|
||||
auto&& [a_zero_point_storage, a_zero_point_mem_desc] =
|
||||
get_runtime_memory_ptr(3);
|
||||
a_zero_point_storage->set_data_handle((void*)args.a_zero_points_ptr);
|
||||
}
|
||||
|
||||
if (args.use_bias) {
|
||||
auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(4);
|
||||
bias_storage->set_data_handle((void*)args.bias_ptr);
|
||||
}
|
||||
|
||||
dnnl::matmul matmul = get_matmul_cache(args);
|
||||
|
||||
auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(5);
|
||||
scratchpad_storage->set_data_handle(
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>());
|
||||
|
||||
matmul.execute(default_stream(), memory_cache_);
|
||||
default_stream().wait();
|
||||
}
|
||||
|
||||
dnnl::matmul W8A8MatMulPrimitiveHandler::get_matmul_cache(
|
||||
const MSizeCacheKey& key) {
|
||||
if (m_size_cache_.get() == nullptr) {
|
||||
ClassMatmulCacheKey key = {.b_n_size = b_n_size_,
|
||||
.b_k_size = b_k_size_,
|
||||
.a_qs = a_qs_,
|
||||
.b_qs = b_qs_,
|
||||
.use_azp = use_azp_,
|
||||
.c_type = c_type_};
|
||||
m_size_cache_ = get_w8a8_class_primitive_cache(key, primitive_cache_size_);
|
||||
}
|
||||
|
||||
return m_size_cache_->get_or_create(key, [&]() {
|
||||
dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
|
||||
auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager();
|
||||
manager->realloc(desc.scratchpad_desc().get_size());
|
||||
return dnnl::matmul(desc);
|
||||
});
|
||||
}
|
||||
|
||||
void W8A8MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
|
||||
memory_cache_[DNNL_ARG_SRC] = dnnl::memory({{1, b_k_size_},
|
||||
dnnl::memory::data_type::s8,
|
||||
dnnl::memory::format_tag::ab},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get());
|
||||
memory_cache_[DNNL_ARG_DST] =
|
||||
dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get());
|
||||
|
||||
// For PER_TOKEN, scales will be applied in outside epilogue
|
||||
if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
|
||||
memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC] = dnnl::memory(
|
||||
{{1}, dnnl::memory::data_type::f32, {1}}, default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(
|
||||
2, memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_SRC].get());
|
||||
if (use_azp_) {
|
||||
memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC] = dnnl::memory(
|
||||
{{1}, dnnl::memory::data_type::s32, {1}}, default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(
|
||||
3, memory_cache_[DNNL_ARG_ATTR_ZERO_POINTS | DNNL_ARG_SRC].get());
|
||||
}
|
||||
}
|
||||
|
||||
if (b_qs_ == QuantizationStrategy::PER_TENSOR) {
|
||||
memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] =
|
||||
dnnl::memory({{1}, dnnl::memory::data_type::f32, {1}}, default_engine(),
|
||||
(void*)args.b_scales_ptr);
|
||||
} else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) {
|
||||
memory_cache_[DNNL_ARG_ATTR_SCALES | DNNL_ARG_WEIGHTS] =
|
||||
dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
|
||||
default_engine(), (void*)args.b_scales_ptr);
|
||||
}
|
||||
|
||||
memory_cache_[DNNL_ARG_BIAS] =
|
||||
dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(4, memory_cache_[DNNL_ARG_BIAS].get());
|
||||
|
||||
memory_cache_[DNNL_ARG_SCRATCHPAD] =
|
||||
dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(5, memory_cache_[DNNL_ARG_SCRATCHPAD].get());
|
||||
}
|
||||
|
||||
dnnl::matmul::primitive_desc W8A8MatMulPrimitiveHandler::create_primitive_desc(
|
||||
const MSizeCacheKey& key, bool first_time) {
|
||||
dnnl::memory::desc a_md({key.a_m_size, b_k_size_},
|
||||
dnnl::memory::data_type::s8,
|
||||
dnnl::memory::format_tag::ab);
|
||||
dnnl::memory::desc b_md;
|
||||
if (first_time) {
|
||||
b_md =
|
||||
dnnl::memory::desc({b_k_size_, b_n_size_}, dnnl::memory::data_type::s8,
|
||||
dnnl::memory::format_tag::any);
|
||||
} else {
|
||||
b_md = b_target_mem_desc_;
|
||||
}
|
||||
dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_,
|
||||
dnnl::memory::format_tag::ab);
|
||||
|
||||
dnnl::primitive_attr attr;
|
||||
|
||||
attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
|
||||
|
||||
// For PER_TOKEN, scales will be applied in outside epilogue
|
||||
if (a_qs_ == QuantizationStrategy::PER_TENSOR) {
|
||||
attr.set_scales_mask(DNNL_ARG_SRC, 0);
|
||||
if (use_azp_) {
|
||||
attr.set_zero_points_mask(DNNL_ARG_SRC, 0);
|
||||
}
|
||||
}
|
||||
|
||||
if (b_qs_ == QuantizationStrategy::PER_TENSOR) {
|
||||
attr.set_scales_mask(DNNL_ARG_WEIGHTS, 0);
|
||||
} else if (b_qs_ == QuantizationStrategy::PER_OUTPUT_CHANNEL) {
|
||||
attr.set_scales_mask(DNNL_ARG_WEIGHTS, 2);
|
||||
}
|
||||
|
||||
if (key.use_bias) {
|
||||
// For PER_TOKEN, bias will be applied in epilogue
|
||||
assert(a_qs_ == QuantizationStrategy::PER_TENSOR);
|
||||
dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1});
|
||||
return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md,
|
||||
c_md, attr);
|
||||
} else {
|
||||
return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
|
||||
attr);
|
||||
}
|
||||
}
|
||||
|
||||
MatMulPrimitiveHandler::MatMulPrimitiveHandler(const Args& args)
|
||||
: DNNLMatMulPrimitiveHandler(
|
||||
static_cast<DNNLMatMulPrimitiveHandler::Args>(args), args.ab_type),
|
||||
m_size_cache_(nullptr) {
|
||||
assert(b_type_ == dnnl::memory::data_type::f32 ||
|
||||
b_type_ == dnnl::memory::data_type::bf16 ||
|
||||
b_type_ == dnnl::memory::data_type::f16);
|
||||
|
||||
dnnl::memory::desc original_b_md({b_k_size_, b_n_size_}, b_type_,
|
||||
{b_k_stride_, b_n_stride_});
|
||||
|
||||
// dummy M size for prepacking weights
|
||||
// Prepacking weights improves performance and avoid runtime reorders
|
||||
constexpr dnnl_dim_t kProbeM = 128;
|
||||
|
||||
prepack_weight(args.b_ptr, original_b_md,
|
||||
create_primitive_desc(
|
||||
MSizeCacheKey{// Use a concrete M so oneDNN's kernel
|
||||
// selector can choose an optimally blocked
|
||||
// weight layout.
|
||||
.a_m_size = kProbeM,
|
||||
.a_m_stride = b_k_size_,
|
||||
.use_bias = false,
|
||||
.bias_type = dnnl::memory::data_type::undef},
|
||||
true)
|
||||
.weights_desc());
|
||||
init_runtime_memory_cache(args);
|
||||
}
|
||||
|
||||
static std::shared_ptr<MatMulPrimitiveHandler::MSizeCache>
|
||||
get_matul_class_primitive_cache(
|
||||
const MatMulPrimitiveHandler::ClassMatmulCacheKey& key,
|
||||
int64_t cache_size) {
|
||||
static MatMulPrimitiveHandler::ClassMatmulCache cache(128);
|
||||
assert(cache_size > 0);
|
||||
return cache.get_or_create(key, [&]() {
|
||||
return std::make_shared<MatMulPrimitiveHandler::MSizeCache>(cache_size);
|
||||
});
|
||||
}
|
||||
|
||||
void MatMulPrimitiveHandler::execute(ExecArgs& args) {
|
||||
auto&& [a_storage, a_mem_desc] = get_runtime_memory_ptr(0);
|
||||
auto&& [c_storage, c_mem_desc] = get_runtime_memory_ptr(1);
|
||||
a_storage->set_data_handle((void*)args.a_ptr);
|
||||
a_mem_desc->dims[0] = args.a_m_size;
|
||||
a_mem_desc->format_desc.blocking.strides[0] = args.a_m_stride;
|
||||
c_storage->set_data_handle((void*)args.c_ptr);
|
||||
c_mem_desc->dims[0] = args.a_m_size;
|
||||
|
||||
#ifndef VLLM_USE_ACL
|
||||
// We do not support in ACL backend of oneDNN, we handle bias by:
|
||||
// 1. copying it into the result tensor
|
||||
// 2. attaching a fused-sum post-op to the matmul primitive
|
||||
if (args.use_bias) {
|
||||
auto&& [bias_storage, bias_mem_desc] = get_runtime_memory_ptr(2);
|
||||
bias_storage->set_data_handle((void*)args.bias_ptr);
|
||||
}
|
||||
#endif
|
||||
dnnl::matmul matmul = get_matmul_cache(args);
|
||||
|
||||
// With ACL backend of oneDNN, the required memory format might change when the
|
||||
// source tensor dims change. This does not really happen in practice, so isn't
|
||||
// a performance hit, but we need to support it because the API allows for it.
|
||||
#ifdef VLLM_USE_ACL
|
||||
auto new_expected_wei_desc =
|
||||
dnnl::matmul::primitive_desc(
|
||||
const_cast<dnnl_primitive_desc_t>(matmul.get_primitive_desc()))
|
||||
.weights_desc();
|
||||
if (new_expected_wei_desc != b_target_mem_desc_) {
|
||||
prepack_weight(memory_cache_[DNNL_ARG_WEIGHTS].get_data_handle(),
|
||||
b_target_mem_desc_, new_expected_wei_desc);
|
||||
}
|
||||
#endif
|
||||
|
||||
auto&& [scratchpad_storage, scratchpad_mem_desc] = get_runtime_memory_ptr(3);
|
||||
scratchpad_storage->set_data_handle(
|
||||
cpu_utils::ScratchPadManager::get_scratchpad_manager()->get_data<void>());
|
||||
|
||||
matmul.execute(default_stream(), memory_cache_);
|
||||
default_stream().wait();
|
||||
}
|
||||
|
||||
dnnl::matmul MatMulPrimitiveHandler::get_matmul_cache(
|
||||
const MSizeCacheKey& key) {
|
||||
if (m_size_cache_.get() == nullptr) {
|
||||
ClassMatmulCacheKey class_key = {
|
||||
.b_n_size = b_n_size_, .b_k_size = b_k_size_, .b_type = b_type_};
|
||||
m_size_cache_ =
|
||||
get_matul_class_primitive_cache(class_key, primitive_cache_size_);
|
||||
}
|
||||
return m_size_cache_->get_or_create(key, [&]() {
|
||||
dnnl::matmul::primitive_desc desc = this->create_primitive_desc(key, false);
|
||||
auto manager = cpu_utils::ScratchPadManager::get_scratchpad_manager();
|
||||
manager->realloc(desc.scratchpad_desc().get_size());
|
||||
return dnnl::matmul(desc);
|
||||
});
|
||||
}
|
||||
|
||||
dnnl::matmul::primitive_desc MatMulPrimitiveHandler::create_primitive_desc(
|
||||
const MSizeCacheKey& key, bool first_time) {
|
||||
dnnl::memory::desc a_md;
|
||||
dnnl::memory::desc b_md;
|
||||
if (first_time) {
|
||||
a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_,
|
||||
dnnl::memory::format_tag::ab);
|
||||
b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_,
|
||||
dnnl::memory::format_tag::any);
|
||||
} else {
|
||||
a_md = dnnl::memory::desc({key.a_m_size, b_k_size_}, b_type_,
|
||||
{key.a_m_stride, 1});
|
||||
#ifdef VLLM_USE_ACL
|
||||
// ACL's backend of oneDNN always expects the weight format to be "any"
|
||||
b_md = dnnl::memory::desc({b_k_size_, b_n_size_}, b_type_,
|
||||
dnnl::memory::format_tag::any);
|
||||
#else
|
||||
b_md = b_target_mem_desc_;
|
||||
#endif
|
||||
}
|
||||
dnnl::memory::desc c_md({key.a_m_size, b_n_size_}, c_type_,
|
||||
dnnl::memory::format_tag::ab);
|
||||
|
||||
dnnl::primitive_attr attr;
|
||||
attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
|
||||
|
||||
if (key.use_bias) {
|
||||
dnnl::memory::desc bias_md({1, b_n_size_}, key.bias_type, {b_n_size_, 1});
|
||||
// Since ACL's matmuls don't support passing a bias_md, we apply the bias
|
||||
// through a fused-sum post-op
|
||||
#ifdef VLLM_USE_ACL
|
||||
dnnl::post_ops post_ops;
|
||||
post_ops.append_sum();
|
||||
attr.set_post_ops(post_ops);
|
||||
return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
|
||||
attr);
|
||||
#else
|
||||
return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, bias_md,
|
||||
c_md, attr);
|
||||
#endif
|
||||
} else {
|
||||
return dnnl::matmul::primitive_desc(default_engine(), a_md, b_md, c_md,
|
||||
attr);
|
||||
}
|
||||
}
|
||||
|
||||
void MatMulPrimitiveHandler::init_runtime_memory_cache(const Args& args) {
|
||||
memory_cache_[DNNL_ARG_SRC] = dnnl::memory(
|
||||
{{1, b_k_size_}, b_type_, {b_k_size_, 1}}, default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(0, memory_cache_[DNNL_ARG_SRC].get());
|
||||
memory_cache_[DNNL_ARG_DST] =
|
||||
dnnl::memory({{1, b_n_size_}, c_type_, dnnl::memory::format_tag::ab},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(1, memory_cache_[DNNL_ARG_DST].get());
|
||||
|
||||
// ACL matmuls don't support bias_md, so we don't need these
|
||||
#ifndef VLLM_USE_ACL
|
||||
memory_cache_[DNNL_ARG_BIAS] =
|
||||
dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(2, memory_cache_[DNNL_ARG_BIAS].get());
|
||||
#endif
|
||||
memory_cache_[DNNL_ARG_SCRATCHPAD] =
|
||||
dnnl::memory({{b_n_size_}, dnnl::memory::data_type::f32, {1}},
|
||||
default_engine(), nullptr);
|
||||
set_runtime_memory_ptr(3, memory_cache_[DNNL_ARG_SCRATCHPAD].get());
|
||||
}
|
||||
|
||||
bool is_onednn_acl_supported() {
|
||||
#ifdef VLLM_USE_ACL
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
#ifndef DNNL_HELPER_H
|
||||
#define DNNL_HELPER_H
|
||||
|
||||
#include <optional>
|
||||
#include <cassert>
|
||||
|
||||
#include "oneapi/dnnl/dnnl.hpp"
|
||||
|
||||
namespace c10 {
|
||||
struct BFloat16;
|
||||
struct Half;
|
||||
} // namespace c10
|
||||
|
||||
namespace dnnl {
|
||||
namespace impl {
|
||||
struct memory_storage_t;
|
||||
struct matmul_pd_t;
|
||||
struct matmul_desc_t;
|
||||
} // namespace impl
|
||||
} // namespace dnnl
|
||||
struct dnnl_memory_desc;
|
||||
|
||||
template <typename KT, typename VT>
|
||||
class DNNLPrimitiveCache;
|
||||
|
||||
template <typename T>
|
||||
struct DNNLType {
|
||||
static constexpr dnnl::memory::data_type type =
|
||||
dnnl::memory::data_type::undef;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DNNLType<int8_t> {
|
||||
static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s8;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DNNLType<int32_t> {
|
||||
static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::s32;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DNNLType<float> {
|
||||
static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f32;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DNNLType<c10::BFloat16> {
|
||||
static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::bf16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DNNLType<c10::Half> {
|
||||
static constexpr dnnl::memory::data_type type = dnnl::memory::data_type::f16;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
constexpr inline dnnl::memory::data_type get_dnnl_type() {
|
||||
return DNNLType<std::decay_t<T>>::type;
|
||||
}
|
||||
|
||||
class DNNLMatMulPrimitiveHandler {
|
||||
public:
|
||||
virtual ~DNNLMatMulPrimitiveHandler() = default;
|
||||
|
||||
protected:
|
||||
struct Args {
|
||||
dnnl_dim_t b_n_size;
|
||||
dnnl_dim_t b_n_stride;
|
||||
dnnl_dim_t b_k_size;
|
||||
dnnl_dim_t b_k_stride;
|
||||
void* b_ptr;
|
||||
dnnl::memory::data_type c_type;
|
||||
size_t primitive_cache_size;
|
||||
};
|
||||
|
||||
protected:
|
||||
DNNLMatMulPrimitiveHandler(const Args& args, dnnl::memory::data_type b_type);
|
||||
|
||||
void prepack_weight(void* original_b_ptr, dnnl::memory::desc original_b_md,
|
||||
dnnl::memory::desc b_target_mem_desc);
|
||||
|
||||
void set_runtime_memory_ptr(size_t index, dnnl_memory* memory_ptr);
|
||||
|
||||
std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>
|
||||
get_runtime_memory_ptr(size_t index);
|
||||
|
||||
protected:
|
||||
const dnnl_dim_t b_n_size_;
|
||||
const dnnl_dim_t b_n_stride_;
|
||||
const dnnl_dim_t b_k_size_;
|
||||
const dnnl_dim_t b_k_stride_;
|
||||
dnnl::memory::data_type b_type_;
|
||||
dnnl::memory::data_type c_type_;
|
||||
std::unordered_map<int, dnnl::memory> memory_cache_;
|
||||
std::vector<std::pair<dnnl::impl::memory_storage_t*, dnnl_memory_desc*>>
|
||||
runtime_memory_ptrs_;
|
||||
dnnl::memory::desc b_target_mem_desc_;
|
||||
int64_t primitive_cache_size_;
|
||||
};
|
||||
|
||||
class W8A8MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
|
||||
public:
|
||||
enum class QuantizationStrategy { PER_TOKEN, PER_TENSOR, PER_OUTPUT_CHANNEL };
|
||||
|
||||
struct Args : public DNNLMatMulPrimitiveHandler::Args {
|
||||
bool use_a_zero_point;
|
||||
QuantizationStrategy a_quantization_strategy;
|
||||
QuantizationStrategy b_quantization_strategy;
|
||||
float* b_scales_ptr;
|
||||
};
|
||||
|
||||
struct ClassMatmulCacheKey {
|
||||
dnnl_dim_t b_n_size;
|
||||
dnnl_dim_t b_k_size;
|
||||
QuantizationStrategy a_qs;
|
||||
QuantizationStrategy b_qs;
|
||||
bool use_azp;
|
||||
dnnl::memory::data_type c_type;
|
||||
|
||||
friend bool operator==(const ClassMatmulCacheKey& l,
|
||||
const ClassMatmulCacheKey& r);
|
||||
};
|
||||
|
||||
struct MSizeCacheKey {
|
||||
dnnl_dim_t a_m_size;
|
||||
bool use_bias;
|
||||
dnnl::memory::data_type bias_type;
|
||||
|
||||
friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
|
||||
};
|
||||
|
||||
using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
|
||||
using ClassMatmulCache =
|
||||
DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
|
||||
|
||||
struct ExecArgs : public MSizeCacheKey {
|
||||
const int8_t* a_ptr;
|
||||
const float* a_scales_ptr;
|
||||
const int32_t* a_zero_points_ptr;
|
||||
const void* bias_ptr;
|
||||
void* c_ptr;
|
||||
};
|
||||
|
||||
public:
|
||||
W8A8MatMulPrimitiveHandler(const Args& args);
|
||||
|
||||
QuantizationStrategy get_input_scale_strategy() const { return a_qs_; }
|
||||
|
||||
bool get_input_use_zero_point() const { return use_azp_; }
|
||||
|
||||
void execute(ExecArgs& args);
|
||||
|
||||
private:
|
||||
dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
|
||||
bool first_time);
|
||||
|
||||
void init_runtime_memory_cache(const Args& args);
|
||||
|
||||
dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
|
||||
|
||||
private:
|
||||
const bool use_azp_;
|
||||
const QuantizationStrategy a_qs_;
|
||||
const QuantizationStrategy b_qs_;
|
||||
std::shared_ptr<MSizeCache> m_size_cache_;
|
||||
};
|
||||
|
||||
class MatMulPrimitiveHandler : public DNNLMatMulPrimitiveHandler {
|
||||
public:
|
||||
struct Args : public DNNLMatMulPrimitiveHandler::Args {
|
||||
dnnl::memory::data_type ab_type;
|
||||
};
|
||||
|
||||
struct ClassMatmulCacheKey {
|
||||
dnnl_dim_t b_n_size;
|
||||
dnnl_dim_t b_k_size;
|
||||
dnnl::memory::data_type b_type;
|
||||
|
||||
friend bool operator==(const ClassMatmulCacheKey& l,
|
||||
const ClassMatmulCacheKey& r);
|
||||
};
|
||||
|
||||
struct MSizeCacheKey {
|
||||
dnnl_dim_t a_m_size;
|
||||
dnnl_dim_t a_m_stride;
|
||||
bool use_bias;
|
||||
dnnl::memory::data_type bias_type;
|
||||
|
||||
friend bool operator==(const MSizeCacheKey& l, const MSizeCacheKey& r);
|
||||
};
|
||||
|
||||
using MSizeCache = DNNLPrimitiveCache<MSizeCacheKey, dnnl::matmul>;
|
||||
using ClassMatmulCache =
|
||||
DNNLPrimitiveCache<ClassMatmulCacheKey, std::shared_ptr<MSizeCache>>;
|
||||
|
||||
struct ExecArgs : public MSizeCacheKey {
|
||||
const void* a_ptr;
|
||||
const void* bias_ptr;
|
||||
void* c_ptr;
|
||||
};
|
||||
|
||||
public:
|
||||
MatMulPrimitiveHandler(const Args& args);
|
||||
|
||||
void execute(ExecArgs& args);
|
||||
|
||||
private:
|
||||
dnnl::matmul::primitive_desc create_primitive_desc(const MSizeCacheKey& key,
|
||||
bool first_time);
|
||||
|
||||
void init_runtime_memory_cache(const Args& args);
|
||||
|
||||
dnnl::matmul get_matmul_cache(const MSizeCacheKey& key);
|
||||
|
||||
private:
|
||||
std::shared_ptr<MSizeCache> m_size_cache_;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,570 @@
|
||||
#include "cpu_types.hpp"
|
||||
#include "dnnl_helper.h"
|
||||
|
||||
namespace {
|
||||
template <typename scalar_t>
|
||||
struct KernelVecType {
|
||||
using load_vec_type = void;
|
||||
using cvt_vec_type = void;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<float> {
|
||||
using load_vec_type = vec_op::FP32Vec16;
|
||||
using cvt_vec_type = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<c10::BFloat16> {
|
||||
using load_vec_type = vec_op::BF16Vec16;
|
||||
using cvt_vec_type = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<c10::Half> {
|
||||
#if defined(__powerpc64__) || defined(__s390x__)
|
||||
// Power architecture-specific vector type
|
||||
using load_vec_type = vec_op::FP32Vec16;
|
||||
#else
|
||||
// Fallback for other architectures
|
||||
using load_vec_type = vec_op::FP16Vec16;
|
||||
#endif
|
||||
using cvt_vec_type = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <bool AZP, typename scalar_t>
|
||||
void static_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
|
||||
const float* scale, const int32_t* azp,
|
||||
const int64_t num_tokens,
|
||||
const int64_t input_stride,
|
||||
const int64_t hidden_size) {
|
||||
using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
|
||||
using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
|
||||
constexpr int64_t vec_elem_num = load_vec_t::VEC_ELEM_NUM;
|
||||
|
||||
constexpr float i8_min =
|
||||
static_cast<float>(std::numeric_limits<int8_t>::min());
|
||||
constexpr float i8_max =
|
||||
static_cast<float>(std::numeric_limits<int8_t>::max());
|
||||
const cvt_vec_t inv_scale(1.0 / *scale);
|
||||
const cvt_vec_t i8_min_vec(i8_min);
|
||||
const cvt_vec_t i8_max_vec(i8_max);
|
||||
|
||||
cvt_vec_t zp_vec;
|
||||
if constexpr (AZP) {
|
||||
zp_vec = cvt_vec_t(static_cast<float>(*azp));
|
||||
}
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < num_tokens; ++i) {
|
||||
int64_t j = 0;
|
||||
const scalar_t* input_ptr = input + i * input_stride;
|
||||
int8_t* output_ptr = output + i * hidden_size;
|
||||
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
elems_fp32 = elems_fp32 * inv_scale;
|
||||
|
||||
if constexpr (AZP) {
|
||||
elems_fp32 = elems_fp32 + zp_vec;
|
||||
}
|
||||
|
||||
elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
|
||||
vec_op::INT8Vec16 elems_int8(elems_fp32);
|
||||
elems_int8.save(output_ptr + j);
|
||||
}
|
||||
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
elems_fp32 = elems_fp32 * inv_scale;
|
||||
|
||||
if constexpr (AZP) {
|
||||
elems_fp32 = elems_fp32 + zp_vec;
|
||||
}
|
||||
|
||||
elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
|
||||
vec_op::INT8Vec16 elems_int8(elems_fp32);
|
||||
elems_int8.save(output_ptr + j, hidden_size - j);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool AZP, typename scalar_t>
|
||||
void dynamic_scaled_int8_quant_impl(const scalar_t* input, int8_t* output,
|
||||
float* scale, int32_t* azp,
|
||||
const int64_t num_tokens,
|
||||
const int64_t input_stride,
|
||||
const int64_t hidden_size) {
|
||||
using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
|
||||
using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
|
||||
constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
|
||||
|
||||
constexpr float i8_min =
|
||||
static_cast<float>(std::numeric_limits<int8_t>::min());
|
||||
constexpr float i8_max =
|
||||
static_cast<float>(std::numeric_limits<int8_t>::max());
|
||||
const cvt_vec_t i8_min_vec(i8_min);
|
||||
const cvt_vec_t i8_max_vec(i8_max);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < num_tokens; ++i) {
|
||||
cvt_vec_t max_value(std::numeric_limits<float>::lowest());
|
||||
cvt_vec_t min_value(std::numeric_limits<float>::max());
|
||||
{
|
||||
int64_t j = 0;
|
||||
const scalar_t* input_ptr = input + i * input_stride;
|
||||
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
if constexpr (AZP) {
|
||||
max_value = max_value.max(elems_fp32);
|
||||
min_value = min_value.min(elems_fp32);
|
||||
} else {
|
||||
max_value = max_value.max(elems_fp32.abs());
|
||||
}
|
||||
}
|
||||
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
|
||||
if (j + vec_elem_num == hidden_size) {
|
||||
if constexpr (AZP) {
|
||||
max_value = max_value.max(elems_fp32);
|
||||
min_value = min_value.min(elems_fp32);
|
||||
} else {
|
||||
max_value = max_value.max(elems_fp32.abs());
|
||||
}
|
||||
} else {
|
||||
if constexpr (AZP) {
|
||||
max_value = max_value.max(elems_fp32, hidden_size - j);
|
||||
min_value = min_value.min(elems_fp32, hidden_size - j);
|
||||
} else {
|
||||
max_value = max_value.max(elems_fp32.abs(), hidden_size - j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
float scale_val;
|
||||
float azp_val = 0.0f;
|
||||
if constexpr (AZP) {
|
||||
float max_scalar = max_value.reduce_max();
|
||||
float min_scalar = min_value.reduce_min();
|
||||
scale_val = (max_scalar - min_scalar) / 255.0f;
|
||||
azp_val = std::nearbyint(-128.0f - min_scalar / scale_val);
|
||||
azp[i] = azp_val;
|
||||
scale[i] = scale_val;
|
||||
} else {
|
||||
scale_val = max_value.reduce_max() / 127.0f;
|
||||
scale[i] = scale_val;
|
||||
}
|
||||
|
||||
const cvt_vec_t inv_scale(1.0 / scale_val);
|
||||
const cvt_vec_t azp_vec(azp_val);
|
||||
|
||||
{
|
||||
int64_t j = 0;
|
||||
const scalar_t* input_ptr = input + i * input_stride;
|
||||
int8_t* output_ptr = output + i * hidden_size;
|
||||
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
elems_fp32 = (elems_fp32 * inv_scale);
|
||||
|
||||
if constexpr (AZP) {
|
||||
elems_fp32 = elems_fp32 + azp_vec;
|
||||
}
|
||||
elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
|
||||
vec_op::INT8Vec16 elems_int8(elems_fp32);
|
||||
elems_int8.save(output_ptr + j);
|
||||
}
|
||||
|
||||
load_vec_t elems(input_ptr + j);
|
||||
cvt_vec_t elems_fp32(elems);
|
||||
elems_fp32 = (elems_fp32 * inv_scale);
|
||||
|
||||
if constexpr (AZP) {
|
||||
elems_fp32 = elems_fp32 + azp_vec;
|
||||
}
|
||||
elems_fp32 = elems_fp32.clamp(i8_min_vec, i8_max_vec);
|
||||
vec_op::INT8Vec16 elems_int8(elems_fp32);
|
||||
elems_int8.save(output_ptr + j, hidden_size - j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <bool AZP, bool Bias, typename scalar_t>
|
||||
void dynamic_quant_epilogue(const float* input, scalar_t* output,
|
||||
const float* a_scale, const int32_t* azp,
|
||||
const float* azp_adj, const scalar_t* bias,
|
||||
const int64_t num_tokens,
|
||||
const int64_t hidden_size) {
|
||||
CPU_KERNEL_GUARD_IN(dynamic_quant_epilogue)
|
||||
using load_vec_t = typename KernelVecType<scalar_t>::load_vec_type;
|
||||
using cvt_vec_t = typename KernelVecType<scalar_t>::cvt_vec_type;
|
||||
constexpr int vec_elem_num = load_vec_t::VEC_ELEM_NUM;
|
||||
|
||||
const int64_t thread_num = cpu_utils::get_max_threads();
|
||||
if (num_tokens > thread_num) {
|
||||
#pragma omp parallel for
|
||||
for (int64_t i = 0; i < num_tokens; ++i) {
|
||||
const float* input_ptr = input + i * hidden_size;
|
||||
scalar_t* output_ptr = output + i * hidden_size;
|
||||
int64_t j = 0;
|
||||
cvt_vec_t token_scale_vec(a_scale[i]);
|
||||
cvt_vec_t token_zp_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
|
||||
token_zp_scale_vec = cvt_vec_t(zp_scale_val);
|
||||
}
|
||||
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
|
||||
cvt_vec_t elems_fp32(input_ptr + j);
|
||||
elems_fp32 = elems_fp32 * token_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
cvt_vec_t azp_adj_fp32(azp_adj + j);
|
||||
elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
|
||||
}
|
||||
if constexpr (Bias) {
|
||||
load_vec_t bias_vec(bias + j);
|
||||
cvt_vec_t bias_vec_fp32(bias_vec);
|
||||
elems_fp32 = elems_fp32 + bias_vec_fp32;
|
||||
}
|
||||
load_vec_t elems_out(elems_fp32);
|
||||
elems_out.save(output_ptr + j);
|
||||
}
|
||||
cvt_vec_t elems_fp32(input_ptr + j);
|
||||
elems_fp32 = elems_fp32 * token_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
cvt_vec_t azp_adj_fp32(azp_adj + j);
|
||||
elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
|
||||
}
|
||||
if constexpr (Bias) {
|
||||
load_vec_t bias_vec(bias + j);
|
||||
cvt_vec_t bias_vec_fp32(bias_vec);
|
||||
elems_fp32 = elems_fp32 + bias_vec_fp32;
|
||||
}
|
||||
load_vec_t elems_out(elems_fp32);
|
||||
elems_out.save(output_ptr + j, hidden_size - j);
|
||||
}
|
||||
} else {
|
||||
const int64_t vec_iteration =
|
||||
(hidden_size + vec_elem_num - 1) / vec_elem_num;
|
||||
const int64_t vec_iteration_per_thread =
|
||||
(vec_iteration + thread_num - 1) / thread_num;
|
||||
const int64_t elem_num_per_thread = vec_iteration_per_thread * vec_elem_num;
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int64_t i = 0; i < thread_num; ++i) {
|
||||
const int64_t start = elem_num_per_thread * i;
|
||||
const int64_t end = std::min(hidden_size, elem_num_per_thread + start);
|
||||
for (int64_t j = 0; j < num_tokens; ++j) {
|
||||
cvt_vec_t token_scale_vec(a_scale[j]);
|
||||
cvt_vec_t token_zp_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
float zp_scale_val = a_scale[j] * static_cast<float>(azp[j]);
|
||||
token_zp_scale_vec = cvt_vec_t(zp_scale_val);
|
||||
}
|
||||
int64_t k = start;
|
||||
const float* input_ptr = input + j * hidden_size;
|
||||
scalar_t* output_ptr = output + j * hidden_size;
|
||||
for (; k < end - vec_elem_num; k += vec_elem_num) {
|
||||
cvt_vec_t elems_fp32(input_ptr + k);
|
||||
elems_fp32 = elems_fp32 * token_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
cvt_vec_t azp_adj_fp32(azp_adj + k);
|
||||
elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
|
||||
}
|
||||
if constexpr (Bias) {
|
||||
load_vec_t bias_vec(bias + k);
|
||||
cvt_vec_t bias_vec_fp32(bias_vec);
|
||||
elems_fp32 = elems_fp32 + bias_vec_fp32;
|
||||
}
|
||||
load_vec_t elems_out(elems_fp32);
|
||||
elems_out.save(output_ptr + k);
|
||||
}
|
||||
if (k < end) {
|
||||
cvt_vec_t elems_fp32(input_ptr + k);
|
||||
elems_fp32 = elems_fp32 * token_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
cvt_vec_t azp_adj_fp32(azp_adj + k);
|
||||
elems_fp32 = elems_fp32 - azp_adj_fp32 * token_zp_scale_vec;
|
||||
}
|
||||
if constexpr (Bias) {
|
||||
load_vec_t bias_vec(bias + k);
|
||||
cvt_vec_t bias_vec_fp32(bias_vec);
|
||||
elems_fp32 = elems_fp32 + bias_vec_fp32;
|
||||
}
|
||||
load_vec_t elems_out(elems_fp32);
|
||||
elems_out.save(output_ptr + k, end - k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
int64_t create_onednn_scaled_mm_handler(
|
||||
const torch::Tensor& b, // [IC, OC], column-major
|
||||
const torch::Tensor& b_scales, // [1] or [OC]
|
||||
at::ScalarType output_type, bool dynamic_act_quant, bool use_azp,
|
||||
int64_t primitive_cache_size) {
|
||||
TORCH_CHECK(b.dim() == 2);
|
||||
TORCH_CHECK(b.stride(0) == 1); // Column-major
|
||||
TORCH_CHECK(b_scales.is_contiguous());
|
||||
|
||||
W8A8MatMulPrimitiveHandler::Args args;
|
||||
args.primitive_cache_size = primitive_cache_size;
|
||||
|
||||
if (b_scales.numel() == 1) {
|
||||
args.b_quantization_strategy =
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR;
|
||||
} else {
|
||||
TORCH_CHECK_EQ(b_scales.numel(), b.size(1));
|
||||
args.b_quantization_strategy =
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_OUTPUT_CHANNEL;
|
||||
}
|
||||
args.b_scales_ptr = b_scales.data_ptr<float>();
|
||||
args.b_k_size = b.size(0);
|
||||
args.b_k_stride = b.stride(0);
|
||||
args.b_n_size = b.size(1);
|
||||
args.b_n_stride = b.stride(1);
|
||||
args.b_ptr = b.data_ptr<int8_t>();
|
||||
|
||||
if (dynamic_act_quant) {
|
||||
// dynamic per-token, bias, A scales and A zps will be applied in outside.
|
||||
args.a_quantization_strategy =
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN;
|
||||
args.use_a_zero_point = false;
|
||||
} else {
|
||||
// static per-tensor
|
||||
args.a_quantization_strategy =
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR;
|
||||
args.use_a_zero_point = use_azp;
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(output_type, "create_onednn_scaled_mm_handler",
|
||||
[&] {
|
||||
if (dynamic_act_quant) {
|
||||
args.c_type = get_dnnl_type<float>();
|
||||
} else {
|
||||
args.c_type = get_dnnl_type<scalar_t>();
|
||||
}
|
||||
});
|
||||
|
||||
return reinterpret_cast<int64_t>(new W8A8MatMulPrimitiveHandler(args));
|
||||
}
|
||||
|
||||
void onednn_scaled_mm(
|
||||
torch::Tensor& c, // [M, OC], row-major
|
||||
const torch::Tensor& a, // [M, IC], row-major
|
||||
const torch::Tensor& a_scales, // [M] or [1]
|
||||
const std::optional<torch::Tensor>& azp, // [M] or [1]
|
||||
const std::optional<torch::Tensor>& azp_adj, // [M] or [1]
|
||||
const std::optional<torch::Tensor>& bias, // [N]
|
||||
const torch::Tensor& handler_tensor) {
|
||||
CPU_KERNEL_GUARD_IN(onednn_scaled_mm)
|
||||
TORCH_CHECK(a.dim() == 2);
|
||||
TORCH_CHECK(a.is_contiguous());
|
||||
TORCH_CHECK(c.is_contiguous());
|
||||
W8A8MatMulPrimitiveHandler* ptr =
|
||||
reinterpret_cast<W8A8MatMulPrimitiveHandler*>(
|
||||
handler_tensor.item<int64_t>());
|
||||
const int32_t* azp_ptr = nullptr;
|
||||
if (azp.has_value()) {
|
||||
azp_ptr = azp->data_ptr<int32_t>();
|
||||
}
|
||||
if (ptr->get_input_scale_strategy() ==
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) {
|
||||
TORCH_CHECK_EQ(a_scales.numel(), 1);
|
||||
}
|
||||
|
||||
W8A8MatMulPrimitiveHandler::ExecArgs exec_args;
|
||||
exec_args.a_ptr = a.data_ptr<int8_t>();
|
||||
exec_args.a_m_size = a.size(0);
|
||||
exec_args.bias_ptr = nullptr;
|
||||
exec_args.bias_type = get_dnnl_type<void>();
|
||||
exec_args.use_bias = false;
|
||||
exec_args.a_scales_ptr = nullptr;
|
||||
exec_args.a_zero_points_ptr = nullptr;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(c.scalar_type(), "onednn_scaled_mm", [&] {
|
||||
if (ptr->get_input_scale_strategy() ==
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TENSOR) {
|
||||
if (bias.has_value()) {
|
||||
exec_args.bias_ptr = bias->data_ptr<scalar_t>();
|
||||
exec_args.bias_type = get_dnnl_type<scalar_t>();
|
||||
exec_args.use_bias = true;
|
||||
}
|
||||
exec_args.a_scales_ptr = a_scales.data_ptr<float>();
|
||||
exec_args.a_zero_points_ptr = azp_ptr;
|
||||
exec_args.c_ptr = c.data_ptr<scalar_t>();
|
||||
ptr->execute(exec_args);
|
||||
} else if (ptr->get_input_scale_strategy() ==
|
||||
W8A8MatMulPrimitiveHandler::QuantizationStrategy::PER_TOKEN) {
|
||||
torch::Tensor tmp_fp32_out =
|
||||
torch::empty_like(c, ::at::ScalarType::Float);
|
||||
exec_args.c_ptr = tmp_fp32_out.data_ptr<float>();
|
||||
ptr->execute(exec_args);
|
||||
if (bias.has_value()) {
|
||||
if (azp.has_value()) {
|
||||
dynamic_quant_epilogue<true, true>(
|
||||
tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
|
||||
a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(),
|
||||
bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
|
||||
} else {
|
||||
dynamic_quant_epilogue<false, true>(
|
||||
tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
|
||||
a_scales.data_ptr<float>(), azp_ptr, nullptr,
|
||||
bias->data_ptr<scalar_t>(), c.size(0), c.size(1));
|
||||
}
|
||||
} else {
|
||||
if (azp.has_value()) {
|
||||
dynamic_quant_epilogue<true, false>(
|
||||
tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
|
||||
a_scales.data_ptr<float>(), azp_ptr, azp_adj->data_ptr<float>(),
|
||||
(scalar_t*)nullptr, c.size(0), c.size(1));
|
||||
} else {
|
||||
dynamic_quant_epilogue<false, false>(
|
||||
tmp_fp32_out.data_ptr<float>(), c.data_ptr<scalar_t>(),
|
||||
a_scales.data_ptr<float>(), azp_ptr, nullptr, (scalar_t*)nullptr,
|
||||
c.size(0), c.size(1));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
TORCH_CHECK(false, "invalid act quant type.");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// static-per-tensor quantization.
|
||||
void static_scaled_int8_quant(
|
||||
torch::Tensor& out, // [batch, hidden_size]
|
||||
const torch::Tensor& input, // [batch, hidden_size]
|
||||
const torch::Tensor& scale, std::optional<torch::Tensor> const& azp) {
|
||||
CPU_KERNEL_GUARD_IN(static_scaled_int8_quant)
|
||||
TORCH_CHECK(out.is_contiguous());
|
||||
TORCH_CHECK_EQ(input.dim(), 2);
|
||||
TORCH_CHECK_EQ(input.stride(1), 1);
|
||||
TORCH_CHECK(scale.numel() == 1);
|
||||
TORCH_CHECK(!azp.has_value() || azp->numel() == 1);
|
||||
|
||||
const int64_t stride = input.stride(0);
|
||||
const int64_t hidden_size = input.size(1);
|
||||
const int64_t num_tokens = input.size(0);
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "static_scaled_int8_quant_impl", [&] {
|
||||
if (azp.has_value()) {
|
||||
static_scaled_int8_quant_impl<true>(
|
||||
input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
|
||||
scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
|
||||
stride, hidden_size);
|
||||
} else {
|
||||
static_scaled_int8_quant_impl<false>(input.data_ptr<scalar_t>(),
|
||||
out.data_ptr<int8_t>(),
|
||||
scale.data_ptr<float>(), nullptr,
|
||||
num_tokens, stride, hidden_size);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// dynamic-per-token quantization.
|
||||
void dynamic_scaled_int8_quant(
|
||||
torch::Tensor& out, // [batch, hidden_size]
|
||||
const torch::Tensor& input, // [batch, hidden_size]
|
||||
torch::Tensor& scale, // [batch, 1]
|
||||
std::optional<torch::Tensor> const& azp) {
|
||||
CPU_KERNEL_GUARD_IN(dynamic_scaled_int8_quant)
|
||||
TORCH_CHECK(out.is_contiguous());
|
||||
TORCH_CHECK_EQ(input.dim(), 2);
|
||||
TORCH_CHECK_EQ(input.stride(1), 1);
|
||||
|
||||
const int64_t hidden_size = input.size(1);
|
||||
const int64_t num_tokens = input.size(0);
|
||||
const int64_t stride = input.stride(0);
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "dynamic_scaled_int8_quant_impl", [&] {
|
||||
if (azp.has_value()) {
|
||||
dynamic_scaled_int8_quant_impl<true>(
|
||||
input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
|
||||
scale.data_ptr<float>(), azp->data_ptr<int32_t>(), num_tokens,
|
||||
stride, hidden_size);
|
||||
} else {
|
||||
dynamic_scaled_int8_quant_impl<false>(
|
||||
input.data_ptr<scalar_t>(), out.data_ptr<int8_t>(),
|
||||
scale.data_ptr<float>(), nullptr, num_tokens, stride,
|
||||
hidden_size);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
int64_t create_onednn_mm_handler(const torch::Tensor& b,
|
||||
int64_t primitive_cache_size) {
|
||||
TORCH_CHECK(b.dim() == 2);
|
||||
|
||||
MatMulPrimitiveHandler::Args args;
|
||||
args.primitive_cache_size = primitive_cache_size;
|
||||
|
||||
args.b_k_size = b.size(0);
|
||||
args.b_k_stride = b.stride(0);
|
||||
args.b_n_size = b.size(1);
|
||||
args.b_n_stride = b.stride(1);
|
||||
args.b_ptr = b.data_ptr();
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(b.scalar_type(), "create_onednn_mm_handler",
|
||||
[&] {
|
||||
args.c_type = get_dnnl_type<scalar_t>();
|
||||
args.ab_type = get_dnnl_type<scalar_t>();
|
||||
});
|
||||
|
||||
return reinterpret_cast<int64_t>(new MatMulPrimitiveHandler(args));
|
||||
}
|
||||
|
||||
void onednn_mm(torch::Tensor& c, // [M, OC], row-major
|
||||
const torch::Tensor& a, // [M, IC], row-major
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const torch::Tensor& handler_tensor) {
|
||||
CPU_KERNEL_GUARD_IN(onednn_mm)
|
||||
TORCH_CHECK(a.dim() == 2);
|
||||
TORCH_CHECK(a.stride(-1) == 1);
|
||||
TORCH_CHECK(c.stride(-1) == 1);
|
||||
MatMulPrimitiveHandler* ptr =
|
||||
reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
|
||||
|
||||
// ACL matmuls expect contiguous source tensors
|
||||
#ifdef VLLM_USE_ACL
|
||||
torch::Tensor a_contig = a.contiguous();
|
||||
#endif
|
||||
|
||||
MatMulPrimitiveHandler::ExecArgs exec_args;
|
||||
|
||||
#ifdef VLLM_USE_ACL
|
||||
exec_args.a_m_size = a_contig.size(0);
|
||||
exec_args.a_m_stride = a_contig.stride(0);
|
||||
#else
|
||||
exec_args.a_m_size = a.size(0);
|
||||
exec_args.a_m_stride = a.stride(0);
|
||||
#endif
|
||||
VLLM_DISPATCH_FLOATING_TYPES(a.scalar_type(), "onednn_mm", [&] {
|
||||
if (bias.has_value()) {
|
||||
exec_args.use_bias = true;
|
||||
exec_args.bias_type = get_dnnl_type<scalar_t>();
|
||||
#ifdef VLLM_USE_ACL
|
||||
// ACL matmuls in oneDNN do not support a bias.
|
||||
// We handle a matmul with bias by doing: c = bias; c += matmul(a, b)
|
||||
c.copy_(bias.value());
|
||||
#else
|
||||
exec_args.bias_ptr = bias->data_ptr<scalar_t>();
|
||||
#endif
|
||||
} else {
|
||||
exec_args.use_bias = false;
|
||||
exec_args.bias_type = get_dnnl_type<void>();
|
||||
exec_args.bias_ptr = nullptr;
|
||||
}
|
||||
#ifdef VLLM_USE_ACL
|
||||
exec_args.a_ptr = a_contig.data_ptr<scalar_t>();
|
||||
#else
|
||||
exec_args.a_ptr = a.data_ptr<scalar_t>();
|
||||
|
||||
#endif
|
||||
exec_args.c_ptr = c.data_ptr<scalar_t>();
|
||||
|
||||
ptr->execute(exec_args);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
#pragma once
|
||||
|
||||
#include <bit>
|
||||
#include <cstdint>
|
||||
|
||||
inline float bf16_to_float(uint16_t bf16) {
|
||||
uint32_t bits = static_cast<uint32_t>(bf16) << 16;
|
||||
return std::bit_cast<float>(bits);
|
||||
}
|
||||
|
||||
inline uint16_t float_to_bf16(float fp32) {
|
||||
uint32_t bits = std::bit_cast<uint32_t>(fp32);
|
||||
return static_cast<uint16_t>(bits >> 16);
|
||||
}
|
||||
|
||||
/************************************************
|
||||
* Copyright (c) 2015 Princeton Vision Group
|
||||
* Licensed under the MIT license.
|
||||
* Codes below copied from
|
||||
* https://github.com/PrincetonVision/marvin/tree/master/tools/tensorIO_matlab
|
||||
*************************************************/
|
||||
inline uint16_t float_to_fp16(float fp32) {
|
||||
uint16_t fp16;
|
||||
|
||||
unsigned u, remainder, shift, lsb, lsb_s1, lsb_m1;
|
||||
unsigned sign, exponent, mantissa;
|
||||
|
||||
uint32_t x = std::bit_cast<uint32_t>(fp32);
|
||||
u = (x & 0x7fffffff);
|
||||
|
||||
// Get rid of +NaN/-NaN case first.
|
||||
if (u > 0x7f800000) {
|
||||
fp16 = 0x7fffU;
|
||||
return fp16;
|
||||
}
|
||||
|
||||
sign = ((x >> 16) & 0x8000);
|
||||
|
||||
// Get rid of +Inf/-Inf, +0/-0.
|
||||
if (u > 0x477fefff) {
|
||||
fp16 = sign | 0x7c00U;
|
||||
return fp16;
|
||||
}
|
||||
if (u < 0x33000001) {
|
||||
fp16 = (sign | 0x0000);
|
||||
return fp16;
|
||||
}
|
||||
|
||||
exponent = ((u >> 23) & 0xff);
|
||||
mantissa = (u & 0x7fffff);
|
||||
|
||||
if (exponent > 0x70) {
|
||||
shift = 13;
|
||||
exponent -= 0x70;
|
||||
} else {
|
||||
shift = 0x7e - exponent;
|
||||
exponent = 0;
|
||||
mantissa |= 0x800000;
|
||||
}
|
||||
lsb = (1 << shift);
|
||||
lsb_s1 = (lsb >> 1);
|
||||
lsb_m1 = (lsb - 1);
|
||||
|
||||
// Round to nearest even.
|
||||
remainder = (mantissa & lsb_m1);
|
||||
mantissa >>= shift;
|
||||
if (remainder > lsb_s1 || (remainder == lsb_s1 && (mantissa & 0x1))) {
|
||||
++mantissa;
|
||||
if (!(mantissa & 0x3ff)) {
|
||||
++exponent;
|
||||
mantissa = 0;
|
||||
}
|
||||
}
|
||||
|
||||
fp16 = (sign | (exponent << 10) | mantissa);
|
||||
|
||||
return fp16;
|
||||
}
|
||||
|
||||
inline float fp16_to_float(uint16_t fp16) {
|
||||
unsigned sign = ((fp16 >> 15) & 1);
|
||||
unsigned exponent = ((fp16 >> 10) & 0x1f);
|
||||
unsigned mantissa = ((fp16 & 0x3ff) << 13);
|
||||
uint32_t temp;
|
||||
if (exponent == 0x1f) { /* NaN or Inf */
|
||||
mantissa = (mantissa ? (sign = 0, 0x7fffff) : 0);
|
||||
exponent = 0xff;
|
||||
} else if (!exponent) { /* Denorm or Zero */
|
||||
if (mantissa) {
|
||||
unsigned int msb;
|
||||
exponent = 0x71;
|
||||
do {
|
||||
msb = (mantissa & 0x400000);
|
||||
mantissa <<= 1; /* normalize */
|
||||
--exponent;
|
||||
} while (!msb);
|
||||
mantissa &= 0x7fffff; /* 1.mantissa is implicit */
|
||||
}
|
||||
} else {
|
||||
exponent += 0x70;
|
||||
}
|
||||
temp = ((sign << 31) | (exponent << 23) | mantissa);
|
||||
return std::bit_cast<float>(temp);
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Generate CPU attention dispatch switch cases and kernel instantiations.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
# Head dimensions divisible by 32 (support all ISAs)
|
||||
HEAD_DIMS_32 = [32, 64, 96, 128, 160, 192, 224, 256, 512]
|
||||
|
||||
# Head dimensions divisible by 16 but not 32 (VEC16 only)
|
||||
HEAD_DIMS_16 = [48, 80, 112]
|
||||
|
||||
# ISA types
|
||||
ISA_TYPES = {
|
||||
"AMX": 0,
|
||||
"VEC": 1,
|
||||
"VEC16": 2,
|
||||
"NEON": 3,
|
||||
"VXE": 4,
|
||||
"RVV": 5,
|
||||
"VSX": 6,
|
||||
}
|
||||
|
||||
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
||||
KV_CACHE_IDX = {
|
||||
"auto": 0,
|
||||
"fp8_e4m3": 1,
|
||||
"fp8_e5m2": 2,
|
||||
}
|
||||
|
||||
# C++ type for each kv_cache index
|
||||
KV_CACHE_CPP_TYPES = {
|
||||
"auto": "scalar_t",
|
||||
"fp8_e4m3": "c10::Float8_e4m3fn",
|
||||
"fp8_e5m2": "c10::Float8_e5m2",
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "RVV", "VSX"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
|
||||
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
|
||||
ISA_FOR_FP8 = ["AMX", "VEC"]
|
||||
|
||||
|
||||
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
|
||||
"""Encode head_dim, ISA type, and KV cache type into a single int64_t."""
|
||||
isa_val = ISA_TYPES[isa_type]
|
||||
kv_val = KV_CACHE_IDX[kv_cache]
|
||||
# Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
|
||||
# This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
|
||||
return (head_dim << 16) | (kv_val << 8) | isa_val
|
||||
|
||||
|
||||
def _make_case(
|
||||
head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
|
||||
) -> str:
|
||||
"""Generate a single switch case line."""
|
||||
encoded = encode_params(head_dim, isa, kv_cache)
|
||||
actual_isa = isa_override if isa_override else isa
|
||||
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
|
||||
attn_impl = (
|
||||
f"cpu_attention::AttentionImpl<"
|
||||
f"cpu_attention::ISA::{actual_isa}, \\\n"
|
||||
f" "
|
||||
f"scalar_t, head_dim, {cpp_type}>"
|
||||
)
|
||||
comment = (
|
||||
f"head_dim={head_dim}, isa={isa}"
|
||||
if kv_cache == "auto"
|
||||
else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
|
||||
)
|
||||
return (
|
||||
f""" case {encoded}LL: {{ """
|
||||
f"""/* {comment} */ \\"""
|
||||
f"""
|
||||
constexpr size_t head_dim = {head_dim}; \\"""
|
||||
f"""
|
||||
using attn_impl = {attn_impl}; \\"""
|
||||
f"""
|
||||
return __VA_ARGS__(); \\"""
|
||||
f"""
|
||||
}} \\"""
|
||||
)
|
||||
|
||||
|
||||
def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
|
||||
"""Generate switch cases for a specific ISA group."""
|
||||
cases = []
|
||||
|
||||
# Non-FP8 cases for head_dims divisible by 32
|
||||
for head_dim in HEAD_DIMS_32:
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_32:
|
||||
continue
|
||||
cases.append(_make_case(head_dim, isa, "auto"))
|
||||
|
||||
# Non-FP8 cases for head_dims divisible by 16 only
|
||||
for head_dim in HEAD_DIMS_16:
|
||||
for isa in isa_list:
|
||||
cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
|
||||
|
||||
# FP8 cases: only AMX and VEC, only head_dims divisible by 32
|
||||
if include_fp8:
|
||||
for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
|
||||
for head_dim in HEAD_DIMS_32:
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_FP8:
|
||||
continue
|
||||
cases.append(_make_case(head_dim, isa, fp8_type))
|
||||
|
||||
return "\n".join(cases)
|
||||
|
||||
|
||||
def generate_helper_function() -> str:
|
||||
"""Generate helper function to encode parameters."""
|
||||
return """
|
||||
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
|
||||
int64_t kv_cache_idx = 0) {
|
||||
return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
def generate_header_file() -> str:
|
||||
"""Generate the complete header file content."""
|
||||
header = """// auto generated by generate_cpu_attn_dispatch.py
|
||||
// clang-format off
|
||||
|
||||
#ifndef CPU_ATTN_DISPATCH_GENERATED_H
|
||||
#define CPU_ATTN_DISPATCH_GENERATED_H
|
||||
|
||||
#include "cpu_attn_vec.hpp"
|
||||
#include "cpu_attn_vec16.hpp"
|
||||
|
||||
#ifdef CPU_CAPABILITY_AMXBF16
|
||||
#include "cpu_attn_amx.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __aarch64__
|
||||
#include "cpu_attn_neon.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __s390x__
|
||||
#include "cpu_attn_vxe.hpp"
|
||||
#endif
|
||||
|
||||
// cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256 via RVVI() macros.
|
||||
// Other VLENs and scalar RISC-V builds skip it entirely.
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
||||
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
#include "cpu_attn_rvv.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __powerpc__
|
||||
#include "cpu_attn_vsx.hpp"
|
||||
#endif
|
||||
|
||||
"""
|
||||
|
||||
header += generate_helper_function()
|
||||
|
||||
# Generate dispatch macro with conditional compilation for different ISA sets
|
||||
header += """
|
||||
// Dispatch macro using encoded parameters.
|
||||
// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
|
||||
// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
|
||||
// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
|
||||
// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
|
||||
"""
|
||||
|
||||
def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
|
||||
"""Return one CPU_ATTN_DISPATCH macro block for a given guard."""
|
||||
enc = (
|
||||
" int64_t encoded_params = encode_cpu_attn_params("
|
||||
"HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
|
||||
)
|
||||
cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
|
||||
tail = (
|
||||
"\n"
|
||||
" default: { \\\n"
|
||||
" TORCH_CHECK(false, "
|
||||
'"Unsupported CPU attention configuration: head_dim=" + \\\n'
|
||||
' std::to_string(HEAD_DIM) + " isa=" + \\\n'
|
||||
" std::to_string(static_cast<int>(ISA_TYPE))"
|
||||
" + \\\n"
|
||||
' " kv_cache_idx=" + '
|
||||
"std::to_string(KV_CACHE_IDX)); \\\n"
|
||||
" } \\\n"
|
||||
" } \\\n"
|
||||
" }()\n\n"
|
||||
)
|
||||
return (
|
||||
f"{guard}\n"
|
||||
"#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
|
||||
" [&] { \\\n"
|
||||
f"{enc}\n"
|
||||
" switch (encoded_params) { \\\n"
|
||||
f"{cases}"
|
||||
f"{tail}"
|
||||
)
|
||||
|
||||
header += _macro_block(
|
||||
"#if defined(CPU_CAPABILITY_AMXBF16)",
|
||||
["AMX", "VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__aarch64__)",
|
||||
["NEON", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__s390x__)",
|
||||
["VXE", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
# RISC-V with RVV. cpu_attn_rvv.hpp supports VLEN=128 and VLEN=256
|
||||
# via RVVI() macros. Builds with a supported VLEN get
|
||||
# RVV+VEC+VEC16; other RISC-V builds fall back to VEC/VEC16 only.
|
||||
header += _macro_block(
|
||||
"#elif defined(__riscv) && defined(__riscv_v_min_vlen) "
|
||||
"&& (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)",
|
||||
["RVV", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__riscv)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__powerpc__)",
|
||||
["VSX", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__AVX512F__)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__AVX2__)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#else",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += (
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ /"
|
||||
" __riscv / __powerpc__ */\n\n"
|
||||
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
|
||||
)
|
||||
|
||||
return header
|
||||
|
||||
|
||||
def main():
|
||||
output_path = os.path.join(
|
||||
os.path.dirname(__file__), "cpu_attn_dispatch_generated.h"
|
||||
)
|
||||
|
||||
with open(output_path, "w") as f:
|
||||
f.write(generate_header_file())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,136 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
namespace {
|
||||
template <typename scalar_t>
|
||||
void rms_norm_impl(scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
const scalar_t* __restrict__ weight, const bool has_weight,
|
||||
const float epsilon, const int num_tokens,
|
||||
const int hidden_size) {
|
||||
using scalar_vec_t = vec_op::vec_t<scalar_t>;
|
||||
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
|
||||
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int i = 0; i < num_tokens; ++i) {
|
||||
vec_op::FP32Vec8 variance(0.0);
|
||||
auto input_p = input + i * hidden_size;
|
||||
auto output_p = out + i * hidden_size;
|
||||
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
|
||||
scalar_vec_t x(input_p + j);
|
||||
vec_op::FP32Vec8 fp32_x(x);
|
||||
variance = variance + fp32_x * fp32_x;
|
||||
}
|
||||
|
||||
float s_variance =
|
||||
1.0f / sqrtf(variance.reduce_sum() / (float)hidden_size + epsilon);
|
||||
vec_op::FP32Vec8 fp32_s_variance(s_variance);
|
||||
|
||||
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
|
||||
scalar_vec_t x(input_p + j);
|
||||
vec_op::FP32Vec8 fp32_x(x);
|
||||
vec_op::FP32Vec8 fp32_out;
|
||||
if (has_weight) {
|
||||
scalar_vec_t w(weight + j);
|
||||
vec_op::FP32Vec8 fp32_w(w);
|
||||
fp32_out = fp32_x * fp32_s_variance * fp32_w;
|
||||
} else {
|
||||
fp32_out = fp32_x * fp32_s_variance;
|
||||
}
|
||||
|
||||
scalar_vec_t out(fp32_out);
|
||||
out.save(output_p + j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void fused_add_rms_norm_impl(scalar_t* __restrict__ input,
|
||||
scalar_t* __restrict__ residual,
|
||||
const scalar_t* __restrict__ weight,
|
||||
const bool has_weight, const float epsilon,
|
||||
const int num_tokens, const int hidden_size) {
|
||||
using scalar_vec_t = vec_op::vec_t<scalar_t>;
|
||||
constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
|
||||
TORCH_CHECK(hidden_size % VEC_ELEM_NUM == 0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int i = 0; i < num_tokens; ++i) {
|
||||
vec_op::FP32Vec8 variance(0.0);
|
||||
auto input_p = input + i * hidden_size;
|
||||
auto residual_p = residual + i * hidden_size;
|
||||
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
|
||||
scalar_vec_t x(input_p + j);
|
||||
scalar_vec_t res(residual_p + j);
|
||||
vec_op::FP32Vec8 fp32_x(x);
|
||||
vec_op::FP32Vec8 fp32_res(res);
|
||||
|
||||
fp32_x = fp32_x + fp32_res;
|
||||
variance = variance + fp32_x * fp32_x;
|
||||
scalar_vec_t out(fp32_x);
|
||||
out.save(residual_p + j);
|
||||
}
|
||||
|
||||
float s_variance =
|
||||
1.0f / sqrtf(variance.reduce_sum() / (float)hidden_size + epsilon);
|
||||
vec_op::FP32Vec8 fp32_s_variance(s_variance);
|
||||
|
||||
for (int j = 0; j < hidden_size; j += VEC_ELEM_NUM) {
|
||||
vec_op::FP32Vec8 fp32_out;
|
||||
if (has_weight) {
|
||||
scalar_vec_t w(weight + j);
|
||||
scalar_vec_t res(residual_p + j);
|
||||
vec_op::FP32Vec8 fp32_w(w);
|
||||
vec_op::FP32Vec8 fp32_res(res);
|
||||
fp32_out = fp32_res * fp32_s_variance * fp32_w;
|
||||
} else {
|
||||
scalar_vec_t res(residual_p + j);
|
||||
vec_op::FP32Vec8 fp32_res(res);
|
||||
fp32_out = fp32_res * fp32_s_variance;
|
||||
}
|
||||
|
||||
scalar_vec_t out(fp32_out);
|
||||
out.save(input_p + j);
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void rms_norm(torch::Tensor& out, torch::Tensor& input,
|
||||
std::optional<torch::Tensor> weight, double epsilon) {
|
||||
int hidden_size = input.size(-1);
|
||||
int num_tokens = input.numel() / hidden_size;
|
||||
const bool has_weight = weight.has_value();
|
||||
if (has_weight) {
|
||||
TORCH_CHECK(weight->is_contiguous());
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(rms_norm_impl)
|
||||
rms_norm_impl(out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
|
||||
has_weight ? weight->data_ptr<scalar_t>() : nullptr,
|
||||
has_weight, epsilon, num_tokens, hidden_size);
|
||||
CPU_KERNEL_GUARD_OUT(rms_norm_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
|
||||
std::optional<torch::Tensor> weight, double epsilon) {
|
||||
int hidden_size = input.size(-1);
|
||||
int num_tokens = input.numel() / hidden_size;
|
||||
const bool has_weight = weight.has_value();
|
||||
if (has_weight) {
|
||||
TORCH_CHECK(weight->scalar_type() == input.scalar_type());
|
||||
TORCH_CHECK(weight->is_contiguous());
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "fused_add_rms_norm_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(fused_add_rms_norm_impl)
|
||||
fused_add_rms_norm_impl(
|
||||
input.data_ptr<scalar_t>(), residual.data_ptr<scalar_t>(),
|
||||
has_weight ? weight->data_ptr<scalar_t>() : nullptr, has_weight,
|
||||
epsilon, num_tokens, hidden_size);
|
||||
CPU_KERNEL_GUARD_OUT(fused_add_rms_norm_impl)
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
#ifndef CPU_MICRO_GEMM_AMX_HPP
|
||||
#define CPU_MICRO_GEMM_AMX_HPP
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
namespace {
|
||||
// AMX specific
|
||||
constexpr static int64_t AMX_TILE_ROW_BYTES = 64;
|
||||
constexpr static int64_t AMX_TILE_ROW_NUM = 16;
|
||||
constexpr static int64_t AMX_TILE_BYTES = AMX_TILE_ROW_BYTES * AMX_TILE_ROW_NUM;
|
||||
|
||||
typedef struct __tile_config {
|
||||
uint8_t palette_id = 1;
|
||||
uint8_t start_row = 0;
|
||||
uint8_t reserved_0[14] = {0};
|
||||
uint16_t colsb[16] = {0};
|
||||
uint8_t rows[16] = {0};
|
||||
} __tilecfg;
|
||||
|
||||
// 2-2-4 pattern, for 16 < m <= 32
|
||||
// TILE 0, 1: load A matrix, row num should be 16, m - 16
|
||||
// TILE 2, 3: load B matrix, row num should be 16
|
||||
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
|
||||
// - 16
|
||||
template <typename scalar_t>
|
||||
class TileGemm224 {
|
||||
public:
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TORCH_CHECK(false, "Unsupported data type for TileGemm224");
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
TORCH_CHECK(false, "Unsupported data type for TileGemm224");
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class TileGemm224<c10::BFloat16> {
|
||||
public:
|
||||
using scalar_t = c10::BFloat16;
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
|
||||
c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
|
||||
c10::BFloat16* __restrict__ a_tile_1 = a_ptr + lda * AMX_TILE_ROW_NUM;
|
||||
const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
|
||||
|
||||
// B is always packed as 16 output channels block
|
||||
c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
|
||||
c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
|
||||
const int32_t b_tile_stride = AMX_TILE_ROW_BYTES;
|
||||
|
||||
float* __restrict__ c_tile_4 = c_ptr;
|
||||
float* __restrict__ c_tile_5 =
|
||||
c_tile_4 + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
float* __restrict__ c_tile_6 = c_ptr + AMX_TILE_ROW_NUM * ldc;
|
||||
float* __restrict__ c_tile_7 =
|
||||
c_tile_6 + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
const int32_t c_tile_stride = ldc * sizeof(float);
|
||||
|
||||
if (accum_c) {
|
||||
_tile_loadd(4, c_tile_4, c_tile_stride);
|
||||
_tile_loadd(5, c_tile_5, c_tile_stride);
|
||||
_tile_loadd(6, c_tile_6, c_tile_stride);
|
||||
_tile_loadd(7, c_tile_7, c_tile_stride);
|
||||
} else {
|
||||
_tile_zero(4);
|
||||
_tile_zero(5);
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
for (int32_t k = 0; k < k_times; ++k) {
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_tile_stride);
|
||||
_tile_dpbf16ps(4, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_tile_stride);
|
||||
_tile_dpbf16ps(5, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_dpbf16ps(6, 1, 2);
|
||||
_tile_dpbf16ps(7, 1, 3);
|
||||
|
||||
// update ptrs
|
||||
a_tile_0 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_2 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_3 += AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
}
|
||||
|
||||
_tile_stored(4, c_tile_4, c_tile_stride);
|
||||
_tile_stored(5, c_tile_5, c_tile_stride);
|
||||
_tile_stored(6, c_tile_6, c_tile_stride);
|
||||
_tile_stored(7, c_tile_7, c_tile_stride);
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
const int32_t m_0 = AMX_TILE_ROW_NUM;
|
||||
const int32_t m_1 = m - AMX_TILE_ROW_NUM;
|
||||
config.rows[0] = m_0;
|
||||
config.rows[1] = m_1;
|
||||
config.rows[2] = AMX_TILE_ROW_NUM;
|
||||
config.rows[3] = AMX_TILE_ROW_NUM;
|
||||
config.rows[4] = m_0;
|
||||
config.rows[5] = m_0;
|
||||
config.rows[6] = m_1;
|
||||
config.rows[7] = m_1;
|
||||
_tile_loadconfig(&config);
|
||||
}
|
||||
};
|
||||
|
||||
// 1-2-2 pattern, for 0 < m <= 16
|
||||
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
|
||||
// m, m
|
||||
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
|
||||
// num should be 16
|
||||
// TILE 6, 7, (6, 7): store results C matrix, row num should be
|
||||
// m
|
||||
template <typename scalar_t>
|
||||
class TileGemm122 {
|
||||
public:
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TORCH_CHECK(false, "Unsupported data type for TileGemm122");
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
TORCH_CHECK(false, "Unsupported data type for TileGemm122");
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class TileGemm122<c10::BFloat16> {
|
||||
public:
|
||||
using scalar_t = c10::BFloat16;
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
c10::BFloat16* __restrict__ a_tile_0 = a_ptr;
|
||||
c10::BFloat16* __restrict__ a_tile_1 =
|
||||
a_ptr + AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
const int64_t a_tile_stride = lda * sizeof(c10::BFloat16);
|
||||
|
||||
c10::BFloat16* __restrict__ b_tile_2 = b_ptr;
|
||||
c10::BFloat16* __restrict__ b_tile_3 = b_ptr + b_n_group_stride;
|
||||
c10::BFloat16* __restrict__ b_tile_4 =
|
||||
b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
c10::BFloat16* __restrict__ b_tile_5 =
|
||||
b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
int64_t b_stride = AMX_TILE_ROW_BYTES;
|
||||
|
||||
float* __restrict__ c_tile_6 = c_ptr;
|
||||
float* __restrict__ c_tile_7 = c_ptr + AMX_TILE_ROW_BYTES / sizeof(float);
|
||||
int64_t c_stride = ldc * sizeof(float);
|
||||
|
||||
const int32_t k_times = k / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
|
||||
const int32_t k_group_times = k_times / 2;
|
||||
const bool has_tail = (k_times % 2 == 1);
|
||||
|
||||
if (accum_c) {
|
||||
_tile_loadd(6, c_tile_6, c_stride);
|
||||
_tile_loadd(7, c_tile_7, c_stride);
|
||||
} else {
|
||||
_tile_zero(6);
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
for (int32_t k = 0; k < k_group_times; ++k) {
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_stream_loadd(4, b_tile_4, b_stride);
|
||||
_tile_dpbf16ps(6, 1, 4);
|
||||
_tile_stream_loadd(5, b_tile_5, b_stride);
|
||||
_tile_dpbf16ps(7, 1, 5);
|
||||
|
||||
// update ptrs
|
||||
a_tile_0 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
a_tile_1 += 2 * AMX_TILE_ROW_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_2 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_3 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_4 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
b_tile_5 += 2 * AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
}
|
||||
|
||||
if (has_tail) {
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
}
|
||||
|
||||
_tile_stored(6, c_tile_6, c_stride);
|
||||
_tile_stored(7, c_tile_7, c_stride);
|
||||
}
|
||||
|
||||
FORCE_INLINE static void init_tile_config(int32_t m, __tilecfg& config) {
|
||||
config.rows[0] = m;
|
||||
config.rows[1] = m;
|
||||
config.rows[2] = AMX_TILE_ROW_NUM;
|
||||
config.rows[3] = AMX_TILE_ROW_NUM;
|
||||
config.rows[4] = AMX_TILE_ROW_NUM;
|
||||
config.rows[5] = AMX_TILE_ROW_NUM;
|
||||
config.rows[6] = m;
|
||||
config.rows[7] = m;
|
||||
_tile_loadconfig(&config);
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
// Gemm kernel uses AMX, requires B matrix to be packed
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::AMX, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 32;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = 16;
|
||||
static constexpr bool PackA = false;
|
||||
|
||||
public:
|
||||
MicroGemm() : curr_m_(-1) {
|
||||
vec_op::unroll_loop<int, 8>([&](int i) { amx_tile_config_.colsb[i] = 64; });
|
||||
}
|
||||
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
if (m > AMX_TILE_ROW_NUM) {
|
||||
if (m != curr_m_) {
|
||||
curr_m_ = m;
|
||||
TileGemm224<scalar_t>::init_tile_config(m, amx_tile_config_);
|
||||
}
|
||||
TileGemm224<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
|
||||
} else {
|
||||
if (m != curr_m_) {
|
||||
curr_m_ = m;
|
||||
TileGemm122<scalar_t>::init_tile_config(m, amx_tile_config_);
|
||||
}
|
||||
TileGemm122<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
|
||||
}
|
||||
}
|
||||
|
||||
static void pack_weight(const scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
constexpr int32_t elem_num_per_group = 4 / sizeof(scalar_t);
|
||||
TORCH_CHECK_EQ(output_size % 16, 0);
|
||||
TORCH_CHECK_EQ(input_size % (16 * elem_num_per_group), 0);
|
||||
|
||||
const int32_t output_group_num = output_size / 16;
|
||||
const int32_t input_32b_num = input_size / elem_num_per_group;
|
||||
for (int32_t output_group_idx = 0; output_group_idx < output_group_num;
|
||||
++output_group_idx) {
|
||||
const int32_t* __restrict__ weight_32b =
|
||||
reinterpret_cast<const int32_t*>(weight);
|
||||
int32_t* __restrict__ packed_weight_32b =
|
||||
reinterpret_cast<int32_t*>(packed_weight);
|
||||
for (int32_t output_idx = 0; output_idx < 16; ++output_idx) {
|
||||
for (int32_t weight_offset = 0, packed_offset = 0;
|
||||
weight_offset < input_32b_num;
|
||||
++weight_offset, packed_offset += 16) {
|
||||
packed_weight_32b[packed_offset] = weight_32b[weight_offset];
|
||||
}
|
||||
|
||||
// update
|
||||
weight_32b += input_32b_num;
|
||||
packed_weight_32b += 1;
|
||||
}
|
||||
|
||||
// update
|
||||
weight += 16 * input_size;
|
||||
packed_weight += 16 * input_size;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
alignas(64) __tilecfg amx_tile_config_;
|
||||
int32_t curr_m_;
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,132 @@
|
||||
#ifndef CPU_MICRO_GEMM_IMPL_HPP
|
||||
#define CPU_MICRO_GEMM_IMPL_HPP
|
||||
#include "cpu/utils.hpp"
|
||||
#include "cpu/cpu_types.hpp"
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
#define DEFINE_CPU_MICRO_GEMM_PARAMS \
|
||||
scalar_t *__restrict__ a_ptr, scalar_t *__restrict__ b_ptr, \
|
||||
float *__restrict__ c_ptr, const int32_t m, const int32_t k, \
|
||||
const int64_t lda, const int64_t b_n_group_stride, const int64_t ldc, \
|
||||
const bool accum_c
|
||||
|
||||
#define CPU_MICRO_GEMM_PARAMS \
|
||||
a_ptr, b_ptr, c_ptr, m, k, lda, b_n_group_stride, ldc, accum_c
|
||||
|
||||
// Note: weights for MicroGemm should be packed as (output_size / 16) contiguous
|
||||
// blocks, means the logical shape of blocks is [16, input_size]. And the actual
|
||||
// layout of blocks can be ISA-specific.
|
||||
template <cpu_utils::ISA isa, typename scalar_t>
|
||||
class MicroGemm {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 16;
|
||||
static constexpr int32_t NSize = 16;
|
||||
static constexpr int32_t WeightOCGroupSize = 16;
|
||||
// callers must pack A matrix before GEMM
|
||||
static constexpr bool PackA = false;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TORCH_CHECK(false, "Unimplemented MicroGemm.");
|
||||
}
|
||||
};
|
||||
|
||||
template <int32_t n_size, typename scalar_t>
|
||||
FORCE_INLINE void default_epilogue(float* __restrict__ c_ptr,
|
||||
scalar_t* __restrict__ d_ptr,
|
||||
const int32_t m, const int64_t ldc,
|
||||
const int64_t ldd) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
static_assert(n_size % 16 == 0);
|
||||
|
||||
float* __restrict__ curr_c = c_ptr;
|
||||
scalar_t* __restrict__ curr_d = d_ptr;
|
||||
for (int32_t i = 0; i < m; ++i) {
|
||||
float* __restrict__ curr_c_iter = curr_c;
|
||||
scalar_t* __restrict__ curr_d_iter = curr_d;
|
||||
vec_op::unroll_loop<int32_t, n_size / 16>([&](int32_t n_g_idx) {
|
||||
vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
|
||||
scalar_vec_t c_vec(c_vec_fp32);
|
||||
c_vec.save(curr_d_iter);
|
||||
curr_c_iter += 16;
|
||||
curr_d_iter += 16;
|
||||
});
|
||||
curr_c += ldc;
|
||||
curr_d += ldd;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t n_size, typename scalar_t>
|
||||
FORCE_INLINE void bias_epilogue(float* __restrict__ c_ptr,
|
||||
scalar_t* __restrict__ d_ptr,
|
||||
scalar_t* __restrict__ bias_ptr,
|
||||
const int32_t m, const int64_t ldc,
|
||||
const int64_t ldd) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
static_assert(n_size % 16 == 0);
|
||||
constexpr int32_t n_group_num = n_size / 16;
|
||||
static_assert(n_group_num <= 16);
|
||||
|
||||
vec_op::FP32Vec16 bias_vecs[n_group_num];
|
||||
scalar_t* __restrict__ curr_bias = bias_ptr;
|
||||
vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
|
||||
scalar_vec_t vec(curr_bias);
|
||||
bias_vecs[i] = vec_op::FP32Vec16(vec);
|
||||
curr_bias += 16;
|
||||
});
|
||||
|
||||
float* __restrict__ curr_c = c_ptr;
|
||||
scalar_t* __restrict__ curr_d = d_ptr;
|
||||
for (int32_t i = 0; i < m; ++i) {
|
||||
float* __restrict__ curr_c_iter = curr_c;
|
||||
scalar_t* __restrict__ curr_d_iter = curr_d;
|
||||
vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
|
||||
vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
|
||||
c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
|
||||
scalar_vec_t c_vec(c_vec_fp32);
|
||||
c_vec.save(curr_d_iter);
|
||||
curr_c_iter += 16;
|
||||
curr_d_iter += 16;
|
||||
});
|
||||
curr_c += ldc;
|
||||
curr_d += ldd;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t n_size, typename scalar_t>
|
||||
FORCE_INLINE void add_bias_epilogue(float* c_ptr, float* d_ptr,
|
||||
scalar_t* __restrict__ bias_ptr,
|
||||
const int32_t m, const int64_t ldc,
|
||||
const int64_t ldd) {
|
||||
using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
static_assert(n_size % 16 == 0);
|
||||
constexpr int32_t n_group_num = n_size / 16;
|
||||
static_assert(n_group_num <= 16);
|
||||
|
||||
vec_op::FP32Vec16 bias_vecs[n_group_num];
|
||||
scalar_t* __restrict__ curr_bias = bias_ptr;
|
||||
vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t i) {
|
||||
scalar_vec_t vec(curr_bias);
|
||||
bias_vecs[i] = vec_op::FP32Vec16(vec);
|
||||
curr_bias += 16;
|
||||
});
|
||||
|
||||
float* curr_c = c_ptr;
|
||||
float* curr_d = d_ptr;
|
||||
for (int32_t i = 0; i < m; ++i) {
|
||||
float* curr_c_iter = curr_c;
|
||||
float* curr_d_iter = curr_d;
|
||||
vec_op::unroll_loop<int32_t, n_group_num>([&](int32_t n_g_idx) {
|
||||
vec_op::FP32Vec16 c_vec_fp32(curr_c_iter);
|
||||
c_vec_fp32 = c_vec_fp32 + bias_vecs[n_g_idx];
|
||||
c_vec_fp32.save(curr_d_iter);
|
||||
curr_c_iter += 16;
|
||||
curr_d_iter += 16;
|
||||
});
|
||||
curr_c += ldc;
|
||||
curr_d += ldd;
|
||||
}
|
||||
}
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,503 @@
|
||||
#ifndef CPU_MICRO_GEMM_NEON_HPP
|
||||
#define CPU_MICRO_GEMM_NEON_HPP
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
#include <arm_bf16.h>
|
||||
#include <arm_neon.h>
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int32_t K = 4;
|
||||
constexpr int32_t Cols = 2;
|
||||
constexpr int32_t TileSize = K * Cols;
|
||||
constexpr int32_t Mr = 8;
|
||||
constexpr int32_t Nr = 8;
|
||||
constexpr int32_t Nr_gemv = 16;
|
||||
|
||||
// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a0, a1, b0, b1]
|
||||
FORCE_INLINE float32x4_t zip1_f32x4(const float32x4_t a, const float32x4_t b) {
|
||||
return vreinterpretq_f32_f64(
|
||||
vzip1q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
|
||||
}
|
||||
|
||||
// a = [a0, a1, a2, a3], b = [b0, b1, b2, b3] -> [a2, a3, b2, b3]
|
||||
FORCE_INLINE float32x4_t zip2_f32x4(const float32x4_t a, const float32x4_t b) {
|
||||
return vreinterpretq_f32_f64(
|
||||
vzip2q_f64(vreinterpretq_f64_f32(a), vreinterpretq_f64_f32(b)));
|
||||
}
|
||||
|
||||
FORCE_INLINE void init_acc_rowpair(float32x4_t& acc01, float32x4_t& acc23,
|
||||
float32x4_t& acc45, float32x4_t& acc67,
|
||||
const float* __restrict__ c_ptr,
|
||||
const int64_t ldc, const int32_t m_rows,
|
||||
const bool accum_c) {
|
||||
if (!accum_c || m_rows == 0) {
|
||||
acc01 = vdupq_n_f32(0.0f);
|
||||
acc23 = vdupq_n_f32(0.0f);
|
||||
acc45 = vdupq_n_f32(0.0f);
|
||||
acc67 = vdupq_n_f32(0.0f);
|
||||
return;
|
||||
}
|
||||
|
||||
const float32x4_t row0_0123 = vld1q_f32(c_ptr);
|
||||
const float32x4_t row0_4567 = vld1q_f32(c_ptr + 4);
|
||||
const float32x4_t row1_0123 =
|
||||
(m_rows == 2) ? vld1q_f32(c_ptr + ldc) : vdupq_n_f32(0.0f);
|
||||
const float32x4_t row1_4567 =
|
||||
(m_rows == 2) ? vld1q_f32(c_ptr + ldc + 4) : vdupq_n_f32(0.0f);
|
||||
|
||||
acc01 = zip1_f32x4(row0_0123, row1_0123);
|
||||
acc23 = zip2_f32x4(row0_0123, row1_0123);
|
||||
acc45 = zip1_f32x4(row0_4567, row1_4567);
|
||||
acc67 = zip2_f32x4(row0_4567, row1_4567);
|
||||
}
|
||||
|
||||
FORCE_INLINE void store_acc_rowpair(const float32x4_t acc01,
|
||||
const float32x4_t acc23,
|
||||
const float32x4_t acc45,
|
||||
const float32x4_t acc67,
|
||||
float* __restrict__ c_ptr,
|
||||
const int64_t ldc, const int32_t m_rows) {
|
||||
if (m_rows == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
vst1q_f32(c_ptr, zip1_f32x4(acc01, acc23));
|
||||
vst1q_f32(c_ptr + 4, zip1_f32x4(acc45, acc67));
|
||||
|
||||
if (m_rows == 2) {
|
||||
vst1q_f32(c_ptr + ldc, zip2_f32x4(acc01, acc23));
|
||||
vst1q_f32(c_ptr + ldc + 4, zip2_f32x4(acc45, acc67));
|
||||
}
|
||||
}
|
||||
|
||||
FORCE_INLINE void gemm_micro_bfmmla_8x8_packed_a(
|
||||
const bfloat16_t* __restrict__ a_packed,
|
||||
const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
|
||||
const int32_t m, const int32_t k_size, const int64_t ldc,
|
||||
const bool accum_c) {
|
||||
float32x4_t acc0101, acc0123, acc0145, acc0167;
|
||||
float32x4_t acc2301, acc2323, acc2345, acc2367;
|
||||
float32x4_t acc4501, acc4523, acc4545, acc4567;
|
||||
float32x4_t acc6701, acc6723, acc6745, acc6767;
|
||||
|
||||
init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
|
||||
std::min(2, m), accum_c);
|
||||
init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 2)), accum_c);
|
||||
init_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 4)), accum_c);
|
||||
init_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 6)), accum_c);
|
||||
|
||||
const bfloat16_t* __restrict__ a_tile = a_packed;
|
||||
const bfloat16_t* __restrict__ b_tile = b_packed;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
|
||||
const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
|
||||
const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
|
||||
const bfloat16x8_t a_tile45 = vld1q_bf16(a_tile + 2 * TileSize);
|
||||
const bfloat16x8_t a_tile67 = vld1q_bf16(a_tile + 3 * TileSize);
|
||||
|
||||
const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile);
|
||||
const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile + TileSize);
|
||||
const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile + 2 * TileSize);
|
||||
const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile + 3 * TileSize);
|
||||
|
||||
acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
|
||||
acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
|
||||
acc4501 = vbfmmlaq_f32(acc4501, a_tile45, b_tile01);
|
||||
acc6701 = vbfmmlaq_f32(acc6701, a_tile67, b_tile01);
|
||||
|
||||
acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
|
||||
acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
|
||||
acc4523 = vbfmmlaq_f32(acc4523, a_tile45, b_tile23);
|
||||
acc6723 = vbfmmlaq_f32(acc6723, a_tile67, b_tile23);
|
||||
|
||||
acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
|
||||
acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
|
||||
acc4545 = vbfmmlaq_f32(acc4545, a_tile45, b_tile45);
|
||||
acc6745 = vbfmmlaq_f32(acc6745, a_tile67, b_tile45);
|
||||
|
||||
acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
|
||||
acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
|
||||
acc4567 = vbfmmlaq_f32(acc4567, a_tile45, b_tile67);
|
||||
acc6767 = vbfmmlaq_f32(acc6767, a_tile67, b_tile67);
|
||||
|
||||
a_tile += 4 * TileSize;
|
||||
b_tile += Nr * K;
|
||||
}
|
||||
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc,
|
||||
std::min(2, m));
|
||||
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 2)));
|
||||
store_acc_rowpair(acc4501, acc4523, acc4545, acc4567, c_ptr + 4 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 4)));
|
||||
store_acc_rowpair(acc6701, acc6723, acc6745, acc6767, c_ptr + 6 * ldc, ldc,
|
||||
std::min(2, std::max(0, m - 6)));
|
||||
}
|
||||
|
||||
FORCE_INLINE void gemm_micro_bfmmla_4x16_packed_a(
|
||||
const bfloat16_t* __restrict__ a_packed,
|
||||
const bfloat16_t* __restrict__ b_packed, float* __restrict__ c_ptr,
|
||||
const int32_t m, const int32_t k_size, const int64_t b_n_group_stride,
|
||||
const int64_t ldc, const bool accum_c) {
|
||||
const int32_t m_rows_01 = std::min(2, m);
|
||||
const int32_t m_rows_23 = std::min(2, std::max(0, m - 2));
|
||||
|
||||
float32x4_t acc0101, acc0123, acc0145, acc0167;
|
||||
float32x4_t acc2301, acc2323, acc2345, acc2367;
|
||||
float32x4_t acc0189, acc011011, acc011213, acc011415;
|
||||
float32x4_t acc2389, acc231011, acc231213, acc231415;
|
||||
|
||||
init_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01,
|
||||
accum_c);
|
||||
init_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
m_rows_23, accum_c);
|
||||
init_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
|
||||
m_rows_01, accum_c);
|
||||
init_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
|
||||
c_ptr + 2 * ldc + 8, ldc, m_rows_23, accum_c);
|
||||
|
||||
const bfloat16_t* __restrict__ a_tile = a_packed;
|
||||
const bfloat16_t* __restrict__ b_tile0 = b_packed;
|
||||
const bfloat16_t* __restrict__ b_tile1 = b_packed + b_n_group_stride;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int32_t k_idx = 0; k_idx < k_size; k_idx += K) {
|
||||
const bfloat16x8_t a_tile01 = vld1q_bf16(a_tile);
|
||||
const bfloat16x8_t a_tile23 = vld1q_bf16(a_tile + TileSize);
|
||||
const bfloat16x8_t b_tile01 = vld1q_bf16(b_tile0);
|
||||
const bfloat16x8_t b_tile23 = vld1q_bf16(b_tile0 + TileSize);
|
||||
const bfloat16x8_t b_tile45 = vld1q_bf16(b_tile0 + 2 * TileSize);
|
||||
const bfloat16x8_t b_tile67 = vld1q_bf16(b_tile0 + 3 * TileSize);
|
||||
const bfloat16x8_t b_tile89 = vld1q_bf16(b_tile1);
|
||||
const bfloat16x8_t b_tile1011 = vld1q_bf16(b_tile1 + TileSize);
|
||||
const bfloat16x8_t b_tile1213 = vld1q_bf16(b_tile1 + 2 * TileSize);
|
||||
const bfloat16x8_t b_tile1415 = vld1q_bf16(b_tile1 + 3 * TileSize);
|
||||
|
||||
acc0101 = vbfmmlaq_f32(acc0101, a_tile01, b_tile01);
|
||||
acc2301 = vbfmmlaq_f32(acc2301, a_tile23, b_tile01);
|
||||
acc0123 = vbfmmlaq_f32(acc0123, a_tile01, b_tile23);
|
||||
acc2323 = vbfmmlaq_f32(acc2323, a_tile23, b_tile23);
|
||||
|
||||
acc0145 = vbfmmlaq_f32(acc0145, a_tile01, b_tile45);
|
||||
acc2345 = vbfmmlaq_f32(acc2345, a_tile23, b_tile45);
|
||||
acc0167 = vbfmmlaq_f32(acc0167, a_tile01, b_tile67);
|
||||
acc2367 = vbfmmlaq_f32(acc2367, a_tile23, b_tile67);
|
||||
|
||||
acc0189 = vbfmmlaq_f32(acc0189, a_tile01, b_tile89);
|
||||
acc2389 = vbfmmlaq_f32(acc2389, a_tile23, b_tile89);
|
||||
acc011011 = vbfmmlaq_f32(acc011011, a_tile01, b_tile1011);
|
||||
acc231011 = vbfmmlaq_f32(acc231011, a_tile23, b_tile1011);
|
||||
|
||||
acc011213 = vbfmmlaq_f32(acc011213, a_tile01, b_tile1213);
|
||||
acc231213 = vbfmmlaq_f32(acc231213, a_tile23, b_tile1213);
|
||||
acc011415 = vbfmmlaq_f32(acc011415, a_tile01, b_tile1415);
|
||||
acc231415 = vbfmmlaq_f32(acc231415, a_tile23, b_tile1415);
|
||||
|
||||
a_tile += 2 * TileSize;
|
||||
b_tile0 += Nr * K;
|
||||
b_tile1 += Nr * K;
|
||||
}
|
||||
|
||||
store_acc_rowpair(acc0101, acc0123, acc0145, acc0167, c_ptr, ldc, m_rows_01);
|
||||
store_acc_rowpair(acc2301, acc2323, acc2345, acc2367, c_ptr + 2 * ldc, ldc,
|
||||
m_rows_23);
|
||||
store_acc_rowpair(acc0189, acc011011, acc011213, acc011415, c_ptr + 8, ldc,
|
||||
m_rows_01);
|
||||
store_acc_rowpair(acc2389, acc231011, acc231213, acc231415,
|
||||
c_ptr + 2 * ldc + 8, ldc, m_rows_23);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::NEON, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = Nr;
|
||||
static constexpr bool PackA = false;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
|
||||
}
|
||||
|
||||
static void pack_weight(const scalar_t* __restrict__ /*weight*/,
|
||||
scalar_t* __restrict__ /*packed_weight*/,
|
||||
const int32_t /*output_size*/,
|
||||
const int32_t /*input_size*/) {
|
||||
TORCH_CHECK(false, "NEON BFMMLA MicroGemm only supports bfloat16.");
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class MicroGemm<cpu_utils::ISA::NEON, c10::BFloat16> {
|
||||
public:
|
||||
using scalar_t = c10::BFloat16;
|
||||
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = Nr;
|
||||
static constexpr bool PackA = true;
|
||||
|
||||
public:
|
||||
// physical layout [
|
||||
// M / 8; Mr is 8
|
||||
// K / 4; K for bfmmla is 4
|
||||
// 4, ; 4 row-pairs for each 8 rows
|
||||
// 2, ; row-pair is 2 rows
|
||||
// 4 ; 4 elements per row
|
||||
// ]
|
||||
|
||||
static void pack_input_from_rows(const scalar_t* const* __restrict__ rows,
|
||||
scalar_t* __restrict__ a_packed,
|
||||
const int32_t m, const int32_t k) {
|
||||
TORCH_CHECK(m > 0 && m <= MaxMSize);
|
||||
TORCH_CHECK_EQ(k % K, 0);
|
||||
|
||||
auto* __restrict__ out = reinterpret_cast<bfloat16_t*>(a_packed);
|
||||
const bfloat16x8_t zero_q = vdupq_n_bf16(bfloat16_t{});
|
||||
const bfloat16x4_t zero = vget_low_bf16(zero_q);
|
||||
|
||||
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
|
||||
const int32_t actual_m = std::min(Mr, m - row_base);
|
||||
const bfloat16_t* __restrict__ row[Mr];
|
||||
for (int32_t i = 0; i < actual_m; ++i) {
|
||||
row[i] = reinterpret_cast<const bfloat16_t*>(rows[row_base + i]);
|
||||
}
|
||||
|
||||
if (actual_m == 8) {
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 8 <= k; k_idx += 8) {
|
||||
bfloat16_t* __restrict__ block0 = out;
|
||||
bfloat16_t* __restrict__ block1 = out + 4 * TileSize;
|
||||
|
||||
bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
|
||||
bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
|
||||
vst1q_bf16(block0,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = vld1q_bf16(row[2] + k_idx);
|
||||
a1 = vld1q_bf16(row[3] + k_idx);
|
||||
vst1q_bf16(block0 + TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = vld1q_bf16(row[4] + k_idx);
|
||||
a1 = vld1q_bf16(row[5] + k_idx);
|
||||
vst1q_bf16(block0 + 2 * TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + 2 * TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = vld1q_bf16(row[6] + k_idx);
|
||||
a1 = vld1q_bf16(row[7] + k_idx);
|
||||
vst1q_bf16(block0 + 3 * TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + 3 * TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
out += 8 * TileSize;
|
||||
}
|
||||
|
||||
for (; k_idx < k; k_idx += K) {
|
||||
bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
|
||||
bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
|
||||
vst1q_bf16(out, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = vld1_bf16(row[2] + k_idx);
|
||||
a1 = vld1_bf16(row[3] + k_idx);
|
||||
vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = vld1_bf16(row[4] + k_idx);
|
||||
a1 = vld1_bf16(row[5] + k_idx);
|
||||
vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = vld1_bf16(row[6] + k_idx);
|
||||
a1 = vld1_bf16(row[7] + k_idx);
|
||||
vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
out += 4 * TileSize;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (actual_m == 4) {
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 8 <= k; k_idx += 8) {
|
||||
bfloat16_t* __restrict__ block0 = out;
|
||||
bfloat16_t* __restrict__ block1 = out + 2 * TileSize;
|
||||
|
||||
bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
|
||||
bfloat16x8_t a1 = vld1q_bf16(row[1] + k_idx);
|
||||
vst1q_bf16(block0,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = vld1q_bf16(row[2] + k_idx);
|
||||
a1 = vld1q_bf16(row[3] + k_idx);
|
||||
vst1q_bf16(block0 + TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
out += 4 * TileSize;
|
||||
}
|
||||
|
||||
for (; k_idx < k; k_idx += K) {
|
||||
bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
|
||||
bfloat16x4_t a1 = vld1_bf16(row[1] + k_idx);
|
||||
vst1q_bf16(out, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = vld1_bf16(row[2] + k_idx);
|
||||
a1 = vld1_bf16(row[3] + k_idx);
|
||||
vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
out += 2 * TileSize;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const int32_t row_pair_count = (actual_m <= 4) ? 2 : Mr / 2;
|
||||
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 8 <= k; k_idx += 8) {
|
||||
bfloat16_t* __restrict__ block0 = out;
|
||||
bfloat16_t* __restrict__ block1 = out + row_pair_count * TileSize;
|
||||
|
||||
bfloat16x8_t a0 = vld1q_bf16(row[0] + k_idx);
|
||||
bfloat16x8_t a1 = (actual_m > 1) ? vld1q_bf16(row[1] + k_idx) : zero_q;
|
||||
vst1q_bf16(block0, vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = (actual_m > 2) ? vld1q_bf16(row[2] + k_idx) : zero_q;
|
||||
a1 = (actual_m > 3) ? vld1q_bf16(row[3] + k_idx) : zero_q;
|
||||
vst1q_bf16(block0 + TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
if (actual_m > 4) {
|
||||
a0 = vld1q_bf16(row[4] + k_idx);
|
||||
a1 = (actual_m > 5) ? vld1q_bf16(row[5] + k_idx) : zero_q;
|
||||
vst1q_bf16(block0 + 2 * TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + 2 * TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
|
||||
a0 = (actual_m > 6) ? vld1q_bf16(row[6] + k_idx) : zero_q;
|
||||
a1 = (actual_m > 7) ? vld1q_bf16(row[7] + k_idx) : zero_q;
|
||||
vst1q_bf16(block0 + 3 * TileSize,
|
||||
vcombine_bf16(vget_low_bf16(a0), vget_low_bf16(a1)));
|
||||
vst1q_bf16(block1 + 3 * TileSize,
|
||||
vcombine_bf16(vget_high_bf16(a0), vget_high_bf16(a1)));
|
||||
}
|
||||
|
||||
out += 2 * row_pair_count * TileSize;
|
||||
}
|
||||
|
||||
for (; k_idx < k; k_idx += K) {
|
||||
bfloat16x4_t a0 = vld1_bf16(row[0] + k_idx);
|
||||
bfloat16x4_t a1 = (actual_m > 1) ? vld1_bf16(row[1] + k_idx) : zero;
|
||||
vst1q_bf16(out, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = (actual_m > 2) ? vld1_bf16(row[2] + k_idx) : zero;
|
||||
a1 = (actual_m > 3) ? vld1_bf16(row[3] + k_idx) : zero;
|
||||
vst1q_bf16(out + TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
if (actual_m > 4) {
|
||||
a0 = vld1_bf16(row[4] + k_idx);
|
||||
a1 = (actual_m > 5) ? vld1_bf16(row[5] + k_idx) : zero;
|
||||
vst1q_bf16(out + 2 * TileSize, vcombine_bf16(a0, a1));
|
||||
|
||||
a0 = (actual_m > 6) ? vld1_bf16(row[6] + k_idx) : zero;
|
||||
a1 = (actual_m > 7) ? vld1_bf16(row[7] + k_idx) : zero;
|
||||
vst1q_bf16(out + 3 * TileSize, vcombine_bf16(a0, a1));
|
||||
}
|
||||
out += row_pair_count * TileSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
(void)lda; // A is packed, so lda is not needed
|
||||
TORCH_CHECK_EQ(k % K, 0);
|
||||
|
||||
for (int32_t n_idx = 0; n_idx < NSize; n_idx += Nr_gemv) {
|
||||
const bfloat16_t* __restrict__ b_panel =
|
||||
reinterpret_cast<const bfloat16_t*>(b_ptr) + n_idx * k;
|
||||
|
||||
for (int32_t row_base = 0; row_base < m; row_base += Mr) {
|
||||
const int32_t panel_m = std::min(Mr, m - row_base);
|
||||
const bfloat16_t* __restrict__ a_panel =
|
||||
reinterpret_cast<const bfloat16_t*>(a_ptr) + row_base * k;
|
||||
float* __restrict__ c_panel = c_ptr + row_base * ldc + n_idx;
|
||||
|
||||
if (panel_m <= 4) {
|
||||
gemm_micro_bfmmla_4x16_packed_a(a_panel, b_panel, c_panel, panel_m, k,
|
||||
b_n_group_stride, ldc, accum_c);
|
||||
} else {
|
||||
gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel, c_panel, panel_m, k,
|
||||
ldc, accum_c);
|
||||
gemm_micro_bfmmla_8x8_packed_a(a_panel, b_panel + b_n_group_stride,
|
||||
c_panel + Nr, panel_m, k, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// physical layout [
|
||||
// N / 8; Nr is 8
|
||||
// K / 4; K for bfmmla is 4
|
||||
// 4, ; 4 col-pairs for each 8 cols
|
||||
// 2, ; col-pair is 2 cols
|
||||
// 4 ; 4 elements per col
|
||||
// ]
|
||||
static void pack_weight(const c10::BFloat16* __restrict__ weight,
|
||||
c10::BFloat16* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK_EQ(output_size % NSize, 0);
|
||||
TORCH_CHECK_EQ(input_size % K, 0);
|
||||
|
||||
for (int32_t o_idx = 0; o_idx < output_size; o_idx += Nr) {
|
||||
c10::BFloat16* __restrict__ dst = packed_weight + o_idx * input_size;
|
||||
for (int32_t k_idx = 0; k_idx < input_size; k_idx += K) {
|
||||
for (int32_t pair_idx = 0; pair_idx < Nr; pair_idx += Cols) {
|
||||
const c10::BFloat16* __restrict__ row0 =
|
||||
weight + (o_idx + pair_idx) * input_size;
|
||||
const c10::BFloat16* __restrict__ row1 = row0 + input_size;
|
||||
dst[0] = row0[k_idx + 0];
|
||||
dst[1] = row0[k_idx + 1];
|
||||
dst[2] = row0[k_idx + 2];
|
||||
dst[3] = row0[k_idx + 3];
|
||||
dst[4] = row1[k_idx + 0];
|
||||
dst[5] = row1[k_idx + 1];
|
||||
dst[6] = row1[k_idx + 2];
|
||||
dst[7] = row1[k_idx + 3];
|
||||
dst += TileSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,228 @@
|
||||
#ifndef CPU_MICRO_GEMM_RVV_HPP
|
||||
#define CPU_MICRO_GEMM_RVV_HPP
|
||||
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
#if defined(__riscv_v)
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
namespace {
|
||||
|
||||
constexpr int32_t RVV_MGEMM_N8 = 8;
|
||||
constexpr int32_t RVV_MGEMM_B_GROUP_STRIDE = 16;
|
||||
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32(const scalar_t* ptr);
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t load_row8_b_as_f32<float>(const float* ptr) {
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(ptr, RVV_MGEMM_N8);
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::Half>(const c10::Half* ptr) {
|
||||
#if defined(__riscv_zvfh)
|
||||
fixed_fp16x8_t vec = RVVI(__riscv_vle16_v_f16, LMUL_128)(
|
||||
reinterpret_cast<const _Float16*>(ptr), RVV_MGEMM_N8);
|
||||
return RVVI(__riscv_vfwcvt_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
alignas(32) float values[RVV_MGEMM_N8];
|
||||
for (int32_t i = 0; i < RVV_MGEMM_N8; ++i) {
|
||||
values[i] = static_cast<float>(ptr[i]);
|
||||
}
|
||||
return RVVI(__riscv_vle32_v_f32, LMUL_256)(values, RVV_MGEMM_N8);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
FORCE_INLINE fixed_fp32x8_t
|
||||
load_row8_b_as_f32<c10::BFloat16>(const c10::BFloat16* ptr) {
|
||||
#if defined(__riscv_zvfbfmin)
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
fixed_bf16x8_t vec =
|
||||
RVVI4(__riscv_vreinterpret_v_u16, LMUL_128, _bf16, LMUL_128)(raw);
|
||||
return RVVI(__riscv_vfwcvtbf16_f_f_v_f32, LMUL_256)(vec, RVV_MGEMM_N8);
|
||||
#else
|
||||
fixed_u16x8_t raw = RVVI(__riscv_vle16_v_u16, LMUL_128)(
|
||||
reinterpret_cast<const uint16_t*>(ptr), RVV_MGEMM_N8);
|
||||
auto wide = RVVI(__riscv_vzext_vf2_u32, LMUL_256)(raw, RVV_MGEMM_N8);
|
||||
auto shifted = RVVI(__riscv_vsll_vx_u32, LMUL_256)(wide, 16, RVV_MGEMM_N8);
|
||||
return RVVI4(__riscv_vreinterpret_v_u32, LMUL_256, _f32, LMUL_256)(shifted);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Mx8 RVV kernel. B points at one 8-channel half of a 16-channel packed group,
|
||||
// with rows separated by RVV_MGEMM_B_GROUP_STRIDE scalar elements.
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_fma_mx8_ku4(const scalar_t* __restrict__ a_ptr,
|
||||
const scalar_t* __restrict__ b_ptr,
|
||||
float* __restrict__ c_ptr,
|
||||
const int64_t lda,
|
||||
const int64_t ldc, const int32_t k,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
|
||||
#define RVV_ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define RVV_IF_M(i) if constexpr (M > (i))
|
||||
|
||||
#define RVV_DECL_A(i) const scalar_t* __restrict__ a##i = a_ptr + (i) * lda;
|
||||
RVV_ROWS_APPLY(RVV_DECL_A)
|
||||
#undef RVV_DECL_A
|
||||
|
||||
#define RVV_DECL_ACC(i) fixed_fp32x8_t acc##i;
|
||||
RVV_ROWS_APPLY(RVV_DECL_ACC)
|
||||
#undef RVV_DECL_ACC
|
||||
|
||||
#define RVV_INIT_ACC(i) \
|
||||
RVV_IF_M(i) { \
|
||||
if (accum_c) { \
|
||||
acc##i = RVVI(__riscv_vle32_v_f32, LMUL_256)(c_ptr + (i) * ldc, \
|
||||
RVV_MGEMM_N8); \
|
||||
} else { \
|
||||
acc##i = RVVI(__riscv_vfmv_v_f_f32, LMUL_256)(0.0f, RVV_MGEMM_N8); \
|
||||
} \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_INIT_ACC)
|
||||
#undef RVV_INIT_ACC
|
||||
|
||||
int32_t k_idx = 0;
|
||||
for (; k_idx + 3 < k; k_idx += 4) {
|
||||
#define RVV_FMA_ROW(i, K_OFFSET) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx + (K_OFFSET))), b, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
|
||||
#define RVV_STEP_K(K_OFFSET) \
|
||||
{ \
|
||||
fixed_fp32x8_t b = load_row8_b_as_f32<scalar_t>( \
|
||||
b_ptr + (k_idx + (K_OFFSET)) * RVV_MGEMM_B_GROUP_STRIDE); \
|
||||
RVV_FMA_ROW(0, K_OFFSET) \
|
||||
RVV_FMA_ROW(1, K_OFFSET) \
|
||||
RVV_FMA_ROW(2, K_OFFSET) \
|
||||
RVV_FMA_ROW(3, K_OFFSET) \
|
||||
RVV_FMA_ROW(4, K_OFFSET) \
|
||||
RVV_FMA_ROW(5, K_OFFSET) \
|
||||
RVV_FMA_ROW(6, K_OFFSET) \
|
||||
RVV_FMA_ROW(7, K_OFFSET) \
|
||||
}
|
||||
|
||||
RVV_STEP_K(0)
|
||||
RVV_STEP_K(1)
|
||||
RVV_STEP_K(2)
|
||||
RVV_STEP_K(3)
|
||||
#undef RVV_STEP_K
|
||||
#undef RVV_FMA_ROW
|
||||
}
|
||||
|
||||
for (; k_idx < k; ++k_idx) {
|
||||
fixed_fp32x8_t b =
|
||||
load_row8_b_as_f32<scalar_t>(b_ptr + k_idx * RVV_MGEMM_B_GROUP_STRIDE);
|
||||
#define RVV_TAIL_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
acc##i = RVVI(__riscv_vfmacc_vf_f32, LMUL_256)( \
|
||||
acc##i, static_cast<float>(*(a##i + k_idx)), b, RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_TAIL_ROW)
|
||||
#undef RVV_TAIL_ROW
|
||||
}
|
||||
|
||||
#define RVV_STORE_ROW(i) \
|
||||
RVV_IF_M(i) { \
|
||||
RVVI(__riscv_vse32_v_f32, LMUL_256)(c_ptr + (i) * ldc, acc##i, \
|
||||
RVV_MGEMM_N8); \
|
||||
}
|
||||
RVV_ROWS_APPLY(RVV_STORE_ROW)
|
||||
#undef RVV_STORE_ROW
|
||||
|
||||
#undef RVV_ROWS_APPLY
|
||||
#undef RVV_IF_M
|
||||
}
|
||||
|
||||
template <int32_t M, typename scalar_t>
|
||||
FORCE_INLINE void gemm_micro_rvv_mx32_ku4(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
scalar_t* __restrict__ curr_b_0 = b_ptr;
|
||||
scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
|
||||
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0, c_ptr, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_0 + RVV_MGEMM_N8,
|
||||
c_ptr + RVV_MGEMM_N8, lda, ldc, k, accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1, c_ptr + 16, lda, ldc, k,
|
||||
accum_c);
|
||||
gemm_micro_rvv_fma_mx8_ku4<M>(a_ptr, curr_b_1 + RVV_MGEMM_N8, c_ptr + 24, lda,
|
||||
ldc, k, accum_c);
|
||||
}
|
||||
|
||||
class TileGemmRVV {
|
||||
public:
|
||||
template <typename scalar_t>
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
switch (m) {
|
||||
case 1:
|
||||
gemm_micro_rvv_mx32_ku4<1>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_rvv_mx32_ku4<2>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 3:
|
||||
gemm_micro_rvv_mx32_ku4<3>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_rvv_mx32_ku4<4>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 5:
|
||||
gemm_micro_rvv_mx32_ku4<5>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 6:
|
||||
gemm_micro_rvv_mx32_ku4<6>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 7:
|
||||
gemm_micro_rvv_mx32_ku4<7>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 8:
|
||||
gemm_micro_rvv_mx32_ku4<8>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::RVV, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TileGemmRVV::gemm<scalar_t>(CPU_MICRO_GEMM_PARAMS);
|
||||
}
|
||||
|
||||
static void pack_weight(const scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK_EQ(output_size % 16, 0);
|
||||
for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
|
||||
const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
|
||||
scalar_t* __restrict__ curr_packed_weight =
|
||||
packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
|
||||
for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
|
||||
*curr_packed_weight = *curr_weight;
|
||||
|
||||
curr_packed_weight += 16;
|
||||
++curr_weight;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif // defined(__riscv_v)
|
||||
|
||||
#endif // CPU_MICRO_GEMM_RVV_HPP
|
||||
@@ -0,0 +1,136 @@
|
||||
#ifndef CPU_MICRO_GEMM_VEC_HPP
|
||||
#define CPU_MICRO_GEMM_VEC_HPP
|
||||
#include "cpu/micro_gemm/cpu_micro_gemm_impl.hpp"
|
||||
|
||||
namespace cpu_micro_gemm {
|
||||
namespace {
|
||||
// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
|
||||
template <typename scalar_t>
|
||||
class TileGemm82 {
|
||||
public:
|
||||
FORCE_INLINE static void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
switch (m) {
|
||||
case 1:
|
||||
gemm_micro<1>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro<2>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 3:
|
||||
gemm_micro<3>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro<4>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 5:
|
||||
gemm_micro<5>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 6:
|
||||
gemm_micro<6>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 7:
|
||||
gemm_micro<7>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
case 8:
|
||||
gemm_micro<8>(CPU_MICRO_GEMM_PARAMS);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
template <int32_t M>
|
||||
static void gemm_micro(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
using load_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
scalar_t* __restrict__ curr_b_0 = b_ptr;
|
||||
scalar_t* __restrict__ curr_b_1 = b_ptr + b_n_group_stride;
|
||||
float* __restrict__ curr_c_0 = c_ptr;
|
||||
float* __restrict__ curr_c_1 = c_ptr + 16;
|
||||
|
||||
vec_op::FP32Vec16 c_regs[M * 2];
|
||||
if (accum_c) {
|
||||
float* __restrict__ curr_m_c_0 = curr_c_0;
|
||||
float* __restrict__ curr_m_c_1 = curr_c_1;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i * 2] = vec_op::FP32Vec16(curr_m_c_0);
|
||||
c_regs[i * 2 + 1] = vec_op::FP32Vec16(curr_m_c_1);
|
||||
|
||||
// update
|
||||
curr_m_c_0 += ldc;
|
||||
curr_m_c_1 += ldc;
|
||||
});
|
||||
}
|
||||
|
||||
scalar_t* __restrict__ curr_a = a_ptr;
|
||||
for (int32_t k_idx = 0; k_idx < k; ++k_idx) {
|
||||
load_vec_t b_0_reg(curr_b_0);
|
||||
vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
|
||||
load_vec_t b_1_reg(curr_b_1);
|
||||
vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
|
||||
|
||||
scalar_t* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
scalar_t v = *curr_m_a;
|
||||
load_vec_t a_reg_original(v);
|
||||
vec_op::FP32Vec16 a_reg(a_reg_original);
|
||||
c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
|
||||
c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
|
||||
|
||||
// update
|
||||
curr_m_a += lda;
|
||||
});
|
||||
|
||||
// update
|
||||
curr_a += 1;
|
||||
curr_b_0 += 16;
|
||||
curr_b_1 += 16;
|
||||
}
|
||||
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
c_regs[i * 2].save(curr_c_0);
|
||||
c_regs[i * 2 + 1].save(curr_c_1);
|
||||
|
||||
// update
|
||||
curr_c_0 += ldc;
|
||||
curr_c_1 += ldc;
|
||||
});
|
||||
}
|
||||
};
|
||||
} // namespace
|
||||
|
||||
// Gemm kernel uses vector instructions, requires B matrix to be packed
|
||||
template <typename scalar_t>
|
||||
class MicroGemm<cpu_utils::ISA::VEC, scalar_t> {
|
||||
public:
|
||||
static constexpr int32_t MaxMSize = 8;
|
||||
static constexpr int32_t NSize = 32;
|
||||
static constexpr int32_t WeightOCGroupSize = 16;
|
||||
static constexpr bool PackA = false;
|
||||
|
||||
public:
|
||||
void gemm(DEFINE_CPU_MICRO_GEMM_PARAMS) {
|
||||
TileGemm82<scalar_t>::gemm(CPU_MICRO_GEMM_PARAMS);
|
||||
}
|
||||
|
||||
// Note: pack contiguous weight [output_size, input_size] as contiguous
|
||||
// packed weight [output_size / 16, input_size, 16]
|
||||
static void pack_weight(const scalar_t* __restrict__ weight,
|
||||
scalar_t* __restrict__ packed_weight,
|
||||
const int32_t output_size, const int32_t input_size) {
|
||||
TORCH_CHECK_EQ(output_size % 16, 0);
|
||||
for (int32_t o_idx = 0; o_idx < output_size; ++o_idx) {
|
||||
const scalar_t* __restrict__ curr_weight = weight + o_idx * input_size;
|
||||
scalar_t* __restrict__ curr_packed_weight =
|
||||
packed_weight + (o_idx / 16) * (16 * input_size) + o_idx % 16;
|
||||
for (int32_t i_idx = 0; i_idx < input_size; ++i_idx) {
|
||||
*curr_packed_weight = *curr_weight;
|
||||
|
||||
curr_packed_weight += 16;
|
||||
++curr_weight;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
} // namespace cpu_micro_gemm
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,383 @@
|
||||
#include "cpu_types.hpp"
|
||||
#include <float.h>
|
||||
|
||||
namespace {
|
||||
template <typename scalar_t>
|
||||
struct KernelVecType {
|
||||
using qk_load_vec_type = void;
|
||||
using qk_vec_type = void;
|
||||
using v_load_vec_type = void;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<float> {
|
||||
using qk_load_vec_type = vec_op::FP32Vec16;
|
||||
using qk_vec_type = vec_op::FP32Vec16;
|
||||
using v_load_vec_type = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<c10::Half> {
|
||||
using qk_load_vec_type = vec_op::FP16Vec16;
|
||||
using qk_vec_type = vec_op::FP32Vec16;
|
||||
using v_load_vec_type = vec_op::FP16Vec16;
|
||||
};
|
||||
|
||||
#ifdef __AVX512BF16__
|
||||
template <>
|
||||
struct KernelVecType<c10::BFloat16> {
|
||||
using qk_load_vec_type = vec_op::BF16Vec32;
|
||||
using qk_vec_type = vec_op::BF16Vec32;
|
||||
using v_load_vec_type = vec_op::BF16Vec16;
|
||||
};
|
||||
#else
|
||||
template <>
|
||||
struct KernelVecType<c10::BFloat16> {
|
||||
using qk_load_vec_type = vec_op::BF16Vec16;
|
||||
using qk_vec_type = vec_op::FP32Vec16;
|
||||
using v_load_vec_type = vec_op::BF16Vec16;
|
||||
};
|
||||
#endif
|
||||
|
||||
template <int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE, int HEAD_UNROLL,
|
||||
typename qk_vec_type>
|
||||
void mla_decode_block_head(
|
||||
const qk_vec_type* __restrict__ q_vecs, // [HEAD_UNROLL, head_dim]
|
||||
const qk_vec_type* __restrict__ k_vecs, // [block_size, head_dim]
|
||||
const vec_op::FP32Vec16* __restrict v_vecs_f32, // [block_size, v_head_dim]
|
||||
float* __restrict__ acc_out, // [HEAD_UNROLL, v_head_dim]
|
||||
float* __restrict__ acc_lse, // [HEAD_UNROLL]
|
||||
const float scale, const int num_tokens) {
|
||||
using f32_vec_type = vec_op::FP32Vec16;
|
||||
constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
|
||||
constexpr int V_NUM_ELEM = f32_vec_type::VEC_ELEM_NUM;
|
||||
|
||||
float logits[BLOCK_SIZE][HEAD_UNROLL] = {}; // initialize to zeros
|
||||
float max_val[HEAD_UNROLL];
|
||||
std::fill(max_val, max_val + HEAD_UNROLL, -FLT_MAX);
|
||||
|
||||
f32_vec_type acc_vec[BLOCK_SIZE][HEAD_UNROLL];
|
||||
for (int i = 0; i < HEAD_DIM; i += QK_NUM_ELEM) {
|
||||
// load to registers
|
||||
qk_vec_type q_vec[HEAD_UNROLL];
|
||||
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
|
||||
q_vec[unroll] =
|
||||
qk_vec_type{q_vecs[(i + unroll * HEAD_DIM) / QK_NUM_ELEM]};
|
||||
|
||||
for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
|
||||
qk_vec_type k_vec(k_vecs[(block_offset * HEAD_DIM + i) / QK_NUM_ELEM]);
|
||||
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
|
||||
vec_op::fma(acc_vec[block_offset][unroll], q_vec[unroll], k_vec);
|
||||
}
|
||||
}
|
||||
|
||||
for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
|
||||
const float acc = acc_vec[block_offset][unroll].reduce_sum() * scale;
|
||||
logits[block_offset][unroll] = acc;
|
||||
max_val[unroll] = std::max(max_val[unroll], acc);
|
||||
}
|
||||
}
|
||||
|
||||
float sum_exp[HEAD_UNROLL] = {};
|
||||
for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
|
||||
const float val =
|
||||
std::exp(logits[block_offset][unroll] - max_val[unroll]);
|
||||
logits[block_offset][unroll] = val;
|
||||
sum_exp[unroll] += val;
|
||||
}
|
||||
}
|
||||
|
||||
f32_vec_type this_out[V_HEAD_DIM / V_NUM_ELEM][HEAD_UNROLL];
|
||||
|
||||
for (int block_offset = 0; block_offset < num_tokens; ++block_offset) {
|
||||
// load to registers
|
||||
f32_vec_type scale_[HEAD_UNROLL];
|
||||
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
|
||||
scale_[unroll] =
|
||||
f32_vec_type{logits[block_offset][unroll] / sum_exp[unroll]};
|
||||
|
||||
for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) {
|
||||
f32_vec_type v_vec(
|
||||
v_vecs_f32[(block_offset * HEAD_DIM + i) / V_NUM_ELEM]);
|
||||
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll)
|
||||
vec_op::fma(this_out[i / V_NUM_ELEM][unroll], v_vec, scale_[unroll]);
|
||||
}
|
||||
}
|
||||
|
||||
// merge attention state
|
||||
// section 2.2 in https://arxiv.org/pdf/2501.01005
|
||||
f32_vec_type prev_scale[HEAD_UNROLL];
|
||||
f32_vec_type curr_scale[HEAD_UNROLL];
|
||||
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
|
||||
const float prev_lse = acc_lse[unroll];
|
||||
const float curr_lse = std::log(sum_exp[unroll]) +
|
||||
max_val[unroll]; // add back max_val to get true lse
|
||||
// softmax trick
|
||||
const float max_lse = std::max(prev_lse, curr_lse);
|
||||
const float prev_sum_exp = std::exp(prev_lse - max_lse);
|
||||
const float curr_sum_exp = std::exp(curr_lse - max_lse);
|
||||
|
||||
const float new_sum_exp = prev_sum_exp + curr_sum_exp;
|
||||
acc_lse[unroll] = std::log(new_sum_exp) + max_lse;
|
||||
|
||||
prev_scale[unroll] = f32_vec_type{prev_sum_exp / new_sum_exp};
|
||||
curr_scale[unroll] = f32_vec_type{curr_sum_exp / new_sum_exp};
|
||||
}
|
||||
|
||||
for (int i = 0; i < V_HEAD_DIM; i += V_NUM_ELEM) {
|
||||
#pragma unroll
|
||||
for (int unroll = 0; unroll < HEAD_UNROLL; ++unroll) {
|
||||
f32_vec_type o_vec(acc_out + i + V_HEAD_DIM * unroll);
|
||||
o_vec = o_vec * prev_scale[unroll] +
|
||||
this_out[i / V_NUM_ELEM][unroll] * curr_scale[unroll];
|
||||
o_vec.save(acc_out + i + V_HEAD_DIM * unroll);
|
||||
}
|
||||
}
|
||||
|
||||
q_vecs += HEAD_DIM / QK_NUM_ELEM * HEAD_UNROLL;
|
||||
acc_out += V_HEAD_DIM * HEAD_UNROLL;
|
||||
}
|
||||
|
||||
template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE,
|
||||
typename qk_vec_type>
|
||||
void mla_decode_block(
|
||||
const qk_vec_type* __restrict__ q_vecs, // [num_heads, head_dim]
|
||||
const scalar_t* __restrict__ kv_cache, // [block_size, head_dim]
|
||||
float* __restrict__ acc_out, // [num_heads, v_head_dim]
|
||||
float* __restrict__ acc_lse, // [num_heads]
|
||||
const int num_heads, const float scale, const int num_tokens) {
|
||||
using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type;
|
||||
static_assert(
|
||||
std::is_same<qk_vec_type,
|
||||
typename KernelVecType<scalar_t>::qk_vec_type>::value);
|
||||
using v_load_vec_type = typename KernelVecType<scalar_t>::v_load_vec_type;
|
||||
using f32_vec_type = vec_op::FP32Vec16;
|
||||
static_assert(qk_load_vec_type::VEC_ELEM_NUM == qk_vec_type::VEC_ELEM_NUM);
|
||||
static_assert(v_load_vec_type::VEC_ELEM_NUM == f32_vec_type::VEC_ELEM_NUM);
|
||||
constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
|
||||
constexpr int V_NUM_ELEM = v_load_vec_type::VEC_ELEM_NUM;
|
||||
|
||||
const qk_vec_type* k_vecs;
|
||||
const f32_vec_type* v_vecs_f32;
|
||||
float* kv_cache_f32 = nullptr;
|
||||
|
||||
if constexpr (!std::is_same<scalar_t, float>::value) {
|
||||
// convert KV cache block to FP32 to reuse it across query heads and
|
||||
// attn @ V computation, since FP16/BF16->FP32 is expensive.
|
||||
// TODO: move malloc outside of this fn to reuse across iterations.
|
||||
const int nbytes = BLOCK_SIZE * HEAD_DIM * sizeof(float);
|
||||
kv_cache_f32 = static_cast<float*>(std::aligned_alloc(64, nbytes));
|
||||
|
||||
for (int block_offset = 0; block_offset < num_tokens; ++block_offset)
|
||||
for (int i = 0; i < HEAD_DIM; i += V_NUM_ELEM) {
|
||||
v_load_vec_type kv_load_vec(kv_cache + block_offset * HEAD_DIM + i);
|
||||
f32_vec_type kv_vec_f32(kv_load_vec);
|
||||
kv_vec_f32.save(kv_cache_f32 + block_offset * HEAD_DIM + i);
|
||||
}
|
||||
|
||||
if constexpr (std::is_same<qk_load_vec_type, qk_vec_type>::value) {
|
||||
// for AVX512_BF16, Q @ K.T uses BF16 for K (no conversion)
|
||||
// NOTE: in this case, we only need to convert the V section to FP32.
|
||||
// But for simplicity, we will convert the whole KV block to FP32.
|
||||
k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache);
|
||||
} else {
|
||||
k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache_f32);
|
||||
}
|
||||
|
||||
// attn @ V always use FP32 for V, since attn is FP32.
|
||||
v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache_f32);
|
||||
|
||||
} else {
|
||||
// KV cache is FP32. don't need to do anything.
|
||||
k_vecs = reinterpret_cast<const qk_vec_type*>(kv_cache);
|
||||
v_vecs_f32 = reinterpret_cast<const f32_vec_type*>(kv_cache);
|
||||
}
|
||||
|
||||
// compute 2 heads at the same time to improve ILP and
|
||||
// take advantage of register cache for K and V.
|
||||
constexpr int HEAD_UNROLL = 2;
|
||||
for (int iter = 0; iter < num_heads / HEAD_UNROLL; ++iter) {
|
||||
mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, HEAD_UNROLL>(
|
||||
q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens);
|
||||
|
||||
q_vecs += HEAD_UNROLL * HEAD_DIM / QK_NUM_ELEM;
|
||||
acc_out += HEAD_UNROLL * V_HEAD_DIM;
|
||||
acc_lse += HEAD_UNROLL;
|
||||
}
|
||||
|
||||
// take care of the remaining heads
|
||||
for (int iter = 0; iter < num_heads % HEAD_UNROLL; ++iter) {
|
||||
mla_decode_block_head<HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE, 1>(
|
||||
q_vecs, k_vecs, v_vecs_f32, acc_out, acc_lse, scale, num_tokens);
|
||||
|
||||
q_vecs += HEAD_DIM / QK_NUM_ELEM;
|
||||
acc_out += V_HEAD_DIM;
|
||||
acc_lse += 1;
|
||||
}
|
||||
|
||||
if (kv_cache_f32 != nullptr) {
|
||||
std::free(kv_cache_f32);
|
||||
}
|
||||
}
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int HEAD_DIM, int V_HEAD_DIM, int BLOCK_SIZE>
|
||||
void mla_decode_kvcache_cpu_impl(
|
||||
scalar_t* __restrict__ out, // [num_seqs, num_heads, v_head_dim]
|
||||
const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_dim]
|
||||
const scalar_t* __restrict__ kv_cache, // [num_blocks, block_size,
|
||||
// head_dim]
|
||||
const int num_heads, const float scale,
|
||||
const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
|
||||
const int* __restrict__ seq_lens, // [num_seqs]
|
||||
const int max_num_blocks_per_seq, const int o_stride, const int q_stride,
|
||||
const int kv_stride, const int num_seqs) {
|
||||
using qk_load_vec_type = typename KernelVecType<scalar_t>::qk_load_vec_type;
|
||||
using qk_vec_type = typename KernelVecType<scalar_t>::qk_vec_type;
|
||||
constexpr int QK_NUM_ELEM = qk_vec_type::VEC_ELEM_NUM;
|
||||
|
||||
// shared across threads
|
||||
const int max_threads = cpu_utils::get_max_threads();
|
||||
const int acc_out_nbytes =
|
||||
max_threads * num_heads * V_HEAD_DIM * sizeof(float);
|
||||
float* acc_out = static_cast<float*>(std::aligned_alloc(64, acc_out_nbytes));
|
||||
std::vector<float> acc_lse(max_threads * num_heads);
|
||||
|
||||
// allocate memory to pre-convert query to FP32 later
|
||||
float* q_f32;
|
||||
constexpr bool PRE_CONVERT_QUERY =
|
||||
!std::is_same<scalar_t, float>::value &&
|
||||
std::is_same<qk_vec_type, vec_op::FP32Vec16>::value;
|
||||
if constexpr (PRE_CONVERT_QUERY) {
|
||||
const int q_f32_nbytes = num_heads * HEAD_DIM * sizeof(float);
|
||||
q_f32 = static_cast<float*>(std::aligned_alloc(64, q_f32_nbytes));
|
||||
}
|
||||
|
||||
#pragma omp parallel
|
||||
{
|
||||
const int num_threads = omp_get_num_threads();
|
||||
const int thread_id = omp_get_thread_num();
|
||||
float* __restrict__ acc_out_thread =
|
||||
acc_out + thread_id * num_heads * V_HEAD_DIM;
|
||||
float* __restrict__ acc_lse_thread = acc_lse.data() + thread_id * num_heads;
|
||||
|
||||
for (int seq_idx = 0; seq_idx < num_seqs; ++seq_idx) {
|
||||
// reset accumulator
|
||||
std::fill(acc_out_thread, acc_out_thread + num_heads * V_HEAD_DIM, 0.0f);
|
||||
std::fill(acc_lse_thread, acc_lse_thread + num_heads, -FLT_MAX);
|
||||
|
||||
const int seq_len = seq_lens[seq_idx];
|
||||
const int block_num = (seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
|
||||
const int last_block_size = seq_len - (block_num - 1) * BLOCK_SIZE;
|
||||
|
||||
const qk_vec_type* q_vecs;
|
||||
if constexpr (PRE_CONVERT_QUERY) {
|
||||
// pre-convert query to FP32 since FP16/BF16->FP32 is slow.
|
||||
#pragma omp for
|
||||
for (int i = 0; i < num_heads * HEAD_DIM; i += QK_NUM_ELEM) {
|
||||
qk_load_vec_type q_load_vec(q + seq_idx * q_stride + i);
|
||||
qk_vec_type q_vec(q_load_vec);
|
||||
q_vec.save(q_f32 + i);
|
||||
}
|
||||
q_vecs = reinterpret_cast<const qk_vec_type*>(q_f32);
|
||||
} else {
|
||||
q_vecs = reinterpret_cast<const qk_vec_type*>(q + seq_idx * q_stride);
|
||||
}
|
||||
|
||||
#pragma omp for
|
||||
for (int block_idx = 0; block_idx < block_num; ++block_idx) {
|
||||
const int physical_block_idx =
|
||||
block_tables[seq_idx * max_num_blocks_per_seq + block_idx];
|
||||
const int num_tokens =
|
||||
block_idx < block_num - 1 ? BLOCK_SIZE : last_block_size;
|
||||
|
||||
mla_decode_block<scalar_t, HEAD_DIM, V_HEAD_DIM, BLOCK_SIZE>(
|
||||
q_vecs, kv_cache + physical_block_idx * kv_stride, acc_out_thread,
|
||||
acc_lse_thread, num_heads, scale, num_tokens);
|
||||
}
|
||||
|
||||
// merge attention states across threads
|
||||
// section 2.2 in https://arxiv.org/pdf/2501.01005
|
||||
// each thread is responsible for 1 head
|
||||
#pragma omp for
|
||||
for (int head_idx = 0; head_idx < num_heads; ++head_idx) {
|
||||
float* acc_lse_head = acc_lse.data() + head_idx;
|
||||
float* acc_out_head = acc_out + head_idx * V_HEAD_DIM;
|
||||
|
||||
float max_val = -FLT_MAX;
|
||||
for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
|
||||
max_val = std::max(max_val, acc_lse_head[thread_id_ * num_heads]);
|
||||
}
|
||||
|
||||
float sum_exp = 0.0f;
|
||||
for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
|
||||
float val = std::exp(acc_lse_head[thread_id_ * num_heads] - max_val);
|
||||
acc_lse_head[thread_id_ * num_heads] = val;
|
||||
sum_exp += val;
|
||||
}
|
||||
|
||||
float inv_sum = 1.0f / sum_exp;
|
||||
float out_head[V_HEAD_DIM] = {};
|
||||
for (int thread_id_ = 0; thread_id_ < num_threads; ++thread_id_) {
|
||||
float scale_ = acc_lse_head[thread_id_ * num_heads] * inv_sum;
|
||||
for (int i = 0; i < V_HEAD_DIM; ++i) {
|
||||
out_head[i] +=
|
||||
acc_out_head[thread_id_ * num_heads * V_HEAD_DIM + i] * scale_;
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < V_HEAD_DIM; ++i) {
|
||||
vec_op::storeFP32(out_head[i], out + seq_idx * o_stride +
|
||||
head_idx * V_HEAD_DIM + i);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (PRE_CONVERT_QUERY) {
|
||||
std::free(q_f32);
|
||||
}
|
||||
std::free(acc_out);
|
||||
}
|
||||
|
||||
void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
|
||||
torch::Tensor& kv_cache, double scale,
|
||||
torch::Tensor& block_tables, torch::Tensor& seq_lens) {
|
||||
const int num_seqs = query.size(0);
|
||||
const int num_heads = query.size(1);
|
||||
const int head_dim = query.size(2);
|
||||
const int block_size = kv_cache.size(1);
|
||||
const int v_head_dim = out.size(2);
|
||||
|
||||
const int max_num_blocks_per_seq = block_tables.size(1);
|
||||
const int o_stride = out.stride(0);
|
||||
const int q_stride = query.stride(0);
|
||||
const int kv_stride = kv_cache.stride(0);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
query.scalar_type(), "mla_decode_kvcache_cpu_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(mla_decode_kvcache_cpu_impl)
|
||||
if (head_dim == 576 && v_head_dim == 512 && block_size == 16)
|
||||
mla_decode_kvcache_cpu_impl<scalar_t, 576, 512, 16>(
|
||||
out.data_ptr<scalar_t>(), query.data_ptr<scalar_t>(),
|
||||
kv_cache.data_ptr<scalar_t>(), num_heads, scale,
|
||||
block_tables.data_ptr<int>(), seq_lens.data_ptr<int>(),
|
||||
max_num_blocks_per_seq, o_stride, q_stride, kv_stride, num_seqs);
|
||||
else
|
||||
TORCH_CHECK(false, "Unsupported block size: ", block_size);
|
||||
CPU_KERNEL_GUARD_OUT(mla_decode_kvcache_cpu_impl)
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,366 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
namespace {
|
||||
template <typename scalar_t>
|
||||
void rotary_embedding_impl(
|
||||
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
|
||||
// [num_tokens]
|
||||
scalar_t* __restrict__ query, /// [batch_size, seq_len, num_heads,
|
||||
/// head_size] or [num_tokens, num_heads,
|
||||
/// head_size]
|
||||
scalar_t* __restrict__ key, // nullptr (optional) or
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
|
||||
// 2]
|
||||
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::vec_t<scalar_t>;
|
||||
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 scalar_t* cache_ptr,
|
||||
scalar_t* 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 scalar_vec_t cos(cache_ptr + x_index);
|
||||
const scalar_vec_t sin(cache_ptr + y_index);
|
||||
|
||||
const scalar_vec_t q_x(qk + out_x);
|
||||
const scalar_vec_t q_y(qk + out_y);
|
||||
|
||||
vec_op::FP32Vec8 fp32_cos(cos);
|
||||
vec_op::FP32Vec8 fp32_sin(sin);
|
||||
|
||||
vec_op::FP32Vec8 fp32_q_x(q_x);
|
||||
vec_op::FP32Vec8 fp32_q_y(q_y);
|
||||
|
||||
auto out1 = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin;
|
||||
scalar_vec_t(out1).save(qk + out_x);
|
||||
|
||||
auto out2 = fp32_q_y * fp32_cos + fp32_q_x * fp32_sin;
|
||||
scalar_vec_t(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 = cache_ptr[x_index];
|
||||
const float fp32_sin = cache_ptr[y_index];
|
||||
|
||||
const float fp32_q_x = qk[out_x];
|
||||
const float fp32_q_y = qk[out_y];
|
||||
|
||||
qk[out_x] = fp32_q_x * fp32_cos - fp32_q_y * fp32_sin;
|
||||
qk[out_y] = 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 scalar_t* 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 <>
|
||||
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
|
||||
// [num_tokens]
|
||||
scalar_t* __restrict__ query, /// [batch_size, seq_len, num_heads,
|
||||
/// head_size] or [num_tokens, num_heads,
|
||||
/// head_size]
|
||||
scalar_t* __restrict__ key, // nullptr (optional) or
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
|
||||
// 2]
|
||||
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 scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
const scalar_t* cos_cache_ptr = cache_ptr;
|
||||
const scalar_t* 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;
|
||||
scalar_t* 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 = cos_cache_ptr[rot_offset];
|
||||
const float sin = sin_cache_ptr[rot_offset];
|
||||
|
||||
const float x = head_query[x_index];
|
||||
const float y = head_query[y_index];
|
||||
|
||||
head_query[x_index] = x * cos - y * sin;
|
||||
head_query[y_index] = 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 scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
const scalar_t* cos_cache_ptr = cache_ptr;
|
||||
const scalar_t* 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;
|
||||
scalar_t* 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 = cos_cache_ptr[rot_offset];
|
||||
const float sin = sin_cache_ptr[rot_offset];
|
||||
|
||||
const float x = head_key[x_index];
|
||||
const float y = head_key[y_index];
|
||||
|
||||
head_key[x_index] = x * cos - y * sin;
|
||||
head_key[y_index] = y * cos + x * sin;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
std::optional<torch::Tensor> key, int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, bool is_neox,
|
||||
int64_t rope_dim_offset, bool inverse) {
|
||||
TORCH_CHECK(rope_dim_offset == 0,
|
||||
"rope_dim_offset != 0 is not supported on CPU");
|
||||
TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
|
||||
|
||||
int num_tokens = positions.numel();
|
||||
int rot_dim = cos_sin_cache.size(1);
|
||||
int num_heads = query.size(-1) / head_size;
|
||||
int num_kv_heads = key.has_value() ? key->size(-1) / head_size : num_heads;
|
||||
int64_t key_stride = key.has_value() ? key->stride(-2) : 0;
|
||||
int64_t query_stride = query.stride(-2);
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
query.scalar_type(), "rotary_embedding_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(rotary_embedding_impl)
|
||||
if (is_neox) {
|
||||
rotary_embedding_impl(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
|
||||
key_stride, num_heads, num_kv_heads, head_size, num_tokens);
|
||||
} else {
|
||||
rotary_embedding_gptj_impl(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
|
||||
key_stride, num_heads, num_kv_heads, head_size, num_tokens);
|
||||
}
|
||||
|
||||
CPU_KERNEL_GUARD_OUT(rotary_embedding_impl)
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
#include <ATen/native/CPUBlas.h>
|
||||
|
||||
// Unlike brgemm, PyTorch does not publicly expose at::native::cpublas::gemm
|
||||
// If OpenBLS is available in the PyTorch wheel, we rely on it for fast
|
||||
// bf16:bf16->fp32 GEMMs Otherwise, we fall back to PyTorch reference BLAS path.
|
||||
#if defined(VLLM_HAS_OPENBLAS)
|
||||
extern "C" void sbgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const at::BFloat16* a, int* lda,
|
||||
const at::BFloat16* b, int* ldb, float* beta, float* c,
|
||||
int* ldc);
|
||||
|
||||
extern "C" void sgemm_(char* transa, char* transb, int* m, int* n, int* k,
|
||||
float* alpha, const float* a, int* lda, const float* b,
|
||||
int* ldb, float* beta, float* c, int* ldc);
|
||||
|
||||
inline char blas_transpose(at::native::TransposeType trans) {
|
||||
switch (trans) {
|
||||
case at::native::TransposeType::NoTranspose:
|
||||
return 'n';
|
||||
case at::native::TransposeType::Transpose:
|
||||
return 't';
|
||||
case at::native::TransposeType::ConjTranspose:
|
||||
return 'c';
|
||||
}
|
||||
return 'n';
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const at::BFloat16* a,
|
||||
int64_t lda, const at::BFloat16* b, int64_t ldb,
|
||||
float beta, float* c, int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sbgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const float* a, int64_t lda,
|
||||
const float* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
char transa_ = blas_transpose(transa);
|
||||
char transb_ = blas_transpose(transb);
|
||||
int m_ = static_cast<int>(m);
|
||||
int n_ = static_cast<int>(n);
|
||||
int k_ = static_cast<int>(k);
|
||||
int lda_ = static_cast<int>(lda);
|
||||
int ldb_ = static_cast<int>(ldb);
|
||||
int ldc_ = static_cast<int>(ldc);
|
||||
sgemm_(&transa_, &transb_, &m_, &n_, &k_, &alpha, a, &lda_, b, &ldb_, &beta,
|
||||
c, &ldc_);
|
||||
}
|
||||
|
||||
inline void blas_gemm(at::native::TransposeType, at::native::TransposeType,
|
||||
int64_t, int64_t, int64_t, float, const at::Half*,
|
||||
int64_t, const at::Half*, int64_t, float, float*,
|
||||
int64_t) {
|
||||
TORCH_CHECK(false, "CPU OpenBLAS hgemm is not available.");
|
||||
}
|
||||
#else
|
||||
template <typename scalar_t>
|
||||
inline void blas_gemm(at::native::TransposeType transa,
|
||||
at::native::TransposeType transb, int64_t m, int64_t n,
|
||||
int64_t k, float alpha, const scalar_t* a, int64_t lda,
|
||||
const scalar_t* b, int64_t ldb, float beta, float* c,
|
||||
int64_t ldc) {
|
||||
auto gemm = at::native::cpublas::gemm_no_downcast_stub.DEFAULT;
|
||||
gemm(c10::CppTypeToScalarType<scalar_t>::value, transa, transb, m, n, k,
|
||||
at::Scalar(alpha), a, lda, b, ldb, at::Scalar(beta), c, ldc);
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,432 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/Parallel.h>
|
||||
|
||||
#if defined(_OPENMP)
|
||||
#include <omp.h>
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
// dispatch bool
|
||||
#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
|
||||
[&] { \
|
||||
if (BOOL_V) { \
|
||||
constexpr bool BOOL_NAME = true; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
constexpr bool BOOL_NAME = false; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
}()
|
||||
|
||||
#define AT_DISPATCH_BOOL2(BOOL_V1, BOOL_NAME1, BOOL_V2, BOOL_NAME2, ...) \
|
||||
[&] { \
|
||||
if (BOOL_V1) { \
|
||||
constexpr bool BOOL_NAME1 = true; \
|
||||
if (BOOL_V2) { \
|
||||
constexpr bool BOOL_NAME2 = true; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
constexpr bool BOOL_NAME2 = false; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
} else { \
|
||||
constexpr bool BOOL_NAME1 = false; \
|
||||
if (BOOL_V2) { \
|
||||
constexpr bool BOOL_NAME2 = true; \
|
||||
return __VA_ARGS__(); \
|
||||
} else { \
|
||||
constexpr bool BOOL_NAME2 = false; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
} \
|
||||
}()
|
||||
|
||||
// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
|
||||
#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \
|
||||
[&] { \
|
||||
switch (TYPE) { \
|
||||
case at::ScalarType::BFloat16: { \
|
||||
using packed_t = at::BFloat16; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Half: { \
|
||||
using packed_t = at::Half; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Char: { \
|
||||
using packed_t = int8_t; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Float8_e4m3fn: { \
|
||||
using packed_t = at::Float8_e4m3fn; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Byte: { \
|
||||
using packed_t = uint8_t; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
default: \
|
||||
TORCH_CHECK(false, "Unsupported floating data type.\n"); \
|
||||
} \
|
||||
}()
|
||||
|
||||
// Helper MICRO for CPU_DISPATCH_FLOATING_TYPES_EXT:
|
||||
// TYPE1: the primary dtype (input, output, weight);
|
||||
// TYPE2: defined as PARAM_T input
|
||||
#define CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, PARAM_T, ...) \
|
||||
switch (TYPE1) { \
|
||||
case at::ScalarType::BFloat16: { \
|
||||
using scalar_t = at::BFloat16; \
|
||||
using param_t = PARAM_T; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Half: { \
|
||||
using scalar_t = at::Half; \
|
||||
using param_t = PARAM_T; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Float: { \
|
||||
using scalar_t = float; \
|
||||
using param_t = PARAM_T; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
default: \
|
||||
TORCH_CHECK(false, "Unsupported floating data type."); \
|
||||
}
|
||||
|
||||
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
|
||||
// TYPE1: the primary dtype (input, output, weight);
|
||||
// TYPE2: defined as PARAM_T input
|
||||
#define CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, PARAM_T, ...) \
|
||||
switch (TYPE1) { \
|
||||
case at::ScalarType::BFloat16: { \
|
||||
using scalar_t = at::BFloat16; \
|
||||
using param_t = PARAM_T; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Half: { \
|
||||
using scalar_t = at::Half; \
|
||||
using param_t = PARAM_T; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
default: \
|
||||
TORCH_CHECK(false, "Unsupported floating data type."); \
|
||||
}
|
||||
|
||||
// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
|
||||
// TYPE1: the dtype both for scalar_t and param_t
|
||||
#define CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, ...) \
|
||||
switch (TYPE1) { \
|
||||
case at::ScalarType::BFloat16: { \
|
||||
using scalar_t = at::BFloat16; \
|
||||
using param_t = at::BFloat16; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
case at::ScalarType::Half: { \
|
||||
using scalar_t = at::Half; \
|
||||
using param_t = at::Half; \
|
||||
return __VA_ARGS__(); \
|
||||
} \
|
||||
default: \
|
||||
TORCH_CHECK(false, "Unsupported reduced floating data type."); \
|
||||
}
|
||||
|
||||
// dispatch with mixed dtypes (TYPE1, TYPE2):
|
||||
// TYPE1: the primary dtype (input, output, weight);
|
||||
// TYPE2: the secondary dtype (bias, etc.).
|
||||
#define CPU_DISPATCH_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
|
||||
[&] { \
|
||||
if (TYPE2 == at::kFloat) { \
|
||||
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, float, __VA_ARGS__) \
|
||||
} else if (TYPE2 == at::ScalarType::BFloat16) { \
|
||||
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::BFloat16, __VA_ARGS__) \
|
||||
} else if (TYPE2 == at::ScalarType::Half) { \
|
||||
CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::Half, __VA_ARGS__) \
|
||||
} else { \
|
||||
TORCH_CHECK(false, "Unsupported floating data type."); \
|
||||
} \
|
||||
}()
|
||||
|
||||
// dispatch with mixed dtypes (reduced one, no float for TYPE1) (TYPE1, TYPE2):
|
||||
// TYPE1: the primary dtype (input, output, weight);
|
||||
// TYPE2: the secondary dtype (bias, etc.).
|
||||
#define CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
|
||||
[&] { \
|
||||
if (TYPE2 == at::kFloat) { \
|
||||
CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, float, __VA_ARGS__) \
|
||||
} else { \
|
||||
TORCH_CHECK(TYPE1 == TYPE2); \
|
||||
CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, __VA_ARGS__) \
|
||||
} \
|
||||
}()
|
||||
|
||||
#define UNUSED(x) (void)(x)
|
||||
|
||||
#define CHECK_CPU(x) TORCH_CHECK(x.device().type() == at::kCPU, #x " must be a CPU tensor")
|
||||
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_LAST_DIM_CONTIGUOUS(x) \
|
||||
TORCH_CHECK(x.strides()[x.strides().size() - 1] == 1, #x "must be contiguous at last dimension")
|
||||
|
||||
#define CHECK_INPUT(x) \
|
||||
CHECK_CPU(x); \
|
||||
CHECK_CONTIGUOUS(x)
|
||||
#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \
|
||||
CHECK_CPU(x); \
|
||||
CHECK_LAST_DIM_CONTIGUOUS(x)
|
||||
|
||||
#define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")
|
||||
|
||||
#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
|
||||
#define CHECK_GT(a, b) TORCH_CHECK((a) > (b), "CHECK_GT(" #a ", " #b ") failed. ", a, " vs ", b)
|
||||
#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)
|
||||
|
||||
template <bool is_only_lastdim_contiguous>
|
||||
static inline void CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const at::IntArrayRef sizes, at::ScalarType st) {
|
||||
TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
|
||||
TORCH_CHECK(tensor.scalar_type() == st, "Input tensor dtype mismatch");
|
||||
if constexpr (is_only_lastdim_contiguous) {
|
||||
CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
|
||||
} else {
|
||||
CHECK_INPUT(tensor);
|
||||
}
|
||||
}
|
||||
|
||||
// [NB] Parallel Routines
|
||||
//
|
||||
// * at::parallel_for - applies for most of generic use cases, this will be compiled
|
||||
// against openmp in default torch release.
|
||||
//
|
||||
// * parallel_for - same function as above, can choose payload partition scheme in
|
||||
// balance211.
|
||||
//
|
||||
// * parallel_2d - parallel for 2 dimensions, used in GEMM, etc.
|
||||
// this one will do payload balance across 2 dimensions.
|
||||
//
|
||||
|
||||
// grain size for each thread
|
||||
constexpr int GRAIN_SIZE = 1024;
|
||||
|
||||
template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
|
||||
inline T div_up(T x, T y) {
|
||||
return (x + y - 1) / y;
|
||||
}
|
||||
|
||||
// you can only use at::get_thread_num() with at::parallel_for()
|
||||
// as it is lazy initialized, otherwise it will always return 0.
|
||||
inline int get_thread_num() {
|
||||
#if defined(_OPENMP)
|
||||
return omp_get_thread_num();
|
||||
#else
|
||||
return 0;
|
||||
#endif
|
||||
}
|
||||
|
||||
// balance payload across each thread
|
||||
template <typename T>
|
||||
inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
|
||||
#if 0
|
||||
// onednn partition pattern
|
||||
T& n_my = n_end;
|
||||
if (nth <= 1 || n == 0) {
|
||||
n_start = 0;
|
||||
n_my = n;
|
||||
} else {
|
||||
T n1 = div_up(n, nth);
|
||||
T n2 = n1 - 1;
|
||||
T T1 = n - n2 * nth;
|
||||
n_my = ith < T1 ? n1 : n2;
|
||||
n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
|
||||
}
|
||||
n_end += n_start;
|
||||
#else
|
||||
// pytorch aten partition pattern
|
||||
T n_my = div_up(n, nth);
|
||||
n_start = ith * n_my;
|
||||
n_end = std::min(n_start + n_my, n);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename func_t>
|
||||
inline void parallel_for(int n, const func_t& f) {
|
||||
#if defined(_OPENMP)
|
||||
#pragma omp parallel
|
||||
{
|
||||
int nth = omp_get_num_threads();
|
||||
int ith = omp_get_thread_num();
|
||||
int tbegin, tend;
|
||||
balance211(n, nth, ith, tbegin, tend);
|
||||
f(tbegin, tend);
|
||||
}
|
||||
#else
|
||||
f(0, n);
|
||||
#endif
|
||||
}
|
||||
|
||||
// for 1d parallel, use `actual_nth`
|
||||
// for 2d parallel, use even nths, e.g. 43->42
|
||||
int inline adjust_num_threads(int m) {
|
||||
int actual_nth = at::get_num_threads();
|
||||
if (m == 1) {
|
||||
return actual_nth;
|
||||
}
|
||||
return std::max(1, (actual_nth >> 1) * 2);
|
||||
}
|
||||
|
||||
template <typename func_t>
|
||||
inline void parallel_2d(int m, int n, const func_t& f) {
|
||||
// make sure we have even num_threads
|
||||
int nth = adjust_num_threads(m);
|
||||
|
||||
// [NOTE] thread blocking:
|
||||
//
|
||||
// 1) prefer square block per thread
|
||||
// 2) use even number of CPU cores
|
||||
// 3) use all `num_threads` cores
|
||||
//
|
||||
// we have:
|
||||
// TM * TN = T
|
||||
// BM / TM = BN / TN
|
||||
// then:
|
||||
// TM = ((BM / BN) * T) ^ 0.5
|
||||
//
|
||||
float r = float(m) / n;
|
||||
int nth_m = std::ceil(std::sqrt(r * nth));
|
||||
int nth_n = 1;
|
||||
for (; nth_m > 0; --nth_m) {
|
||||
nth_n = nth / nth_m;
|
||||
if (nth_m * nth_n == nth) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(_OPENMP)
|
||||
#pragma omp parallel num_threads(nth)
|
||||
{
|
||||
int ith = omp_get_thread_num();
|
||||
int ith_m = ith / nth_n;
|
||||
int ith_n = ith % nth_n;
|
||||
|
||||
int thread_block_m = div_up(m, nth_m);
|
||||
int thread_block_n = div_up(n, nth_n);
|
||||
|
||||
int begin_m = ith_m * thread_block_m;
|
||||
int end_m = std::min(m, begin_m + thread_block_m);
|
||||
int begin_n = ith_n * thread_block_n;
|
||||
int end_n = std::min(n, begin_n + thread_block_n);
|
||||
|
||||
f(begin_m, end_m, begin_n, end_n);
|
||||
}
|
||||
#else
|
||||
f(0, m, 0, n);
|
||||
#endif
|
||||
}
|
||||
|
||||
// limit max cache blocks
|
||||
// when we need to do pre-unpack for weights, e.g. fp8
|
||||
#define MAX_CACHE_BLOCK_SIZE 4
|
||||
|
||||
template <typename T>
|
||||
inline int get_cache_blocks(int chunk_size) {
|
||||
// L2 2MB and ratio of 50%
|
||||
const int L2_size = 2048 * 1024 >> 1;
|
||||
return std::max(1, int(L2_size / (chunk_size * sizeof(T))));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
|
||||
// fp8 uses bf16 as accumulate type
|
||||
int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
|
||||
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
|
||||
}
|
||||
|
||||
template <>
|
||||
inline int get_cache_blocks<uint8_t>(int chunk_size) {
|
||||
// mxfp4 uses bf16 as accumulate type
|
||||
int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
|
||||
return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
|
||||
}
|
||||
|
||||
// 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
|
||||
template <typename T, typename func_t>
|
||||
inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
|
||||
// get number of blocks for L2 in most inner loop
|
||||
int64_t cache_blocks_nb = get_cache_blocks<T>(chunk_size);
|
||||
|
||||
// loop order: [NB / cache_blocks_nb, MB, cache_blocks_nb]
|
||||
// TODO: implement reverse order of [MB / cache_blocks_mb, NB, cache_blocks_mb]
|
||||
for (int64_t nbb = nb0; nbb < nb1; nbb += cache_blocks_nb) {
|
||||
for (int64_t mb = mb0; mb < mb1; ++mb) {
|
||||
for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, nb1); ++nb) {
|
||||
f(mb, nb, nb - nbb);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// data indexing for dimension collapse
|
||||
template <typename T>
|
||||
inline T data_index_init(T offset) {
|
||||
return offset;
|
||||
}
|
||||
|
||||
template <typename T, typename... Args>
|
||||
inline T data_index_init(T offset, T& x, const T& X, Args&&... args) {
|
||||
offset = data_index_init(offset, std::forward<Args>(args)...);
|
||||
x = offset % X;
|
||||
return offset / X;
|
||||
}
|
||||
|
||||
inline bool data_index_step() {
|
||||
return true;
|
||||
}
|
||||
|
||||
template <typename T, typename... Args>
|
||||
inline bool data_index_step(T& x, const T& X, Args&&... args) {
|
||||
if (data_index_step(std::forward<Args>(args)...)) {
|
||||
x = ((x + 1) == X) ? 0 : (x + 1);
|
||||
return x == 0;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// forced unroll for perf critical path
|
||||
|
||||
#if __has_attribute(always_inline)
|
||||
#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
|
||||
#else
|
||||
#define ALWAYS_INLINE inline
|
||||
#endif
|
||||
|
||||
template <int n>
|
||||
struct Unroll {
|
||||
template <typename Func, typename... Args>
|
||||
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
|
||||
Unroll<n - 1>{}(f, args...);
|
||||
f(std::integral_constant<int, n - 1>{}, args...);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct Unroll<1> {
|
||||
template <typename Func, typename... Args>
|
||||
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
|
||||
f(std::integral_constant<int, 0>{}, args...);
|
||||
}
|
||||
};
|
||||
|
||||
// conditional data ptr for optional tensor
|
||||
template <typename T>
|
||||
inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
|
||||
return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
@@ -0,0 +1,722 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "common.h"
|
||||
#include "gemm.h"
|
||||
#include "vec.h"
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_stub(scalar_t* __restrict__ y, const scalar_t* __restrict__ x, int64_t size) {
|
||||
using Vec = at::vec::Vectorized<scalar_t>;
|
||||
const bool is_padding = (x == nullptr);
|
||||
for (int64_t d = 0; d < size; d += Vec::size()) {
|
||||
Vec data_vec = is_padding ? Vec(0.f) : Vec::loadu(x + d);
|
||||
data_vec.store(y + d);
|
||||
}
|
||||
}
|
||||
|
||||
// no remainder
|
||||
template <typename scalar_t>
|
||||
void inline update_conv_state(
|
||||
scalar_t* __restrict__ conv_states,
|
||||
const scalar_t* __restrict__ input,
|
||||
int64_t width,
|
||||
int64_t dim,
|
||||
int64_t seqlen,
|
||||
bool has_initial_states) {
|
||||
// width for `conv_states`
|
||||
int64_t width1 = width - 1;
|
||||
int64_t w = 0;
|
||||
for (; w < width1 - seqlen; ++w) {
|
||||
scalar_t* y = conv_states + w * dim;
|
||||
const scalar_t* x = has_initial_states ? conv_states + (w + seqlen) * dim : nullptr;
|
||||
copy_stub(y, x, dim);
|
||||
}
|
||||
for (; w < width1; ++w) {
|
||||
scalar_t* y = conv_states + w * dim;
|
||||
const scalar_t* x = input + (w + seqlen - width1) * dim;
|
||||
copy_stub(y, x, dim);
|
||||
}
|
||||
}
|
||||
|
||||
// A : [M, BLOCK_N]
|
||||
// B : [BLOCK_N, K], prepacked as [K/2, BLOCK_N, 2]
|
||||
// C : [M, BLOCK_N]
|
||||
// bias : [BLOCK_N]
|
||||
//
|
||||
// lda : leading dimension of `input` and `out`
|
||||
//
|
||||
template <typename scalar_t, int K, int BLOCK_N, bool has_bias, bool has_silu>
|
||||
struct tinygemm_kernel {
|
||||
static inline void apply(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
const scalar_t* __restrict__ bias,
|
||||
const scalar_t* __restrict__ conv_states,
|
||||
bool has_initial_state,
|
||||
int64_t M,
|
||||
int64_t lda,
|
||||
bool is_first_token) {
|
||||
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <int K, int BLOCK_N, bool has_bias, bool has_silu>
|
||||
struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
|
||||
static inline void apply(
|
||||
const at::BFloat16* __restrict__ A,
|
||||
const at::BFloat16* __restrict__ B,
|
||||
at::BFloat16* __restrict__ C,
|
||||
const at::BFloat16* __restrict__ bias,
|
||||
const at::BFloat16* __restrict__ conv_states,
|
||||
bool has_initial_state,
|
||||
int64_t M,
|
||||
int64_t lda,
|
||||
bool is_first_token) {
|
||||
assert(K == 4);
|
||||
constexpr int ROWS = K;
|
||||
constexpr int COLS = BLOCK_N / block_size_n();
|
||||
|
||||
// leading dimension size for b for next block [K/2, 32, 2]
|
||||
constexpr int ldb = block_size_n() * K;
|
||||
|
||||
__m512bh va[ROWS * COLS];
|
||||
__m512bh vb[ROWS * COLS];
|
||||
__m512 vc[COLS * 2];
|
||||
|
||||
// k: {-3, -2, -1} -> {0, 1, 2}
|
||||
auto set_conv_states = [&](int k, int col) -> __m512i {
|
||||
return has_initial_state ? _mm512_loadu_si512(conv_states + (k + K - 1) * lda + col * 32)
|
||||
: _mm512_setzero_si512();
|
||||
};
|
||||
|
||||
#define MM512_LOAD_A(idx) \
|
||||
((idx) < 0 && is_first_token) ? (__m512bh)(set_conv_states((idx), col)) \
|
||||
: (__m512bh)(_mm512_loadu_si512(A + (idx) * lda + col * 32))
|
||||
|
||||
#define MM512_PACK_A(ap, bp, a, b) \
|
||||
do { \
|
||||
__m512i r0 = (__m512i)(a); \
|
||||
__m512i r1 = (__m512i)(b); \
|
||||
__m512i d0 = _mm512_unpacklo_epi16(r0, r1); \
|
||||
__m512i d1 = _mm512_unpackhi_epi16(r0, r1); \
|
||||
r0 = _mm512_shuffle_i32x4(d0, d1, 0x88); \
|
||||
r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd); \
|
||||
(ap) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0x88); \
|
||||
(bp) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0xdd); \
|
||||
} while (0)
|
||||
|
||||
// step 0 : preload a at time step [-3][-2][-1]
|
||||
auto preloada = [&](auto i) {
|
||||
constexpr int col = i;
|
||||
int64_t m = 0;
|
||||
va[1 * COLS + col] = MM512_LOAD_A(m - 3);
|
||||
va[2 * COLS + col] = MM512_LOAD_A(m - 2);
|
||||
va[3 * COLS + col] = MM512_LOAD_A(m - 1);
|
||||
};
|
||||
Unroll<COLS>{}(preloada);
|
||||
|
||||
auto loada = [&](auto i, int64_t m) {
|
||||
constexpr int col = i;
|
||||
// update previous time step
|
||||
va[0 * COLS + col] = va[1 * COLS + col];
|
||||
va[1 * COLS + col] = va[2 * COLS + col];
|
||||
va[2 * COLS + col] = va[3 * COLS + col];
|
||||
// load current time step
|
||||
va[3 * COLS + col] = MM512_LOAD_A(m);
|
||||
};
|
||||
|
||||
// step 1 : load weight for just once
|
||||
auto loadb = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
vb[row * COLS + col] = (__m512bh)(_mm512_loadu_si512(B + col * ldb + row * 32));
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(loadb);
|
||||
|
||||
// [NB] accumulates 4x32 bfloat16 blocks
|
||||
//
|
||||
// +------------+------------+
|
||||
// | col0 | col1 |
|
||||
// +------------+------------+
|
||||
// | va0 va1 | va0 va1 |
|
||||
// | va2 va3 | va2 va3 |
|
||||
// +------------+------------+
|
||||
// | vc0 vc1 | vc0 vc1 |
|
||||
// +------------+------------+
|
||||
//
|
||||
// * va and vb shares the same memory layout
|
||||
// * block_n 32 with 4 rows equals to 4 registers
|
||||
// * 37 uops with avx512bf16 v.s. 57 uops with avx512f
|
||||
//
|
||||
auto compute = [&](auto i) {
|
||||
constexpr int col = i;
|
||||
|
||||
// init accumulators
|
||||
if constexpr (has_bias) {
|
||||
__m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(bias + col * 32));
|
||||
vc[col * 2 + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
|
||||
vc[col * 2 + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
|
||||
} else {
|
||||
vc[col * 2 + 0] = _mm512_set1_ps(0.f);
|
||||
vc[col * 2 + 1] = _mm512_set1_ps(0.f);
|
||||
}
|
||||
|
||||
// convert to vnni2 format
|
||||
__m512bh va0, va1, va2, va3;
|
||||
MM512_PACK_A(va0, va1, va[0 * COLS + col], va[1 * COLS + col]);
|
||||
MM512_PACK_A(va2, va3, va[2 * COLS + col], va[3 * COLS + col]);
|
||||
|
||||
// accumulate
|
||||
vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va0, vb[0 * COLS + col]);
|
||||
vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va2, vb[2 * COLS + col]);
|
||||
vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va1, vb[1 * COLS + col]);
|
||||
vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va3, vb[3 * COLS + col]);
|
||||
};
|
||||
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
using bVec = at::vec::Vectorized<at::BFloat16>;
|
||||
const fVec one = fVec(1.f);
|
||||
auto storec = [&](auto i, int64_t m) {
|
||||
constexpr int col = i;
|
||||
fVec x0 = fVec(vc[col * 2 + 0]);
|
||||
fVec x1 = fVec(vc[col * 2 + 1]);
|
||||
if constexpr (has_silu) {
|
||||
x0 = x0 / (one + x0.neg().exp_u20());
|
||||
x1 = x1 / (one + x1.neg().exp_u20());
|
||||
}
|
||||
bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
|
||||
out_vec.store(C + m * lda + col * 32);
|
||||
};
|
||||
|
||||
for (int64_t m = 0; m < M; ++m) {
|
||||
// step 3.a : load a at current time step
|
||||
Unroll<COLS>{}(loada, m);
|
||||
// step 3.b : accumulate for window size (4)
|
||||
Unroll<COLS>{}(compute);
|
||||
// step 3.c : store c at current time step
|
||||
Unroll<COLS>{}(storec, m);
|
||||
}
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define LAUNCH_TINYGEMM_KERNEL(K, NB_SIZE) \
|
||||
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
|
||||
input + bs * seqlen * dim + mb_start * dim + nb_start, \
|
||||
weight + nb_start * width, \
|
||||
out + bs * seqlen * dim + mb_start * dim + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
|
||||
has_initial_states_value, \
|
||||
mb_size, \
|
||||
dim, \
|
||||
mb_start == 0);
|
||||
|
||||
template <typename scalar_t>
|
||||
void causal_conv1d_fwd_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
const scalar_t* __restrict__ weight,
|
||||
const scalar_t* __restrict__ bias,
|
||||
scalar_t* __restrict__ conv_states,
|
||||
const int32_t* __restrict__ conv_indices,
|
||||
const bool* __restrict__ has_initial_state,
|
||||
bool silu_activation,
|
||||
int64_t batch,
|
||||
int64_t dim,
|
||||
int64_t seqlen,
|
||||
int64_t width,
|
||||
int64_t num_seq_blocks,
|
||||
int64_t conv_state_slot_stride) {
|
||||
// handle 32 x 64 per block
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n() * 2;
|
||||
const int64_t NB = div_up(dim, BLOCK_N);
|
||||
|
||||
const int64_t num_blocks_per_seq = div_up(seqlen, BLOCK_M);
|
||||
const bool has_conv_states = conv_states != nullptr;
|
||||
const bool has_conv_indices = conv_indices != nullptr;
|
||||
|
||||
// parallel on [batch, seq, NB]
|
||||
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
|
||||
at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t mb{0}, nb{0};
|
||||
data_index_init(begin, mb, num_seq_blocks, nb, NB);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
int64_t bs = mb / num_blocks_per_seq;
|
||||
|
||||
int64_t mb_start = (mb % num_blocks_per_seq) * BLOCK_M;
|
||||
int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
|
||||
|
||||
const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
|
||||
switch (width << 4 | nb_size >> 4) {
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL(4, 32);
|
||||
break;
|
||||
case 0x44:
|
||||
LAUNCH_TINYGEMM_KERNEL(4, 64);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
|
||||
}
|
||||
|
||||
// move to the next index
|
||||
data_index_step(mb, num_seq_blocks, nb, NB);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// update conv_states if necessary
|
||||
if (has_conv_states) {
|
||||
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t bs = begin; bs < end; ++bs) {
|
||||
update_conv_state(
|
||||
conv_states + bs * conv_state_slot_stride, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
|
||||
tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
|
||||
input + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
weight + nb_start * width, \
|
||||
out + batch_offset * dim + mb_start * dim + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \
|
||||
has_initial_states_value, \
|
||||
mb_size, \
|
||||
dim, \
|
||||
mb_start == 0);
|
||||
|
||||
template <typename scalar_t>
|
||||
void causal_conv1d_fwd_varlen_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
const scalar_t* __restrict__ weight,
|
||||
const scalar_t* __restrict__ bias,
|
||||
scalar_t* __restrict__ conv_states,
|
||||
const int32_t* __restrict__ query_start_loc,
|
||||
const int32_t* __restrict__ conv_indices,
|
||||
const bool* __restrict__ has_initial_state,
|
||||
const int32_t* __restrict__ block_indices,
|
||||
bool silu_activation,
|
||||
int64_t batch,
|
||||
int64_t dim,
|
||||
int64_t width,
|
||||
int64_t num_seq_blocks,
|
||||
int64_t conv_state_slot_stride) {
|
||||
// handle 32 x 64 per block
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n() * 2;
|
||||
const int64_t NB = div_up(dim, BLOCK_N);
|
||||
|
||||
const bool has_conv_states = conv_states != nullptr;
|
||||
const bool has_conv_indices = conv_indices != nullptr;
|
||||
|
||||
// parallel on [batch, seq, NB]
|
||||
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
|
||||
at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t mb{0}, nb{0};
|
||||
data_index_init(begin, mb, num_seq_blocks, nb, NB);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
int32_t bs = block_indices[mb * 2 + 0];
|
||||
int32_t batch_offset = query_start_loc[bs];
|
||||
int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
|
||||
|
||||
int64_t mb_start = block_indices[mb * 2 + 1] * BLOCK_M;
|
||||
int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
|
||||
|
||||
const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
|
||||
switch (width << 4 | nb_size >> 4) {
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
|
||||
break;
|
||||
case 0x44:
|
||||
LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 64);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
|
||||
}
|
||||
|
||||
// move to the next index
|
||||
data_index_step(mb, num_seq_blocks, nb, NB);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// update conv_states if necessary
|
||||
if (has_conv_states) {
|
||||
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t bs = begin; bs < end; ++bs) {
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
|
||||
int32_t batch_offset = query_start_loc[bs];
|
||||
update_conv_state(
|
||||
conv_states + conv_state_index * conv_state_slot_stride,
|
||||
input + batch_offset * dim,
|
||||
width,
|
||||
dim,
|
||||
seqlen,
|
||||
has_initial_state[bs]);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void causal_conv1d_update_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
scalar_t* __restrict__ conv_states,
|
||||
const scalar_t* __restrict__ weight,
|
||||
const scalar_t* __restrict__ bias,
|
||||
const int32_t* __restrict__ conv_indices,
|
||||
bool silu_activation,
|
||||
int64_t batch,
|
||||
int64_t dim,
|
||||
int64_t seqlen,
|
||||
int64_t width,
|
||||
int64_t conv_state_slot_stride) {
|
||||
// handle 32 x 64 per block
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n() * 2;
|
||||
const int64_t NB = div_up(dim, BLOCK_N);
|
||||
|
||||
const bool has_conv_states = conv_states != nullptr;
|
||||
const bool has_conv_indices = conv_indices != nullptr;
|
||||
|
||||
// parallel on [batch, NB]
|
||||
AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
|
||||
at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t bs{0}, nb{0};
|
||||
data_index_init(begin, bs, batch, nb, NB);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
int64_t mb_start = 0;
|
||||
int64_t mb_size = 1;
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
|
||||
|
||||
const bool has_initial_states_value = true;
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
|
||||
switch (width << 4 | nb_size >> 4) {
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL(4, 32);
|
||||
break;
|
||||
case 0x44:
|
||||
LAUNCH_TINYGEMM_KERNEL(4, 64);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
|
||||
}
|
||||
|
||||
// move to the next index
|
||||
data_index_step(bs, batch, nb, NB);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*conv_state_slot_stride + (w) * dim
|
||||
|
||||
// update conv_states
|
||||
at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t bs = begin; bs < end; ++bs) {
|
||||
// update old states, range [1, width - 1)
|
||||
int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
|
||||
for (int64_t w = 1; w < width - 1; ++w) {
|
||||
std::memcpy(CONV_STATE_INDEXR(w - 1), CONV_STATE_INDEXR(w), dim * sizeof(scalar_t));
|
||||
}
|
||||
// copy new states
|
||||
std::memcpy(CONV_STATE_INDEXR(width - 2), input + bs * dim, dim * sizeof(scalar_t));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// from [dim, width] or [N, K]
|
||||
// to [N/BLOCK_N, K/2, BLOCK_N, 2]
|
||||
at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight) {
|
||||
CHECK_INPUT(weight);
|
||||
|
||||
int64_t dim = weight.size(0);
|
||||
int64_t width = weight.size(1);
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");
|
||||
TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);
|
||||
|
||||
const int64_t N = dim, K2 = width >> 1;
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
auto packed_weight = at::empty_like(weight);
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {
|
||||
// cast to float32 as vnni size is 2
|
||||
const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());
|
||||
float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());
|
||||
|
||||
at::parallel_for(0, NB * K2 * BLOCK_N, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t nb{0}, k2{0}, n{0};
|
||||
data_index_init(begin, nb, NB, k2, K2, n, BLOCK_N);
|
||||
|
||||
// TODO: optimize this if we need to online prepacking.
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
packed_data[i] = w_data[nb * BLOCK_N * K2 + n * K2 + k2];
|
||||
|
||||
// move to the next index
|
||||
data_index_step(nb, NB, k2, K2, n, BLOCK_N);
|
||||
}
|
||||
});
|
||||
});
|
||||
return packed_weight;
|
||||
}
|
||||
|
||||
#define CHECK_OPTIONAL_SHAPE_DTYPE(OPT, SIZE, DTYPE) \
|
||||
if (OPT.has_value()) { \
|
||||
const auto tensor = OPT.value(); \
|
||||
CHECK_CONTIGUOUS(tensor); \
|
||||
CHECK_EQ(tensor.size(0), SIZE); \
|
||||
CHECK_EQ(tensor.scalar_type(), DTYPE); \
|
||||
}
|
||||
|
||||
template <int BLOCK_M>
|
||||
int64_t get_block_count(const std::optional<at::Tensor>& offsets, int64_t batch, int64_t seqlen) {
|
||||
if (offsets.has_value()) {
|
||||
const int32_t* offsets_data = offsets.value().data_ptr<int32_t>();
|
||||
int32_t num_seq_blocks = 0;
|
||||
for (int64_t row = 0; row < batch; ++row) {
|
||||
num_seq_blocks += div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
|
||||
}
|
||||
return num_seq_blocks;
|
||||
}
|
||||
return batch * div_up(seqlen, int64_t(BLOCK_M));
|
||||
}
|
||||
|
||||
template <int BLOCK_M>
|
||||
at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t num_seq_blocks) {
|
||||
if (!offsets.has_value()) {
|
||||
return at::Tensor();
|
||||
}
|
||||
|
||||
const at::Tensor& offsets_ = offsets.value();
|
||||
at::Tensor indices = at::empty({num_seq_blocks, 2}, offsets_.options());
|
||||
|
||||
int64_t batch = offsets_.size(0) - 1;
|
||||
|
||||
const int32_t* offsets_data = offsets_.data_ptr<int32_t>();
|
||||
int32_t* indices_data = indices.data_ptr<int32_t>();
|
||||
|
||||
int64_t idx = 0;
|
||||
for (int32_t row = 0; row < batch; ++row) {
|
||||
int32_t blocks = div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
|
||||
|
||||
for (int32_t col = 0; col < blocks; ++col) {
|
||||
indices_data[idx * 2 + 0] = row;
|
||||
indices_data[idx * 2 + 1] = col;
|
||||
idx++;
|
||||
}
|
||||
}
|
||||
return indices;
|
||||
}
|
||||
|
||||
// API aligned with GPUs
|
||||
//
|
||||
// x: (batch, dim, seqlen) or (dim, cu_seq_len) for varlen
|
||||
// weight: (dim, width)
|
||||
// bias: (dim,)
|
||||
// query_start_loc: (batch + 1) int32
|
||||
// cache_indices: (batch) int32
|
||||
// has_initial_state: (batch) bool
|
||||
// conv_states: (..., dim, width - 1) itype
|
||||
// activation: either None or "silu" or "swish"
|
||||
// pad_slot_id: int
|
||||
//
|
||||
at::Tensor causal_conv1d_fwd_cpu(
|
||||
const at::Tensor& x,
|
||||
const at::Tensor& weight,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
const std::optional<at::Tensor>& conv_states,
|
||||
const std::optional<at::Tensor>& query_start_loc,
|
||||
const std::optional<at::Tensor>& conv_state_indices,
|
||||
const std::optional<at::Tensor>& has_initial_state,
|
||||
bool silu_activation,
|
||||
int64_t pad_slot_id,
|
||||
bool is_vnni) {
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
|
||||
|
||||
const bool is_var_seqlen = query_start_loc.has_value();
|
||||
const int64_t input_ndim = is_var_seqlen ? 2 : 3;
|
||||
TORCH_CHECK(x.dim() == input_ndim, "causal_conv1d_fwd_cpu: expect x to be ", input_ndim, "D tensor.");
|
||||
TORCH_CHECK(x.stride(-2) == 1 && x.stride(-1) == x.size(-2), "causal_conv1d_fwd_cpu: expect x to be transposed.");
|
||||
|
||||
const int64_t batch = is_var_seqlen ? query_start_loc.value().size(0) - 1 : x.size(0);
|
||||
const int64_t dim = x.size(-2);
|
||||
const int64_t seqlen = x.size(-1);
|
||||
const int64_t width = weight.size(-1);
|
||||
|
||||
const auto scalar_type = x.scalar_type();
|
||||
CHECK_EQ(weight.scalar_type(), scalar_type);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(query_start_loc, batch + 1, at::kInt);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(has_initial_state, batch, at::kBool);
|
||||
|
||||
if (conv_states.has_value()) {
|
||||
auto& conv_states_val = conv_states.value();
|
||||
int64_t padded_batch = conv_states_val.size(0);
|
||||
CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
|
||||
CHECK_GE(padded_batch, batch);
|
||||
CHECK_EQ(conv_states_val.size(1), dim);
|
||||
CHECK_EQ(conv_states_val.size(2), width - 1);
|
||||
|
||||
// adjust `conv_states` to be contiguous on `dim`
|
||||
// should happen only once
|
||||
if (conv_states_val.stride(-2) != 1) {
|
||||
auto conv_states_copy = conv_states_val.clone();
|
||||
conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
|
||||
conv_states_val.copy_(conv_states_copy);
|
||||
}
|
||||
}
|
||||
|
||||
// IMPORTANT: To make the kernal compatible with vLLM KV cache layout
|
||||
int64_t conv_state_slot_stride = conv_states->stride(0);
|
||||
|
||||
// block size for sequence blocks, 32
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
|
||||
// total number of sequence blocks
|
||||
int64_t num_seq_blocks = get_block_count<BLOCK_M>(query_start_loc, batch, seqlen);
|
||||
|
||||
at::Tensor out = at::empty_like(x);
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_fwd_kernel_impl", [&] {
|
||||
if (is_var_seqlen) {
|
||||
// record seq blocks in Coordinate format, aka [num_seq_blocks, 2]
|
||||
at::Tensor block_indices = get_block_indices<BLOCK_M>(query_start_loc, num_seq_blocks);
|
||||
|
||||
causal_conv1d_fwd_varlen_kernel_impl(
|
||||
out.data_ptr<scalar_t>(),
|
||||
x.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<scalar_t>(),
|
||||
conditional_data_ptr<scalar_t>(bias),
|
||||
conditional_data_ptr<scalar_t>(conv_states),
|
||||
conditional_data_ptr<int32_t>(query_start_loc),
|
||||
conditional_data_ptr<int32_t>(conv_state_indices),
|
||||
conditional_data_ptr<bool>(has_initial_state),
|
||||
block_indices.data_ptr<int32_t>(),
|
||||
silu_activation,
|
||||
batch,
|
||||
dim,
|
||||
width,
|
||||
num_seq_blocks,
|
||||
conv_state_slot_stride);
|
||||
} else {
|
||||
causal_conv1d_fwd_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
x.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<scalar_t>(),
|
||||
conditional_data_ptr<scalar_t>(bias),
|
||||
conditional_data_ptr<scalar_t>(conv_states),
|
||||
conditional_data_ptr<int32_t>(conv_state_indices),
|
||||
conditional_data_ptr<bool>(has_initial_state),
|
||||
silu_activation,
|
||||
batch,
|
||||
dim,
|
||||
seqlen,
|
||||
width,
|
||||
num_seq_blocks,
|
||||
conv_state_slot_stride);
|
||||
}
|
||||
});
|
||||
return out;
|
||||
}
|
||||
|
||||
// API aligned with GPUs
|
||||
//
|
||||
// x: (batch, dim) or (batch, dim, seqlen)
|
||||
// conv_state: (..., dim, state_len), where state_len >= width - 1
|
||||
// weight: (dim, width)
|
||||
// bias: (dim,)
|
||||
// cache_seqlens: (batch,), dtype int32.
|
||||
// conv_state_indices: (batch,), dtype int32
|
||||
// pad_slot_id: int
|
||||
// out: (batch, dim) or (batch, dim, seqlen)
|
||||
//
|
||||
at::Tensor causal_conv1d_update_cpu(
|
||||
const at::Tensor& x,
|
||||
const at::Tensor& conv_states,
|
||||
const at::Tensor& weight,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
bool silu_activation,
|
||||
const std::optional<at::Tensor>& cache_seqlens,
|
||||
const std::optional<at::Tensor>& conv_state_indices,
|
||||
int64_t pad_slot_id,
|
||||
bool is_vnni) {
|
||||
CHECK_CONTIGUOUS(x);
|
||||
CHECK_CONTIGUOUS(weight);
|
||||
auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
|
||||
|
||||
// TODO: add multi-token prediction support
|
||||
TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
|
||||
TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
|
||||
|
||||
int64_t batch = x.size(0);
|
||||
int64_t dim = x.size(1);
|
||||
int64_t seqlen = 1;
|
||||
int64_t width = weight.size(-1);
|
||||
|
||||
const auto scalar_type = x.scalar_type();
|
||||
CHECK_EQ(weight.scalar_type(), scalar_type);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
|
||||
CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
|
||||
|
||||
CHECK_EQ(conv_states.scalar_type(), scalar_type);
|
||||
CHECK_EQ(conv_states.size(1), dim);
|
||||
CHECK_EQ(conv_states.size(2), width - 1);
|
||||
|
||||
// adjust `conv_states` to be contiguous on `dim`
|
||||
if (conv_states.stride(-2) != 1) {
|
||||
int64_t num_cache_lines = conv_states.size(0);
|
||||
auto conv_states_copy = conv_states.clone();
|
||||
conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
|
||||
conv_states.copy_(conv_states_copy);
|
||||
}
|
||||
|
||||
// IMPORTANT: To make the kernal compatible with vLLM KV cache layout
|
||||
int64_t conv_state_slot_stride = conv_states.stride(0);
|
||||
at::Tensor out = at::empty_like(x);
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
|
||||
causal_conv1d_update_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
x.data_ptr<scalar_t>(),
|
||||
conv_states.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<scalar_t>(),
|
||||
conditional_data_ptr<scalar_t>(bias),
|
||||
conditional_data_ptr<int32_t>(conv_state_indices),
|
||||
silu_activation,
|
||||
batch,
|
||||
dim,
|
||||
seqlen,
|
||||
width,
|
||||
conv_state_slot_stride);
|
||||
});
|
||||
return out;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,853 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "gemm.h"
|
||||
|
||||
#include "common.h"
|
||||
#include "vec.h"
|
||||
|
||||
namespace {
|
||||
|
||||
// packed layout:
|
||||
// quants {N, K} int8_t
|
||||
// comp {N} int32_t
|
||||
template <int BLOCK_N>
|
||||
inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
__m512i vcomp[COLS];
|
||||
|
||||
for (int col = 0; col < COLS; ++col) {
|
||||
vcomp[col] = _mm512_setzero_si512();
|
||||
}
|
||||
|
||||
const int64_t offset = BLOCK_N * K;
|
||||
const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
|
||||
for (int k = 0; k < K / 4; ++k) {
|
||||
for (int col = 0; col < COLS; ++col) {
|
||||
__m512i vb = _mm512_loadu_si512((const __m512i*)(packed + k * BLOCK_N * 4 + col * 64));
|
||||
vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
|
||||
}
|
||||
}
|
||||
|
||||
for (int col = 0; col < COLS; ++col) {
|
||||
_mm512_storeu_si512((__m512i*)(packed + offset + col * 64), vcomp[col]);
|
||||
}
|
||||
#else
|
||||
TORCH_CHECK(false, "s8s8_compensation not implemented!");
|
||||
#endif
|
||||
}
|
||||
|
||||
// convert to vnni format
|
||||
// from [N, K] to [K/2, N, 2] for bfloat16 and float16
|
||||
template <typename packed_t>
|
||||
inline void pack_vnni(packed_t* __restrict__ packed, const packed_t* __restrict__ weight, int N, int K) {
|
||||
const int VNNI_BLK = 2;
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k = 0; k < K / VNNI_BLK; ++k) {
|
||||
for (int d = 0; d < VNNI_BLK; ++d) {
|
||||
packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <>
|
||||
inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restrict__ weight, int N, int K) {
|
||||
constexpr int BLOCK_N = block_size_n();
|
||||
TORCH_CHECK(N == BLOCK_N);
|
||||
|
||||
const int VNNI_BLK = 4;
|
||||
for (int n = 0; n < N; ++n) {
|
||||
for (int k = 0; k < K / VNNI_BLK; ++k) {
|
||||
for (int d = 0; d < VNNI_BLK; ++d) {
|
||||
packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
|
||||
}
|
||||
}
|
||||
}
|
||||
s8s8_compensation<BLOCK_N>(packed, K);
|
||||
}
|
||||
|
||||
// uint8_t: mxfp4 or int4
|
||||
// pack to vnni2 format as they are computed with bfloat16
|
||||
//
|
||||
// from [N, K'/2, 2] to [K'/2, N, 2], view 2x int4 as unit8:
|
||||
// from [N, K ] to [K, N ] where K = K'/2
|
||||
//
|
||||
template <>
|
||||
inline void pack_vnni<uint8_t>(uint8_t* __restrict__ packed, const uint8_t* __restrict__ weight, int N, int K) {
|
||||
constexpr int BLOCK_N = block_size_n();
|
||||
|
||||
uint8_t unpacked[2 * BLOCK_N];
|
||||
|
||||
// 32-way pack (align with BLOCK_N), faster for avx512 unpacking
|
||||
//
|
||||
// for a range of (64):
|
||||
// {0, 1, 2, ..., 63}
|
||||
//
|
||||
// original format:
|
||||
// { 1|0, 3|2, ..., 63|62}
|
||||
//
|
||||
// packed format:
|
||||
// {32|0, 31|1, ..., 63|31}
|
||||
//
|
||||
for (int k = 0; k < K; ++k) {
|
||||
// unpack first
|
||||
for (int n = 0; n < N; ++n) {
|
||||
uint8_t value = weight[n * K + k];
|
||||
unpacked[n * 2 + 0] = value & 0xF; // lower 4 bits
|
||||
unpacked[n * 2 + 1] = value >> 4; // higher 4 bits
|
||||
}
|
||||
// re-pack to 32-way
|
||||
for (int n = 0; n < N; ++n) {
|
||||
packed[k * N + n] = (unpacked[n + BLOCK_N] << 4) | unpacked[n];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
fVec data0 = fVec::loadu(input + d);
|
||||
fVec data1 = fVec::loadu(input + d + fVec::size());
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
fVec data0, data1;
|
||||
bVec b_vec = bVec::loadu(input + d);
|
||||
std::tie(data0, data1) = at::vec::convert_to_float(b_vec);
|
||||
data0.store(out + d);
|
||||
data1.store(out + d + fVec::size());
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<float>(input[d]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_add_stub(
|
||||
scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
|
||||
fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d] + bias[d]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool has_bias>
|
||||
inline void scalar_sigmoid_and_mul(
|
||||
scalar_t* __restrict__ out,
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ bias,
|
||||
const scalar_t* __restrict__ mul,
|
||||
int SIZE) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
// scalar sigmoid
|
||||
const fVec one = fVec(1.f);
|
||||
fVec X;
|
||||
if constexpr (has_bias) {
|
||||
assert(bias != nullptr);
|
||||
X = fVec(input[0] + bias[0]);
|
||||
} else {
|
||||
X = fVec(input[0]);
|
||||
}
|
||||
X = one / (one + X.neg().exp_u20());
|
||||
|
||||
// vec mul
|
||||
constexpr int kVecSize = bVec::size();
|
||||
for (int d = 0; d < SIZE; d += kVecSize) {
|
||||
bVec m_bvec = bVec::loadu(mul + d);
|
||||
fVec m_fvec0, m_fvec1;
|
||||
std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec);
|
||||
m_fvec0 = m_fvec0 * X;
|
||||
m_fvec1 = m_fvec1 * X;
|
||||
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0, m_fvec1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_nn {
|
||||
static inline void apply(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
const float* __restrict__ bias,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <bool has_bias, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
|
||||
static inline void apply(
|
||||
const at::BFloat16* __restrict__ A,
|
||||
const at::BFloat16* __restrict__ B,
|
||||
at::BFloat16* __restrict__ C,
|
||||
const float* __restrict__ bias,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int ROWS = BLOCK_M;
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
|
||||
// prefetch distance
|
||||
constexpr int PREFETCH_SIZE_K = 0;
|
||||
|
||||
__m512bh va;
|
||||
__m512bh vb[COLS];
|
||||
__m512 vc[ROWS * COLS];
|
||||
|
||||
auto loadc = [&](auto i) {
|
||||
constexpr int col = i % COLS;
|
||||
if constexpr (has_bias) {
|
||||
vc[i] = _mm512_loadu_ps(bias + col * 16);
|
||||
} else {
|
||||
vc[i] = _mm512_set1_ps(0.f);
|
||||
}
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(loadc);
|
||||
|
||||
const int64_t K2 = K >> 1;
|
||||
const int64_t lda2 = lda >> 1;
|
||||
const int64_t ldb2 = ldb; // ldb * 2 >> 1;
|
||||
const float* a_ptr = reinterpret_cast<const float*>(A);
|
||||
const float* b_ptr = reinterpret_cast<const float*>(B);
|
||||
|
||||
auto compute = [&](auto i, int64_t k) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
if constexpr (col == 0) {
|
||||
va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
|
||||
}
|
||||
if constexpr (row == 0) {
|
||||
vb[col] = (__m512bh)(_mm512_loadu_si512(b_ptr + k * ldb2 + col * 16));
|
||||
if constexpr (PREFETCH_SIZE_K > 0) {
|
||||
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
|
||||
}
|
||||
}
|
||||
vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
|
||||
};
|
||||
for (int64_t k = 0; k < K2; ++k) {
|
||||
Unroll<ROWS * COLS>{}(compute, k);
|
||||
}
|
||||
|
||||
auto storec = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
// for COLS = 2, 4 use 512bit store
|
||||
// for COLS = 1, 3 use 256bit store
|
||||
if constexpr (COLS % 2 == 0) {
|
||||
if constexpr (col % 2 == 0) {
|
||||
_mm512_storeu_si512(
|
||||
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
|
||||
(__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
|
||||
}
|
||||
} else {
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(C + row * ldc + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc[i])));
|
||||
}
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(storec);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
|
||||
tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
|
||||
A + mb_start * lda, \
|
||||
B + nb_start * 2, \
|
||||
C + mb_start * ldc + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
K, \
|
||||
lda, \
|
||||
ldb, \
|
||||
ldc);
|
||||
|
||||
template <typename scalar_t, bool has_bias>
|
||||
struct brgemm {
|
||||
static inline void apply(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int BLOCK_N = block_size_n();
|
||||
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
|
||||
|
||||
// copy from Ctmp to C
|
||||
for (int64_t m = 0; m < M; ++m) {
|
||||
if constexpr (has_bias) {
|
||||
copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
|
||||
} else {
|
||||
copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
|
||||
}
|
||||
}
|
||||
}
|
||||
static inline void apply(
|
||||
const float* __restrict__ A,
|
||||
const float* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int BLOCK_N = block_size_n();
|
||||
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename scalar_t, bool has_bias>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg) {
|
||||
if (brg) {
|
||||
brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
|
||||
return;
|
||||
}
|
||||
|
||||
// pattern: 1-4-16, N = 16, 32, 48, 64
|
||||
constexpr int64_t BLOCK_M = 4;
|
||||
constexpr int64_t BLOCK_N = 64;
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
for (int mb = 0; mb < MB; ++mb) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
|
||||
for (int64_t nb = 0; nb < NB; ++nb) {
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
|
||||
|
||||
switch (mb_size << 4 | nb_size >> 4) {
|
||||
// mb_size = 1
|
||||
case 0x11:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
|
||||
break;
|
||||
case 0x12:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
|
||||
break;
|
||||
case 0x13:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
|
||||
break;
|
||||
case 0x14:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
|
||||
break;
|
||||
// mb_size = 2
|
||||
case 0x21:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 16);
|
||||
break;
|
||||
case 0x22:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
|
||||
break;
|
||||
case 0x23:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 48);
|
||||
break;
|
||||
case 0x24:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
|
||||
break;
|
||||
// mb_size = 3
|
||||
case 0x31:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 16);
|
||||
break;
|
||||
case 0x32:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
|
||||
break;
|
||||
case 0x33:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 48);
|
||||
break;
|
||||
case 0x34:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
|
||||
break;
|
||||
// mb_size = 4
|
||||
case 0x41:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 16);
|
||||
break;
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
|
||||
break;
|
||||
case 0x43:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 48);
|
||||
break;
|
||||
case 0x44:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", mb_size, " x ", nb_size);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool has_bias>
|
||||
void tinygemm_kernel(
|
||||
const float* __restrict__ A,
|
||||
const float* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg) {
|
||||
TORCH_CHECK(brg, "Expected to use fp32 brgemm for small N GEMM");
|
||||
if (brg) {
|
||||
brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
|
||||
return;
|
||||
}
|
||||
// TODO : add intrinsic path
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void weight_packed_linear_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ mat1,
|
||||
const scalar_t* __restrict__ mat2,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t mat1_strideM,
|
||||
int64_t out_strideM) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
const bool use_brgemm = can_use_brgemm<scalar_t>(M);
|
||||
|
||||
// parallel on [MB, NB]
|
||||
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// for brgemm, use float32 for accumulate
|
||||
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
|
||||
|
||||
loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(N - nb_start, BLOCK_N);
|
||||
|
||||
tinygemm_kernel<scalar_t, has_bias>(
|
||||
/* A */ mat1 + mb_start * mat1_strideM,
|
||||
/* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
|
||||
/* C */ out + mb_start * out_strideM + nb_start,
|
||||
/* Ctmp*/ Ctmp,
|
||||
/* bias*/ bias + nb_start,
|
||||
/* M */ mb_size,
|
||||
/* N */ nb_size,
|
||||
/* K */ K,
|
||||
/* lda */ mat1_strideM,
|
||||
/* ldb */ nb_size,
|
||||
/* ldc */ out_strideM,
|
||||
/* brg */ use_brgemm);
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void weight_packed_linear_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ mat1,
|
||||
const float* __restrict__ mat2,
|
||||
const float* __restrict__ bias,
|
||||
const scalar_t* __restrict__ post_mul_mat,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t mat1_strideM,
|
||||
int64_t out_strideM) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
const bool use_brgemm = true; // TODO: add intrinsic path
|
||||
// parallel on [MB, NB]
|
||||
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// for brgemm, use float32 for accumulate
|
||||
alignas(64) float Atmp[BLOCK_M * K];
|
||||
alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
|
||||
|
||||
loop_2d<float>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(M - mb_start, BLOCK_M);
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(N - nb_start, BLOCK_N);
|
||||
for (int64_t m = 0; m < mb_size; ++m) {
|
||||
copy_stub<scalar_t>(Atmp + m * K, mat1 + mb_start * mat1_strideM + m * K, K);
|
||||
}
|
||||
tinygemm_kernel<scalar_t, has_bias>(
|
||||
/* A */ Atmp,
|
||||
/* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
|
||||
/* C */ out + mb_start * out_strideM + nb_start,
|
||||
/* Ctmp*/ Ctmp,
|
||||
/* bias*/ bias + nb_start,
|
||||
/* M */ mb_size,
|
||||
/* N */ nb_size,
|
||||
/* K */ K,
|
||||
/* lda */ mat1_strideM,
|
||||
/* ldb */ nb_size,
|
||||
/* ldc */ out_strideM,
|
||||
/* brg */ use_brgemm);
|
||||
|
||||
if (post_mul_mat != nullptr) {
|
||||
for (int64_t m = 0; m < mb_size; ++m) {
|
||||
scalar_sigmoid_and_mul<scalar_t, has_bias>(
|
||||
out + mb_start * out_strideM + nb_start + m * out_strideM,
|
||||
Ctmp + m * BLOCK_N,
|
||||
bias + nb_start,
|
||||
post_mul_mat + mb_start * out_strideM + m * out_strideM,
|
||||
out_strideM);
|
||||
}
|
||||
} else {
|
||||
for (int64_t m = 0; m < mb_size; ++m) {
|
||||
if constexpr (has_bias) {
|
||||
copy_add_stub(
|
||||
out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, bias + nb_start, N);
|
||||
} else {
|
||||
copy_stub(out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, N);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// tinygemm interface
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg) {
|
||||
tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
|
||||
}
|
||||
|
||||
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
|
||||
template void tinygemm_kernel<TYPE>( \
|
||||
const TYPE* __restrict__ A, \
|
||||
const TYPE* __restrict__ B, \
|
||||
TYPE* __restrict__ C, \
|
||||
float* __restrict__ Ctmp, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K, \
|
||||
int64_t lda, \
|
||||
int64_t ldb, \
|
||||
int64_t ldc, \
|
||||
bool brg)
|
||||
|
||||
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
|
||||
|
||||
at::Tensor convert_weight_packed(at::Tensor& weight) {
|
||||
// for 3d moe weights
|
||||
// weight : [E, OC, IC]
|
||||
// w1 : [E, 2N, K]
|
||||
// w2 : [E, K, N]
|
||||
CHECK_INPUT(weight);
|
||||
|
||||
const int64_t ndim = weight.ndimension();
|
||||
TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
|
||||
|
||||
if (ndim == 2 && weight.size(0) < TILE_N) {
|
||||
// for 2D weight and small OC shape, we use fma linear path, which needs transpose not pack
|
||||
return weight.to(at::kFloat).t().contiguous();
|
||||
}
|
||||
|
||||
const auto st = weight.scalar_type();
|
||||
const int64_t E = ndim == 3 ? weight.size(0) : 1;
|
||||
const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
|
||||
const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);
|
||||
|
||||
// mxfp4 or int4 are packed with uint8
|
||||
const int64_t actual_IC = st == at::kByte ? IC * 2 : IC;
|
||||
|
||||
// we handle 2 TILE_N at a time.
|
||||
TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
|
||||
TORCH_CHECK(actual_IC % TILE_K == 0, "invalid weight input features ", actual_IC);
|
||||
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
const int64_t NB = div_up(OC, BLOCK_N);
|
||||
|
||||
// use phony sizes here [E, OC, IC], for each [E], [OC, IC] -> [IC / 2, OC, 2]
|
||||
auto packed_weight = at::empty({}, weight.options());
|
||||
const int64_t stride = OC * IC;
|
||||
|
||||
// Note: for `kByte` (uint8), it represents either `mxfp4` or `int4`.
|
||||
TORCH_CHECK(
|
||||
st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn || st == at::kByte,
|
||||
"expect weight to be bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4).");
|
||||
|
||||
CPU_DISPATCH_PACKED_TYPES(st, [&] {
|
||||
// adjust most inner dimension size
|
||||
const int packed_row_size = get_row_size<packed_t>(actual_IC);
|
||||
auto sizes = weight.sizes().vec();
|
||||
sizes[ndim - 1] = packed_row_size;
|
||||
packed_weight.resize_(sizes);
|
||||
|
||||
const packed_t* w_data = weight.data_ptr<packed_t>();
|
||||
packed_t* packed_data = packed_weight.data_ptr<packed_t>();
|
||||
|
||||
// parallel on {E, NB}
|
||||
at::parallel_for(0, E * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t e{0}, nb{0};
|
||||
data_index_init(begin, e, E, nb, NB);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
UNUSED(i);
|
||||
|
||||
int64_t n = nb * BLOCK_N;
|
||||
int64_t n_size = std::min(BLOCK_N, OC - n);
|
||||
pack_vnni<packed_t>(
|
||||
packed_data + e * OC * packed_row_size + n * packed_row_size, w_data + e * stride + n * IC, n_size, IC);
|
||||
|
||||
// move to the next index
|
||||
data_index_step(e, E, nb, NB);
|
||||
}
|
||||
});
|
||||
});
|
||||
return packed_weight;
|
||||
}
|
||||
|
||||
at::Tensor convert_scale_packed(at::Tensor& scale) {
|
||||
CHECK_INPUT(scale);
|
||||
|
||||
const int64_t ndim = scale.ndimension();
|
||||
TORCH_CHECK(ndim == 2 || ndim == 3, "expect scale to be 2d or 3d, got ", ndim, "d tensor.");
|
||||
const auto st = scale.scalar_type();
|
||||
const int64_t E = ndim == 3 ? scale.size(0) : 1;
|
||||
const int64_t N = ndim == 3 ? scale.size(1) : scale.size(0);
|
||||
// number of groups, e.g. K/32
|
||||
const int64_t G = ndim == 3 ? scale.size(2) : scale.size(1);
|
||||
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
TORCH_CHECK(N % BLOCK_N == 0, "invalid weight out features ", N);
|
||||
const int64_t NB = N / BLOCK_N;
|
||||
|
||||
auto packed_scale = at::empty_like(scale);
|
||||
TORCH_CHECK(st == at::kByte, "expect scale to be uint8.");
|
||||
|
||||
const uint8_t* s_data = scale.data_ptr<uint8_t>();
|
||||
uint8_t* packed_data = packed_scale.data_ptr<uint8_t>();
|
||||
|
||||
// parallel on src {E, NB, BLOCK_N, G}, dst {E, NB, G, BLOCK_N}
|
||||
at::parallel_for(0, E * NB * BLOCK_N * G, 0, [&](int64_t begin, int64_t end) {
|
||||
int64_t e{0}, nb{0}, n{0}, g{0};
|
||||
data_index_init(begin, e, E, nb, NB, n, BLOCK_N, g, G);
|
||||
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
packed_data[e * N * G + nb * G * BLOCK_N + g * BLOCK_N + n] = s_data[i];
|
||||
// move to the next index
|
||||
data_index_step(e, E, nb, NB, n, BLOCK_N, g, G);
|
||||
}
|
||||
});
|
||||
return packed_scale;
|
||||
}
|
||||
|
||||
// mat1 : [M, K]
|
||||
// mat2 : [N, K] ([K, N] if use_fma_gemm)
|
||||
// bias : [N]
|
||||
// out : [M, N]
|
||||
//
|
||||
at::Tensor
|
||||
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
|
||||
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
|
||||
bool use_fma_gemm = false;
|
||||
if (packed_w.scalar_type() == at::kFloat) {
|
||||
use_fma_gemm = true;
|
||||
}
|
||||
|
||||
int64_t M = mat1.size(0);
|
||||
int64_t K = mat1.size(1);
|
||||
int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
|
||||
|
||||
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
|
||||
CHECK_INPUT(mat2);
|
||||
CHECK_DIM(2, mat1);
|
||||
CHECK_DIM(2, mat2);
|
||||
if (!use_fma_gemm) {
|
||||
CHECK_EQ(mat1.size(1), K);
|
||||
}
|
||||
|
||||
auto dispatch_type = mat1.scalar_type();
|
||||
auto out = at::empty({M, N}, mat1.options());
|
||||
// strides
|
||||
int64_t out_strideM = out.stride(0);
|
||||
int64_t mat1_strideM = mat1.stride(0);
|
||||
|
||||
const bool has_bias = bias.has_value();
|
||||
const float* bias_data = nullptr;
|
||||
if (has_bias) {
|
||||
CHECK_EQ(bias.value().size(0), N);
|
||||
bias_data = bias.value().data_ptr<float>();
|
||||
}
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] {
|
||||
if (use_fma_gemm) {
|
||||
weight_packed_linear_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
mat1.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<float>(),
|
||||
bias_data,
|
||||
nullptr,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
mat1_strideM,
|
||||
out_strideM);
|
||||
} else {
|
||||
weight_packed_linear_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
mat1.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<scalar_t>(),
|
||||
bias_data,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
mat1_strideM,
|
||||
out_strideM);
|
||||
}
|
||||
});
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
// mat1 : [M, K]
|
||||
// mat2 : [K, 1]
|
||||
// post_mul_mat : [M, K]
|
||||
// bias : [N]
|
||||
// out : [M, N]
|
||||
//
|
||||
at::Tensor fused_linear_sigmoid_mul(
|
||||
at::Tensor& mat1,
|
||||
at::Tensor& mat2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
bool is_vnni,
|
||||
const at::Tensor& post_mul_mat) {
|
||||
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
|
||||
TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight")
|
||||
|
||||
int64_t M = mat1.size(0);
|
||||
int64_t K = mat1.size(1);
|
||||
int64_t N = mat2.size(1);
|
||||
|
||||
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
|
||||
CHECK_INPUT(mat2);
|
||||
CHECK_DIM(2, mat1);
|
||||
CHECK_DIM(2, mat2);
|
||||
|
||||
int64_t out_strideM = post_mul_mat.size(1);
|
||||
int64_t mat1_strideM = mat1.stride(0);
|
||||
auto dispatch_type = mat1.scalar_type();
|
||||
auto out = at::empty({M, out_strideM}, mat1.options());
|
||||
|
||||
TORCH_CHECK(
|
||||
N == 1 && out_strideM % 32 == 0,
|
||||
"post_mul_mat tensor size(1) should be 32 dividable, and the mat2 OC=1 (Mx1 as linear output shape)")
|
||||
|
||||
const bool has_bias = bias.has_value();
|
||||
const float* bias_data = nullptr;
|
||||
if (has_bias) {
|
||||
CHECK_EQ(bias.value().size(0), N);
|
||||
bias_data = bias.value().data_ptr<float>();
|
||||
}
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "fused_linear_sigmoid_mul", [&] {
|
||||
weight_packed_linear_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
mat1.data_ptr<scalar_t>(),
|
||||
packed_w.data_ptr<float>(),
|
||||
bias_data,
|
||||
post_mul_mat.data_ptr<scalar_t>(),
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
mat1_strideM,
|
||||
out_strideM);
|
||||
});
|
||||
|
||||
return out;
|
||||
}
|
||||
@@ -0,0 +1,385 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
#include "common.h"
|
||||
#include "blas_gemm.h"
|
||||
|
||||
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
|
||||
#define CPU_CAPABILITY_AVX512
|
||||
#endif
|
||||
|
||||
// amx-bf16
|
||||
#define TILE_M 16
|
||||
#define TILE_N 16
|
||||
#define TILE_K 32
|
||||
|
||||
// block size for AMX gemm
|
||||
constexpr int block_size_m() {
|
||||
return 2 * TILE_M;
|
||||
}
|
||||
constexpr int block_size_n() {
|
||||
return 2 * TILE_N;
|
||||
}
|
||||
|
||||
constexpr bool brgemm_supported() {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
#endif
|
||||
}
|
||||
|
||||
// define threshold using brgemm (intel AMX)
|
||||
template <typename T>
|
||||
inline bool can_use_brgemm(int M);
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::BFloat16>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Half>(int M) {
|
||||
return brgemm_supported();
|
||||
}
|
||||
// this requires PyTorch 2.7 or above
|
||||
template <>
|
||||
inline bool can_use_brgemm<int8_t>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<uint8_t>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
|
||||
return brgemm_supported() && M > 4;
|
||||
}
|
||||
|
||||
// work around compiler internal error
|
||||
#define BLOCK_K 128 // 4 * TILE_K
|
||||
|
||||
// adjust leading dimension size for K
|
||||
template <typename T>
|
||||
inline int64_t get_row_size(int64_t K) {
|
||||
return K;
|
||||
}
|
||||
|
||||
template <>
|
||||
inline int64_t get_row_size<int8_t>(int64_t K) {
|
||||
return K + sizeof(int32_t);
|
||||
}
|
||||
|
||||
// uint8: mxfp4 or int4
|
||||
template <>
|
||||
inline int64_t get_row_size<uint8_t>(int64_t K) {
|
||||
return K >> 1;
|
||||
}
|
||||
|
||||
inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
|
||||
return use_int8_w8a8 ? K + sizeof(int32_t) : K;
|
||||
}
|
||||
|
||||
enum class CPUAcTMethod : int { silu_and_mul = 0, swiglu = 1 };
|
||||
|
||||
constexpr bool operator==(CPUAcTMethod a, int b) {
|
||||
return static_cast<int>(a) == b;
|
||||
}
|
||||
|
||||
constexpr bool operator==(int a, CPUAcTMethod b) {
|
||||
return a == static_cast<int>(b);
|
||||
}
|
||||
|
||||
enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4 };
|
||||
|
||||
constexpr bool operator==(CPUQuantMethod a, int64_t b) {
|
||||
return static_cast<int64_t>(a) == b;
|
||||
}
|
||||
|
||||
constexpr bool operator==(int64_t a, CPUQuantMethod b) {
|
||||
return a == static_cast<int64_t>(b);
|
||||
}
|
||||
|
||||
enum class CPUQuantAlgo : int64_t { AWQ = 0, GPTQ = 1 };
|
||||
|
||||
constexpr bool operator==(CPUQuantAlgo a, int64_t b) {
|
||||
return static_cast<int64_t>(a) == b;
|
||||
}
|
||||
|
||||
constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
|
||||
return a == static_cast<int64_t>(b);
|
||||
}
|
||||
|
||||
inline int64_t get_4bit_block_k_size(int64_t group_size) {
|
||||
return group_size > 128 ? 128 : group_size;
|
||||
}
|
||||
|
||||
// pack weight to vnni format
|
||||
at::Tensor convert_weight_packed(at::Tensor& weight);
|
||||
|
||||
// pack scale to blocked format for mxfp4
|
||||
at::Tensor convert_scale_packed(at::Tensor& scale);
|
||||
|
||||
// pack weight to vnni format for int4
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor>
|
||||
convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
|
||||
|
||||
// moe implementations for int8 w8a8
|
||||
template <typename scalar_t>
|
||||
void fused_experts_int8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
uint8_t* __restrict__ A_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const int8_t* __restrict__ packed_w1,
|
||||
const int8_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad);
|
||||
|
||||
// moe implementations for fp8 w8a16 and mxfp4
|
||||
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
|
||||
void fused_experts_fp_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
scalar_t* __restrict__ A_tmp,
|
||||
scalar_t* __restrict__ B_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const packed_t* __restrict__ packed_w1,
|
||||
const packed_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1_bias,
|
||||
const float* __restrict__ w2_bias,
|
||||
const param_t* __restrict__ w1s,
|
||||
const param_t* __restrict__ w2s,
|
||||
int64_t block_size_N,
|
||||
int64_t block_size_K,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad,
|
||||
float alpha,
|
||||
float limit,
|
||||
CPUAcTMethod act_func,
|
||||
bool with_bias);
|
||||
|
||||
// shared expert implementation for int8 w8a8
|
||||
template <typename scalar_t>
|
||||
void shared_expert_int8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic1,
|
||||
float* __restrict__ C_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const int8_t* __restrict__ packed_w1,
|
||||
const int8_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
const scalar_t* __restrict__ fused_experts_out,
|
||||
float routed_scaling_factor,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K);
|
||||
|
||||
template <typename scalar_t>
|
||||
void fused_experts_int4_w4a8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
uint8_t* __restrict__ A_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
int32_t* __restrict__ Azp_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
int8_t* __restrict__ dqB_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const uint8_t* __restrict__ packed_w1,
|
||||
const uint8_t* __restrict__ packed_w2,
|
||||
const int8_t* __restrict__ w1z,
|
||||
const int8_t* __restrict__ w2z,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
int group_size,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad);
|
||||
|
||||
template <typename scalar_t>
|
||||
void shared_expert_fp8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ B_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w1,
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
int64_t block_size_N,
|
||||
int64_t block_size_K,
|
||||
const scalar_t* __restrict__ fused_experts_out,
|
||||
float routed_scaling_factor,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K);
|
||||
|
||||
// tinygemm interface
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const scalar_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
float* __restrict__ Ctmp,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg);
|
||||
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
int32_t* __restrict__ Ctmp,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg);
|
||||
|
||||
// block quantization
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const at::Float8_e4m3fn* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
scalar_t* __restrict__ Btmp,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ Bbias,
|
||||
const float* __restrict__ scale,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg,
|
||||
int64_t block_size_K,
|
||||
bool do_unpack = true);
|
||||
|
||||
// per tensor quantization
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const at::Float8_e4m3fn* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
scalar_t* __restrict__ Btmp,
|
||||
float* __restrict__ Ctmp,
|
||||
float scale,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg);
|
||||
|
||||
// mxfp4
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const uint8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
scalar_t* __restrict__ Btmp,
|
||||
float* __restrict__ Ctmp,
|
||||
const float* __restrict__ Bbias,
|
||||
const uint8_t* __restrict__ scale,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg,
|
||||
int64_t block_size_K,
|
||||
bool do_unpack = true);
|
||||
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
scalar_t* C,
|
||||
float* C_temp,
|
||||
const uint8_t* A,
|
||||
const float* scales_a,
|
||||
const int32_t* qzeros_a,
|
||||
const uint8_t* B,
|
||||
const float* scales_b,
|
||||
const int8_t* qzeros_b,
|
||||
const int32_t* compensation,
|
||||
int8_t* dqB_tmp,
|
||||
int64_t M,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldc_f,
|
||||
int64_t ldc_s,
|
||||
bool store_out,
|
||||
bool use_brgemm);
|
||||
|
||||
// mxfp4
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const scalar_t* __restrict__ A,
|
||||
const uint8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
scalar_t* __restrict__ Btmp,
|
||||
float* __restrict__ Ctmp,
|
||||
const uint8_t* __restrict__ scale,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg,
|
||||
int64_t block_size_K,
|
||||
bool do_unpack = true);
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,546 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "common.h"
|
||||
#include "gemm.h"
|
||||
#include "vec.h"
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t, bool has_bias, int BLOCK_N>
|
||||
struct scale_C {
|
||||
static inline void apply(
|
||||
scalar_t* __restrict__ C,
|
||||
const int32_t* __restrict__ Ctmp,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
const float* __restrict__ bias,
|
||||
float As,
|
||||
const float* __restrict__ Bs) {
|
||||
TORCH_CHECK(false, "scale_C: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <bool has_bias, int BLOCK_N>
|
||||
struct scale_C<at::BFloat16, has_bias, BLOCK_N> {
|
||||
static inline void apply(
|
||||
at::BFloat16* __restrict__ C,
|
||||
const int32_t* __restrict__ Ctmp,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
const float* __restrict__ bias,
|
||||
float As,
|
||||
const float* __restrict__ Bs) {
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
__m512 vc[COLS];
|
||||
__m512 vd0 = _mm512_set1_ps(As);
|
||||
|
||||
auto compute = [&](auto col) {
|
||||
__m512 vd1 = _mm512_loadu_ps(Bs + col * 16);
|
||||
__m512i vcomp = _mm512_loadu_si512(Bcomp + col * 16);
|
||||
__m512i vc32 = _mm512_loadu_si512(Ctmp + col * 16);
|
||||
vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp));
|
||||
if constexpr (has_bias) {
|
||||
__m512 vbias = _mm512_loadu_ps(bias + col * 16);
|
||||
vc[col] = _mm512_fmadd_ps(_mm512_mul_ps(vc[col], vd0), vd1, vbias);
|
||||
} else {
|
||||
vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vd0), vd1);
|
||||
}
|
||||
};
|
||||
Unroll<COLS>{}(compute);
|
||||
|
||||
auto storec = [&](auto col) {
|
||||
// for COLS = 2, 4 use 512bit store
|
||||
if constexpr (col % 2 == 0) {
|
||||
_mm512_storeu_si512(
|
||||
reinterpret_cast<__m512i*>((C + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc[col + 1], vc[col + 0])));
|
||||
}
|
||||
};
|
||||
Unroll<COLS>{}(storec);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_nn {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <bool has_bias, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
at::BFloat16* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
const float* __restrict__ bias,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int ROWS = BLOCK_M;
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
// prefetch distance
|
||||
constexpr int PREFETCH_SIZE_K = 0;
|
||||
|
||||
__m512i va;
|
||||
__m512i vb[COLS];
|
||||
__m512i vc[ROWS * COLS];
|
||||
__m512i vcomp[COLS];
|
||||
__m512 vd0;
|
||||
__m512 vd1[COLS];
|
||||
|
||||
// oops! 4x4 spills but we use 4x2
|
||||
__m512 vbias[COLS];
|
||||
|
||||
// [NOTE]: s8s8 igemm compensation in avx512-vnni
|
||||
//
|
||||
// avx512-vnni has no s8s8, so we need to change s8s8 to u8s8 with compensate:
|
||||
//
|
||||
// a * b = (a + 128) * b - 128 * b
|
||||
// s s u s u s
|
||||
//
|
||||
// 1) 128 * b is pre-computed when packing B to vnni formats
|
||||
// 2) a + 128 is fused when dynamically quantize A
|
||||
//
|
||||
auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
|
||||
Unroll<ROWS * COLS>{}(loadc);
|
||||
|
||||
const int64_t K4 = K >> 2;
|
||||
const int64_t lda4 = lda >> 2;
|
||||
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
|
||||
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
|
||||
const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
|
||||
|
||||
auto compute = [&](auto i, int64_t k) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
if constexpr (col == 0) {
|
||||
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
|
||||
}
|
||||
if constexpr (row == 0) {
|
||||
vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
|
||||
if constexpr (PREFETCH_SIZE_K > 0) {
|
||||
_mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb4 + col * 16, _MM_HINT_T0);
|
||||
}
|
||||
}
|
||||
vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
|
||||
};
|
||||
for (int64_t k = 0; k < K4; ++k) {
|
||||
Unroll<ROWS * COLS>{}(compute, k);
|
||||
}
|
||||
|
||||
auto storec = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
// load a scale
|
||||
if constexpr (col == 0) {
|
||||
vd0 = _mm512_set1_ps(As[row]);
|
||||
}
|
||||
// load b scale and vcomp per 2 vectors
|
||||
// also load bias if any
|
||||
if constexpr (row == 0) {
|
||||
if constexpr (col % 2 == 0) {
|
||||
vd1[col + 0] = _mm512_loadu_ps(Bs + col * 16);
|
||||
vd1[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
|
||||
vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
|
||||
vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
|
||||
if constexpr (has_bias) {
|
||||
vbias[col + 0] = _mm512_loadu_ps(bias + col * 16);
|
||||
vbias[col + 1] = _mm512_loadu_ps(bias + col * 16 + 16);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// for COLS = 2, 4 use 512bit store
|
||||
if constexpr (col % 2 == 0) {
|
||||
__m512 vc0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 0], vcomp[col + 0]));
|
||||
__m512 vc1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 1], vcomp[col + 1]));
|
||||
if constexpr (has_bias) {
|
||||
vc0 = _mm512_fmadd_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0], vbias[col + 0]);
|
||||
vc1 = _mm512_fmadd_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1], vbias[col + 1]);
|
||||
} else {
|
||||
vc0 = _mm512_mul_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0]);
|
||||
vc1 = _mm512_mul_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1]);
|
||||
}
|
||||
|
||||
_mm512_storeu_si512(
|
||||
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
|
||||
}
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(storec);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
|
||||
tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
|
||||
A + mb_start * lda, \
|
||||
B + nb_start * 4, \
|
||||
C + mb_start * ldc + nb_start, \
|
||||
As + mb_start, \
|
||||
Bs + nb_start, \
|
||||
Bcomp + nb_start, \
|
||||
has_bias ? bias + nb_start : nullptr, \
|
||||
K, \
|
||||
lda, \
|
||||
ldb, \
|
||||
ldc);
|
||||
|
||||
template <typename scalar_t, bool has_bias>
|
||||
void tinygemm_kernel(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
int32_t* __restrict__ Ctmp,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg) {
|
||||
// B compensation
|
||||
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
|
||||
|
||||
if (brg) {
|
||||
constexpr int BLOCK_N = block_size_n();
|
||||
at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
|
||||
|
||||
// apply compensation and scale
|
||||
for (int64_t m = 0; m < M; ++m) {
|
||||
scale_C<scalar_t, has_bias, BLOCK_N>::apply(C + m * ldc, Ctmp + m * BLOCK_N, Bcomp, bias, As[m], Bs);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// pattern: 1-4-16
|
||||
constexpr int64_t BLOCK_M = 4;
|
||||
constexpr int64_t BLOCK_N = 64;
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
for (int64_t mb = 0; mb < MB; ++mb) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
|
||||
for (int64_t nb = 0; nb < NB; ++nb) {
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
|
||||
|
||||
switch (mb_size << 4 | nb_size >> 4) {
|
||||
// mb_size = 1
|
||||
case 0x12:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
|
||||
break;
|
||||
case 0x14:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
|
||||
break;
|
||||
// mb_size = 2
|
||||
case 0x22:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
|
||||
break;
|
||||
case 0x24:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
|
||||
break;
|
||||
// mb_size = 3
|
||||
case 0x32:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
|
||||
break;
|
||||
case 0x34:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
|
||||
break;
|
||||
// mb_size = 4
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
|
||||
break;
|
||||
case 0x44:
|
||||
LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void int8_scaled_mm_kernel_impl(
|
||||
scalar_t* __restrict__ out,
|
||||
const uint8_t* __restrict__ mat1,
|
||||
const int8_t* __restrict__ mat2,
|
||||
const float* __restrict__ scales1,
|
||||
const float* __restrict__ scales2,
|
||||
const float* __restrict__ bias,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
const bool use_brgemm = can_use_brgemm<int8_t>(M);
|
||||
|
||||
// K + 4 after compensation
|
||||
const int64_t packed_row_size = get_row_size<int8_t>(K);
|
||||
|
||||
AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// for brgemm, use int32_t for accumulate
|
||||
alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
|
||||
|
||||
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int mb_start = mb * BLOCK_M;
|
||||
int mb_size = std::min(M - mb_start, BLOCK_M);
|
||||
int nb_start = nb * BLOCK_N;
|
||||
int nb_size = std::min(N - nb_start, BLOCK_N);
|
||||
|
||||
tinygemm_kernel<scalar_t, has_bias>(
|
||||
/* A */ mat1 + mb_start * K,
|
||||
/* B */ mat2 + nb_start * packed_row_size /* nb * BLOCK_N * (K + 4) */,
|
||||
/* C */ out + mb_start * N + nb_start,
|
||||
/* Ctmp*/ Ctmp,
|
||||
/* As */ scales1 + mb_start,
|
||||
/* Bs */ scales2 + nb_start,
|
||||
/* bias*/ bias + nb_start,
|
||||
/* M */ mb_size,
|
||||
/* N */ nb_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ nb_size,
|
||||
/* ldc */ N,
|
||||
/* brg */ use_brgemm);
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// tinygemm interface
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
scalar_t* __restrict__ C,
|
||||
int32_t* __restrict__ Ctmp,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc,
|
||||
bool brg) {
|
||||
tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
|
||||
}
|
||||
|
||||
#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
|
||||
template void tinygemm_kernel<TYPE>( \
|
||||
const uint8_t* __restrict__ A, \
|
||||
const int8_t* __restrict__ B, \
|
||||
TYPE* __restrict__ C, \
|
||||
int32_t* __restrict__ Ctmp, \
|
||||
const float* __restrict__ As, \
|
||||
const float* __restrict__ Bs, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K, \
|
||||
int64_t lda, \
|
||||
int64_t ldb, \
|
||||
int64_t ldc, \
|
||||
bool brg)
|
||||
|
||||
INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
|
||||
|
||||
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
|
||||
CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
|
||||
CHECK_DIM(2, A);
|
||||
|
||||
int64_t M = A.size(0);
|
||||
int64_t K = A.size(1);
|
||||
int64_t lda = A.stride(0);
|
||||
|
||||
const auto st = A.scalar_type();
|
||||
TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "per_token_quant_int8: expect A to be bfloat16 or half.");
|
||||
|
||||
auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
|
||||
auto As = at::empty({M}, A.options().dtype(at::kFloat));
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "per_token_quant_int8", [&] {
|
||||
uint8_t* __restrict__ Aq_data = Aq.data_ptr<uint8_t>();
|
||||
float* __restrict__ As_data = As.data_ptr<float>();
|
||||
const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();
|
||||
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
|
||||
}
|
||||
});
|
||||
});
|
||||
return std::make_tuple(Aq, As);
|
||||
}
|
||||
|
||||
// weight : static, per-channel, symmetric
|
||||
// activation : dynamic, per-token, symmetric
|
||||
//
|
||||
// mat1 : [M, K]
|
||||
// mat2 : [N, K]
|
||||
// scales1 : [M]
|
||||
// scales2 : [N]
|
||||
// bias : [N]
|
||||
// out : [M, N]
|
||||
//
|
||||
at::Tensor int8_scaled_mm_cpu(
|
||||
at::Tensor& mat1,
|
||||
at::Tensor& mat2,
|
||||
at::Tensor& scales1,
|
||||
at::Tensor& scales2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype,
|
||||
bool is_vnni) {
|
||||
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
|
||||
|
||||
CHECK_INPUT(mat1);
|
||||
CHECK_INPUT(mat2);
|
||||
CHECK_INPUT(scales1);
|
||||
CHECK_INPUT(scales2);
|
||||
CHECK_DIM(2, mat1);
|
||||
CHECK_DIM(2, mat2);
|
||||
|
||||
int64_t M = mat1.size(0);
|
||||
int64_t N = mat2.size(0);
|
||||
int64_t K = mat1.size(1);
|
||||
|
||||
// see [NOTE]: s8s8 igemm compensation in avx512-vnni
|
||||
CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
|
||||
CHECK_EQ(scales1.numel(), M);
|
||||
CHECK_EQ(scales2.numel(), N);
|
||||
|
||||
TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
|
||||
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm: expect mat2 to be int8.");
|
||||
TORCH_CHECK(
|
||||
scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
|
||||
"int8_scaled_mm: expect scales to be float32.");
|
||||
|
||||
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
|
||||
|
||||
const bool has_bias = bias.has_value();
|
||||
const float* bias_data = nullptr;
|
||||
if (has_bias) {
|
||||
CHECK_EQ(bias.value().size(0), N);
|
||||
bias_data = bias.value().data_ptr<float>();
|
||||
}
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_kernel_impl", [&] {
|
||||
int8_scaled_mm_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
mat1.data_ptr<uint8_t>(),
|
||||
packed_w.data_ptr<int8_t>(),
|
||||
scales1.data_ptr<float>(),
|
||||
scales2.data_ptr<float>(),
|
||||
bias_data,
|
||||
M,
|
||||
N,
|
||||
K);
|
||||
});
|
||||
return out;
|
||||
}
|
||||
|
||||
// fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
|
||||
at::Tensor int8_scaled_mm_with_quant(
|
||||
at::Tensor& mat1,
|
||||
at::Tensor& mat2,
|
||||
at::Tensor& scales2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype,
|
||||
bool is_vnni) {
|
||||
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
|
||||
|
||||
CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
|
||||
CHECK_INPUT(mat2);
|
||||
CHECK_INPUT(scales2);
|
||||
CHECK_DIM(2, mat1);
|
||||
CHECK_DIM(2, mat2);
|
||||
|
||||
int64_t M = mat1.size(0);
|
||||
int64_t N = mat2.size(0);
|
||||
int64_t K = mat1.size(1);
|
||||
int64_t lda = mat1.stride(0);
|
||||
|
||||
// see [NOTE]: s8s8 igemm compensation in avx512-vnni
|
||||
CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
|
||||
CHECK_EQ(scales2.numel(), N);
|
||||
|
||||
const auto st = mat1.scalar_type();
|
||||
TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
|
||||
TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
|
||||
TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
|
||||
TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");
|
||||
|
||||
const int64_t buffer_size = M * K + M * sizeof(float);
|
||||
auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
|
||||
auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
|
||||
|
||||
const bool has_bias = bias.has_value();
|
||||
const float* bias_data = nullptr;
|
||||
if (has_bias) {
|
||||
CHECK_EQ(bias.value().size(0), N);
|
||||
bias_data = bias.value().data_ptr<float>();
|
||||
}
|
||||
|
||||
AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
|
||||
uint8_t* __restrict__ Aq_data = buffer.data_ptr<uint8_t>();
|
||||
float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
|
||||
const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
|
||||
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
|
||||
}
|
||||
});
|
||||
|
||||
int8_scaled_mm_kernel_impl<scalar_t>(
|
||||
out.data_ptr<scalar_t>(),
|
||||
Aq_data,
|
||||
packed_w.data_ptr<int8_t>(),
|
||||
As_data,
|
||||
scales2.data_ptr<float>(),
|
||||
bias_data,
|
||||
M,
|
||||
N,
|
||||
K);
|
||||
});
|
||||
return out;
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,284 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
#include "vec.h"
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
|
||||
using Vec = at::vec::Vectorized<scalar_t>;
|
||||
const Vec data_vec(val);
|
||||
at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
|
||||
using Vec = at::vec::Vectorized<scalar_t>;
|
||||
constexpr int kVecSize = Vec::size();
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
Vec data = Vec::loadu(input + d);
|
||||
data.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = input[d];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
auto [x0, x1] = load_float_vec2(input + d);
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d]);
|
||||
}
|
||||
}
|
||||
|
||||
template <>
|
||||
inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
|
||||
// size might be 64x + 32
|
||||
std::memcpy(out, input, size * sizeof(uint8_t));
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename input_t>
|
||||
inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) {
|
||||
static_assert(
|
||||
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
|
||||
"copy_mul_stub only supports input_t == float or input_t == scalar_t");
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
const fVec weight_vec = fVec(weight);
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
auto [x0, x1] = load_float_vec2(input + d);
|
||||
x0 = x0 * weight_vec;
|
||||
x1 = x1 * weight_vec;
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d] * weight);
|
||||
}
|
||||
}
|
||||
|
||||
// acc from [topk, K] to [K]
|
||||
template <typename scalar_t>
|
||||
inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
if (topk == 1) {
|
||||
// do copy for topk = 1
|
||||
copy_stub(out, input, K);
|
||||
} else {
|
||||
// do sum for topk != 1
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= K - kVecSize; d += kVecSize) {
|
||||
fVec sum_fvec0 = fVec(0.f);
|
||||
fVec sum_fvec1 = fVec(0.f);
|
||||
for (int t = 0; t < topk; ++t) {
|
||||
bVec x_bvec = bVec::loadu(input + t * K + d);
|
||||
fVec x_fvec0, x_fvec1;
|
||||
std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
|
||||
|
||||
sum_fvec0 += x_fvec0;
|
||||
sum_fvec1 += x_fvec1;
|
||||
}
|
||||
bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
|
||||
out_bvec.store(out + d);
|
||||
}
|
||||
for (; d < K; ++d) {
|
||||
float sum_val = 0.f;
|
||||
for (int t = 0; t < topk; ++t) {
|
||||
sum_val += static_cast<float>(input[t * K + d]);
|
||||
}
|
||||
out[d] = static_cast<scalar_t>(sum_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// out = input + input2 * scale
|
||||
template <typename scalar_t, typename input_t>
|
||||
inline void add_mul_stub(
|
||||
scalar_t* __restrict__ out,
|
||||
const input_t* __restrict__ input,
|
||||
const scalar_t* __restrict__ input2,
|
||||
float scale,
|
||||
int64_t size) {
|
||||
static_assert(
|
||||
std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
|
||||
"add_mul_stub only supports input_t == float or input_t == scalar_t");
|
||||
|
||||
// out = input (without scale factor)
|
||||
if (input2 == nullptr) {
|
||||
copy_stub(out, input, size);
|
||||
return;
|
||||
}
|
||||
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
const fVec s_vec = fVec(scale);
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
auto [x0, x1] = load_float_vec2(input + d);
|
||||
|
||||
bVec y_bvec = bVec::loadu(input2 + d);
|
||||
fVec y0, y1;
|
||||
std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
|
||||
|
||||
x0 = x0 + y0 * s_vec;
|
||||
x1 = x1 + y1 * s_vec;
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void silu_and_mul_stub(
|
||||
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
const fVec one = fVec(1.f);
|
||||
|
||||
// no remainder
|
||||
#pragma GCC unroll 4
|
||||
for (int64_t d = 0; d < size; d += bVec::size()) {
|
||||
bVec x = bVec::loadu(input + d);
|
||||
fVec x0, x1;
|
||||
std::tie(x0, x1) = at::vec::convert_to_float(x);
|
||||
bVec y = bVec::loadu(input2 + d);
|
||||
fVec y0, y1;
|
||||
std::tie(y0, y1) = at::vec::convert_to_float(y);
|
||||
x0 = x0 / (one + x0.neg().exp_u20());
|
||||
x1 = x1 / (one + x1.neg().exp_u20());
|
||||
x0 = x0 * y0;
|
||||
x1 = x1 * y1;
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
const fVec weight_vec = fVec(weight);
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
fVec data0 = fVec::loadu(input + d) * weight_vec;
|
||||
fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d] * weight);
|
||||
}
|
||||
}
|
||||
|
||||
// input = input + input2
|
||||
inline void add_bias_stub(float* __restrict__ input, const float* __restrict__ input2, int64_t size) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = fVec::size();
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
fVec x_fvec = fVec::loadu(input + d);
|
||||
fVec y_fvec = fVec::loadu(input2 + d);
|
||||
x_fvec = x_fvec + y_fvec;
|
||||
x_fvec.store(input + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
input[d] = input[d] + input2[d];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = bVec::size();
|
||||
const fVec weight_vec = fVec(weight);
|
||||
int64_t d;
|
||||
#pragma GCC unroll 4
|
||||
for (d = 0; d <= size - kVecSize; d += kVecSize) {
|
||||
bVec x = bVec::loadu(input + d);
|
||||
fVec x0, x1;
|
||||
std::tie(x0, x1) = at::vec::convert_to_float(x);
|
||||
x0 = x0 * weight_vec;
|
||||
x1 = x1 * weight_vec;
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(out + d);
|
||||
}
|
||||
for (; d < size; ++d) {
|
||||
out[d] = static_cast<scalar_t>(input[d] * weight);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void clamp_sigmoid_and_mul_stub(
|
||||
scalar_t* __restrict__ out,
|
||||
const scalar_t* __restrict__ input,
|
||||
int64_t size,
|
||||
const float alpha,
|
||||
const float limit) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
const fVec one = fVec(1.f);
|
||||
const fVec zero = fVec(0.f);
|
||||
const fVec limit_v = fVec(limit);
|
||||
const fVec nlimit_v = fVec(-limit);
|
||||
const fVec alpha_v = fVec(alpha);
|
||||
|
||||
// no remainder
|
||||
#pragma GCC unroll 4
|
||||
for (int64_t d = 0; d < size; d += bVec::size()) {
|
||||
bVec x = bVec::loadu(input + d);
|
||||
fVec x0_, y0_;
|
||||
std::tie(x0_, y0_) = at::vec::convert_to_float(x);
|
||||
float tmp_buffer[fVec::size() * 2]; // 32
|
||||
float tmp_glu[fVec::size()]; // 16
|
||||
float tmp_linear[fVec::size()]; // 16
|
||||
x0_.store(tmp_buffer);
|
||||
y0_.store(tmp_buffer + fVec::size());
|
||||
// interleaved: x[2i] = glu, x[2i+1] = linear
|
||||
for (int j = 0; j < fVec::size(); ++j) {
|
||||
// x0 [0,2,..30]
|
||||
tmp_glu[j] = tmp_buffer[j * 2];
|
||||
// y0 [1,3,...31]
|
||||
tmp_linear[j] = tmp_buffer[j * 2 + 1];
|
||||
}
|
||||
fVec x0 = fVec::loadu(tmp_glu);
|
||||
fVec y0 = fVec::loadu(tmp_linear);
|
||||
|
||||
// clamp
|
||||
x0 = at::vec::minimum(x0, limit_v);
|
||||
y0 = at::vec::minimum(limit_v, at::vec::maximum(nlimit_v, y0));
|
||||
// x * sigmoid(x * alpha)
|
||||
x0 = x0 / (one + (x0 * alpha_v).neg().exp_u20());
|
||||
// (y + 1) * x
|
||||
y0 = y0 + one;
|
||||
x0 = x0 * y0;
|
||||
convert_from_float_and_store<scalar_t>(out + d / 2, x0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,404 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "common.h"
|
||||
#include "gemm.h"
|
||||
#include "moe.h"
|
||||
|
||||
template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
|
||||
void fused_experts_fp_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
scalar_t* __restrict__ A_tmp,
|
||||
scalar_t* __restrict__ B_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const packed_t* __restrict__ packed_w1,
|
||||
const packed_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1_bias,
|
||||
const float* __restrict__ w2_bias,
|
||||
const param_t* __restrict__ w1s,
|
||||
const param_t* __restrict__ w2s,
|
||||
int64_t block_size_N,
|
||||
int64_t block_size_K,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad,
|
||||
float alpha,
|
||||
float limit,
|
||||
CPUAcTMethod act_func,
|
||||
bool with_bias) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
|
||||
// stage 1: intermediate_cache0 = hidden_states @ w1
|
||||
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
|
||||
const int64_t NB = div_up(2 * N, BLOCK_N);
|
||||
int64_t scale_size_N = div_up(2 * N, block_size_N);
|
||||
int64_t scale_size_K = div_up(K, block_size_K);
|
||||
int64_t blocks_n_per_group = block_size_N / BLOCK_N;
|
||||
std::function<int64_t(int64_t)> scale_offset_per_block;
|
||||
if constexpr (is_mxfp4) {
|
||||
scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
|
||||
} else {
|
||||
scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
|
||||
}
|
||||
|
||||
const int64_t packed_K = get_row_size<packed_t>(K);
|
||||
|
||||
const int64_t stride_e = 2 * N * packed_K;
|
||||
const int64_t stride_n = packed_K;
|
||||
|
||||
int64_t avg_M = std::max(int64_t(1), M * topk / E);
|
||||
const bool use_brgemm = can_use_brgemm<packed_t>(avg_M);
|
||||
|
||||
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
|
||||
|
||||
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// get local pointers
|
||||
int tid = get_thread_num();
|
||||
scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
|
||||
|
||||
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// B shape [K, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
|
||||
const param_t* __restrict__ Bs =
|
||||
w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
|
||||
const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
|
||||
|
||||
// do unpacking for the first row or a new expert
|
||||
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
|
||||
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
|
||||
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
|
||||
if (nb_offset == 0) {
|
||||
// 1.a load A
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m] / topk;
|
||||
copy_stub(A + m * K, input + index * K, K);
|
||||
}
|
||||
}
|
||||
|
||||
const int64_t offset = offsets[mb];
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
|
||||
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
|
||||
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
|
||||
/* Bbias */ B_bias,
|
||||
/* scale */ Bs,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ 2 * N,
|
||||
/* brg */ use_brgemm,
|
||||
/* block_size_K */ block_size_K,
|
||||
/* do_unpack */ do_unpack);
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
|
||||
if (act_func == CPUAcTMethod::silu_and_mul) {
|
||||
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
|
||||
}
|
||||
});
|
||||
} else if (act_func == CPUAcTMethod::swiglu) {
|
||||
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N, alpha, limit);
|
||||
clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N, alpha, limit);
|
||||
}
|
||||
});
|
||||
}
|
||||
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
|
||||
// w2 : [E, K, N] as [E, OC, IC]
|
||||
const int64_t OC = K; // rename K as OC
|
||||
const int64_t IC = N; // rename N as IC
|
||||
const int64_t MB2 = MB;
|
||||
const int64_t NB2 = div_up(OC, BLOCK_N);
|
||||
scale_size_N = div_up(K, block_size_N);
|
||||
scale_size_K = div_up(N, block_size_K);
|
||||
const int64_t packed_IC = get_row_size<packed_t>(IC);
|
||||
const int64_t stride_e2 = OC * packed_IC;
|
||||
const int64_t stride_oc = packed_IC;
|
||||
|
||||
// parallel on [MB2, NB2]
|
||||
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
int tid = get_thread_num();
|
||||
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
|
||||
|
||||
loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// A ptr from ic1 of [M * topk, N] in sorted order
|
||||
// so as to avoid copy A to tmp buffer again
|
||||
const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
|
||||
// B shape [IC, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
|
||||
const param_t* __restrict__ Bs =
|
||||
w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
|
||||
const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
|
||||
|
||||
// do unpacking for the first row or a new expert
|
||||
int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
|
||||
bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
|
||||
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ C,
|
||||
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
|
||||
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
|
||||
/* Bbias */ B_bias,
|
||||
/* scale */ Bs,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* brg */ use_brgemm,
|
||||
/* block_size_K */ block_size_K,
|
||||
/* do_unpack */ do_unpack);
|
||||
|
||||
// 2.b copy from C to ic2 in original order
|
||||
// and also mul topk_weights in float32
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m];
|
||||
float weight = topk_weights[index];
|
||||
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
// stage 3: out = intermediate_cache2.sum(dim=1)
|
||||
// from [M, topk, K] to [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4) \
|
||||
template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
|
||||
TYPE1* __restrict__ output, \
|
||||
TYPE1* __restrict__ ic0, \
|
||||
TYPE1* __restrict__ ic1, \
|
||||
TYPE1* __restrict__ ic2, \
|
||||
TYPE1* __restrict__ A_tmp, \
|
||||
TYPE1* __restrict__ B_tmp, \
|
||||
float* __restrict__ C_tmp, \
|
||||
const TYPE1* __restrict__ input, \
|
||||
const TYPE2* __restrict__ packed_w1, \
|
||||
const TYPE2* __restrict__ packed_w2, \
|
||||
const float* __restrict__ w1_bias, \
|
||||
const float* __restrict__ w2_bias, \
|
||||
const TYPE3* __restrict__ w1s, \
|
||||
const TYPE3* __restrict__ w2s, \
|
||||
int64_t block_size_N, \
|
||||
int64_t block_size_K, \
|
||||
const float* __restrict__ topk_weights, \
|
||||
const int32_t* __restrict__ sorted_ids, \
|
||||
const int32_t* __restrict__ expert_ids, \
|
||||
const int32_t* __restrict__ offsets, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K, \
|
||||
int64_t E, \
|
||||
int64_t topk, \
|
||||
int64_t num_tokens_post_pad, \
|
||||
float alpha, \
|
||||
float limit, \
|
||||
CPUAcTMethod act_func, \
|
||||
bool with_bias)
|
||||
|
||||
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
|
||||
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
|
||||
INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
|
||||
INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
|
||||
|
||||
template <typename scalar_t>
|
||||
void shared_expert_fp8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ B_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w1,
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
int64_t block_size_N,
|
||||
int64_t block_size_K,
|
||||
const scalar_t* __restrict__ fused_experts_out,
|
||||
float routed_scaling_factor,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
|
||||
// stage 1: intermediate_cache0 = hidden_states @ w1
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(2 * N, BLOCK_N);
|
||||
int64_t scale_size_K = div_up(K, block_size_K);
|
||||
int64_t blocks_n_per_group = block_size_N / BLOCK_N;
|
||||
|
||||
const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
|
||||
const bool apply_scaling_factor = fused_experts_out != nullptr;
|
||||
|
||||
int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
|
||||
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
int tid = get_thread_num();
|
||||
|
||||
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
|
||||
int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// do unpacking for the first row
|
||||
bool do_unpack = (mb == mb0);
|
||||
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ input + mb * BLOCK_M * K,
|
||||
/* B */ packed_w1 + nb * BLOCK_N * K,
|
||||
/* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
|
||||
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
|
||||
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
|
||||
/* Bbias */ nullptr,
|
||||
/* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ 2 * N,
|
||||
/* brg */ use_brgemm,
|
||||
/* block_size_K */ block_size_K,
|
||||
/* do_unpack */ do_unpack);
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
|
||||
// w2 : [K, N] as [OC, IC]
|
||||
const int64_t OC = K; // rename K as OC
|
||||
const int64_t IC = N; // rename N as IC
|
||||
const int64_t MB2 = MB;
|
||||
const int64_t NB2 = div_up(K, BLOCK_N);
|
||||
scale_size_K = div_up(N, block_size_K);
|
||||
|
||||
// parallel on [MB2, NB2]
|
||||
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
int tid = get_thread_num();
|
||||
alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
|
||||
|
||||
loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
|
||||
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// do unpacking for the first row
|
||||
bool do_unpack = (mb == mb0);
|
||||
|
||||
// 2.a gemm: C = A @ B
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ ic1 + mb * BLOCK_M * N,
|
||||
/* B */ packed_w2 + nb * BLOCK_N * N,
|
||||
/* C */ C,
|
||||
/* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
|
||||
/* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
|
||||
/* Bbias */ nullptr,
|
||||
/* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* brg */ use_brgemm,
|
||||
/* block_size_K */ block_size_K,
|
||||
/* do_unpack */ do_unpack);
|
||||
|
||||
// 2.b copy from C to output and add fused_experts_out
|
||||
scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
|
||||
const scalar_t* __restrict__ fused_out =
|
||||
apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
|
||||
add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
}
|
||||
|
||||
#define INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(TYPE) \
|
||||
template void shared_expert_fp8_kernel_impl<TYPE>( \
|
||||
TYPE* __restrict__ output, \
|
||||
TYPE* __restrict__ ic0, \
|
||||
TYPE* __restrict__ ic1, \
|
||||
TYPE* __restrict__ B_tmp, \
|
||||
float* __restrict__ C_tmp, \
|
||||
const TYPE* __restrict__ input, \
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w1, \
|
||||
const at::Float8_e4m3fn* __restrict__ packed_w2, \
|
||||
const float* __restrict__ w1s, \
|
||||
const float* __restrict__ w2s, \
|
||||
int64_t block_size_N, \
|
||||
int64_t block_size_K, \
|
||||
const TYPE* __restrict__ fused_experts_out, \
|
||||
float routed_scaling_factor, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K)
|
||||
|
||||
INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::Half);
|
||||
@@ -0,0 +1,323 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "common.h"
|
||||
#include "gemm.h"
|
||||
#include "moe.h"
|
||||
|
||||
template <int64_t N>
|
||||
inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) {
|
||||
using Vec = at::vec::Vectorized<float>;
|
||||
constexpr int kVecSize = Vec::size();
|
||||
static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized<float>::size()");
|
||||
const bool has_bias = bias_ptr != nullptr;
|
||||
const Vec zero_vec(0.f);
|
||||
for (int i = 0; i < m; ++i) {
|
||||
#pragma GCC unroll 2
|
||||
for (int j = 0; j < N; j += kVecSize) {
|
||||
Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec;
|
||||
vec.store(y_buf + i * ldn + j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void fused_experts_int4_w4a8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic0,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
uint8_t* __restrict__ A_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
int32_t* __restrict__ Azp_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
int8_t* __restrict__ dqB_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const uint8_t* __restrict__ packed_w1,
|
||||
const uint8_t* __restrict__ packed_w2,
|
||||
const int8_t* __restrict__ w1z,
|
||||
const int8_t* __restrict__ w2z,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
int group_size,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad) {
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
int num_threads = at::get_num_threads();
|
||||
// int64_t buffer_size_nbytes = M * topk * N * 2
|
||||
// M * topk * K * 2 +
|
||||
// num_threads * BLOCK_M * K +
|
||||
// num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float) +
|
||||
// M * topk * 2 * N * 2 +
|
||||
// max(M * K, M * topk * N) +
|
||||
// M * topk * sizeof(float);
|
||||
|
||||
// intermediate_cache1 (scalar_t): START + M * topk * N
|
||||
// intermediate_cache2 (scalar_t): + M * topk * K
|
||||
// A_tmp (uint8_t): + num_threads * BLOCK_M * K
|
||||
// C_tmp (float): + num_threads * 2 * BLOCK_M * BLOCK_N
|
||||
// intermediate_cache0 (scalar_t): + M * topk * 2 * N
|
||||
// Aq_tmp (uint8_t): + max(M * K, M * topk * N)
|
||||
// As_tmp (float): + M * topk
|
||||
// dqB_tmp (int8_t) + num_threads * _block_k * BlOCK_N
|
||||
|
||||
// stage 0: quantize input to uint8, [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
|
||||
}
|
||||
});
|
||||
int64_t _block_k = get_4bit_block_k_size(group_size);
|
||||
auto Azp = at::ones({M * topk}).to(at::kInt).mul(128);
|
||||
auto Azp_ptr = Azp.data_ptr<int32_t>();
|
||||
// stage 1: intermediate_cache0 = hidden_states @ w1
|
||||
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
int64_t block_per_group = group_size / _block_k;
|
||||
int64_t Kc = K / _block_k;
|
||||
int64_t num_groups = K / group_size;
|
||||
|
||||
const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
|
||||
const bool sym_quant_act = false;
|
||||
// weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
|
||||
// scales/qzeros shape = [E, Nc, G, block_n]
|
||||
|
||||
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
|
||||
at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
|
||||
// get local pointers
|
||||
int tid = at::get_thread_num();
|
||||
int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
|
||||
int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N;
|
||||
alignas(64) float As[BLOCK_M];
|
||||
uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
|
||||
float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
|
||||
bool is_brgemm_used = false;
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
int64_t mb = i / NB;
|
||||
int64_t nb = i % NB;
|
||||
int64_t nb1 = nb + NB;
|
||||
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
|
||||
// B shape [K, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e;
|
||||
// Bz and Bs: [E, K/gs, 2N]
|
||||
const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N);
|
||||
const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N);
|
||||
|
||||
// 1.a load A
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
|
||||
is_brgemm_used = is_brgemm_used || use_brgemm;
|
||||
// copy to A [BLOCK_M, K]
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m] / topk;
|
||||
copy_stub(A + m * K, Aq_tmp + index * K, K);
|
||||
As[m] = As_tmp[index];
|
||||
}
|
||||
const int64_t offset = offsets[mb];
|
||||
copy_bias<BLOCK_N>(nullptr, C0, m_size, BLOCK_N);
|
||||
copy_bias<BLOCK_N>(nullptr, C1, m_size, BLOCK_N);
|
||||
for (int kci = 0; kci < Kc; ++kci) {
|
||||
int32_t* compensation_ptr =
|
||||
sym_quant_act ? nullptr
|
||||
: (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
|
||||
_block_k * BLOCK_N / 2) /*Bcomp*/;
|
||||
tinygemm_kernel<scalar_t>(
|
||||
ic0 + offset * 2 * N + nb * BLOCK_N,
|
||||
C0,
|
||||
A + kci * _block_k,
|
||||
As,
|
||||
Azp_ptr,
|
||||
B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
|
||||
Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
|
||||
Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
|
||||
compensation_ptr,
|
||||
dqB_tmp1,
|
||||
m_size,
|
||||
_block_k,
|
||||
K,
|
||||
BLOCK_N,
|
||||
2 * N,
|
||||
kci == Kc - 1,
|
||||
use_brgemm);
|
||||
}
|
||||
|
||||
for (int kci = 0; kci < Kc; ++kci) {
|
||||
int32_t* compensation_ptr =
|
||||
sym_quant_act ? nullptr
|
||||
: (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
|
||||
_block_k * BLOCK_N / 2) /*Bcomp*/;
|
||||
tinygemm_kernel<scalar_t>(
|
||||
ic0 + offset * 2 * N + nb1 * BLOCK_N,
|
||||
C1,
|
||||
A + kci * _block_k,
|
||||
As,
|
||||
Azp_ptr,
|
||||
B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
|
||||
Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
|
||||
Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
|
||||
compensation_ptr,
|
||||
dqB_tmp2,
|
||||
m_size,
|
||||
_block_k,
|
||||
K,
|
||||
BLOCK_N,
|
||||
2 * N,
|
||||
kci == Kc - 1,
|
||||
use_brgemm);
|
||||
}
|
||||
}
|
||||
|
||||
if (is_brgemm_used) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
|
||||
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
|
||||
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
|
||||
}
|
||||
});
|
||||
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
|
||||
// w2 : [E, K, N] as [E, OC, IC]
|
||||
const int64_t OC = K; // rename K as OC
|
||||
const int64_t IC = N; // rename N as IC
|
||||
const int64_t MB2 = MB;
|
||||
const int64_t NB2 = div_up(OC, BLOCK_N);
|
||||
const int64_t stride_oc = IC;
|
||||
num_groups = IC / group_size;
|
||||
Kc = IC / _block_k;
|
||||
const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
|
||||
// parallel on [MB2, NB2]
|
||||
at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
|
||||
int tid = at::get_thread_num();
|
||||
int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
|
||||
float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
bool is_brgemm_used = false;
|
||||
for (int64_t i = begin; i < end; ++i) {
|
||||
int64_t mb = i / NB2;
|
||||
int64_t nb = i % NB2;
|
||||
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
|
||||
const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
|
||||
is_brgemm_used = is_brgemm_used || use_brgemm;
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
|
||||
// B shape [IC, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2;
|
||||
|
||||
// Bz and Bs: [E, IC/gs, OC]
|
||||
const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC;
|
||||
const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC;
|
||||
|
||||
// A ptr from ic1 of [M * topk, N] in sorted order
|
||||
// so as to avoid copy A to tmp buffer again
|
||||
const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC;
|
||||
const float* __restrict__ As = As_tmp + offsets[mb];
|
||||
copy_bias<BLOCK_N>(nullptr, C2, m_size, BLOCK_N);
|
||||
for (int kci = 0; kci < Kc; ++kci) {
|
||||
int32_t* compensation_ptr =
|
||||
sym_quant_act ? nullptr
|
||||
: (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
|
||||
_block_k * BLOCK_N / 2) /*Bcomp*/;
|
||||
tinygemm_kernel<scalar_t>(
|
||||
nullptr, /*store_out is false*/
|
||||
C2,
|
||||
A + kci * _block_k,
|
||||
As,
|
||||
Azp_ptr,
|
||||
B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
|
||||
Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
|
||||
Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/,
|
||||
compensation_ptr,
|
||||
dqB_tmp1,
|
||||
m_size,
|
||||
_block_k,
|
||||
IC,
|
||||
BLOCK_N,
|
||||
BLOCK_N,
|
||||
false,
|
||||
use_brgemm);
|
||||
}
|
||||
|
||||
// 2.b copy from C to ic2 in original order
|
||||
// and also mul topk_weights in float32
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m];
|
||||
float weight = topk_weights[index];
|
||||
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size);
|
||||
}
|
||||
}
|
||||
|
||||
if (is_brgemm_used) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 3: out = intermediate_cache2.sum(dim=1)
|
||||
// from [M, topk, K] to [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE) \
|
||||
template void fused_experts_int4_w4a8_kernel_impl<TYPE>( \
|
||||
TYPE* __restrict__ output, \
|
||||
TYPE* __restrict__ ic0, \
|
||||
TYPE* __restrict__ ic1, \
|
||||
TYPE* __restrict__ ic2, \
|
||||
uint8_t* __restrict__ A_tmp, \
|
||||
uint8_t* __restrict__ Aq_tmp, \
|
||||
float* __restrict__ As_tmp, \
|
||||
int32_t* __restrict__ Azp_tmp, \
|
||||
float* __restrict__ C_tmp, \
|
||||
int8_t* __restrict__ dqB_tmp, \
|
||||
const TYPE* __restrict__ input, \
|
||||
const uint8_t* __restrict__ packed_w1, \
|
||||
const uint8_t* __restrict__ packed_w2, \
|
||||
const int8_t* __restrict__ w1z, \
|
||||
const int8_t* __restrict__ w2z, \
|
||||
const float* __restrict__ w1s, \
|
||||
const float* __restrict__ w2s, \
|
||||
int group_size, \
|
||||
const float* __restrict__ topk_weights, \
|
||||
const int32_t* __restrict__ sorted_ids, \
|
||||
const int32_t* __restrict__ expert_ids, \
|
||||
const int32_t* __restrict__ offsets, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K, \
|
||||
int64_t E, \
|
||||
int64_t topk, \
|
||||
int64_t num_tokens_post_pad)
|
||||
|
||||
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half);
|
||||
@@ -0,0 +1,973 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#include "common.h"
|
||||
#include "gemm.h"
|
||||
#include "moe.h"
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename scalar_t, int BLOCK_N>
|
||||
inline void silu_and_mul(
|
||||
scalar_t* __restrict__ C,
|
||||
const int32_t* __restrict__ C0, // x: x0, x1
|
||||
const int32_t* __restrict__ C1, // y: y0, y1
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs0,
|
||||
const float* __restrict__ Bs1,
|
||||
const int32_t* __restrict__ Bcomp0,
|
||||
const int32_t* __restrict__ Bcomp1,
|
||||
int64_t m_size,
|
||||
int64_t N) {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
__m512 vc0[COLS];
|
||||
__m512 vc1[COLS];
|
||||
__m512i vcomp0[COLS];
|
||||
__m512i vcomp1[COLS];
|
||||
__m512 vas;
|
||||
__m512 vbs0[COLS];
|
||||
__m512 vbs1[COLS];
|
||||
|
||||
auto load_scale_and_comp = [&](auto col) {
|
||||
vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
|
||||
vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
|
||||
vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
|
||||
vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
|
||||
};
|
||||
Unroll<COLS>{}(load_scale_and_comp);
|
||||
|
||||
auto scalec = [&](auto col, int64_t m) {
|
||||
// update As
|
||||
vas = _mm512_set1_ps(As[m]);
|
||||
// C = As * (C - Bcomp) * Bs
|
||||
__m512i vc32_0 = _mm512_loadu_si512(C0 + m * BLOCK_N + col * 16);
|
||||
__m512i vc32_1 = _mm512_loadu_si512(C1 + m * BLOCK_N + col * 16);
|
||||
vc0[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_0, vcomp0[col]));
|
||||
vc1[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_1, vcomp1[col]));
|
||||
vc0[col] = _mm512_mul_ps(_mm512_mul_ps(vc0[col], vas), vbs0[col]);
|
||||
vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]);
|
||||
};
|
||||
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
const fVec one = fVec(1.f);
|
||||
auto silu_and_mul = [&](auto col) {
|
||||
fVec x = fVec(vc0[col]);
|
||||
fVec y = fVec(vc1[col]);
|
||||
x = x / (one + x.neg().exp_u20());
|
||||
vc0[col] = x * y;
|
||||
};
|
||||
|
||||
auto storec = [&](auto col, int64_t m) {
|
||||
if constexpr (col % 2 == 0) {
|
||||
fVec x0 = fVec(vc0[col + 0]);
|
||||
fVec x1 = fVec(vc0[col + 1]);
|
||||
bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
|
||||
out_vec.store(C + m * N + col * 16);
|
||||
}
|
||||
};
|
||||
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
Unroll<COLS>{}(scalec, m);
|
||||
Unroll<COLS>{}(silu_and_mul);
|
||||
Unroll<COLS>{}(storec, m);
|
||||
}
|
||||
#else
|
||||
TORCH_CHECK(false, "silu_and_mul: scalar path not implemented!");
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int BLOCK_N>
|
||||
inline void scale_C(
|
||||
float* __restrict__ C,
|
||||
const int32_t* __restrict__ Ctmp,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
int64_t m_size) {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
__m512 vc[COLS];
|
||||
__m512i vcomp[COLS];
|
||||
__m512 vas;
|
||||
__m512 vbs[COLS];
|
||||
|
||||
auto load_scale_and_comp = [&](auto col) {
|
||||
vcomp[col] = _mm512_loadu_si512(Bcomp + col * 16);
|
||||
vbs[col] = _mm512_loadu_ps(Bs + col * 16);
|
||||
};
|
||||
Unroll<COLS>{}(load_scale_and_comp);
|
||||
|
||||
auto scalec = [&](auto col, int64_t m) {
|
||||
// update As
|
||||
vas = _mm512_set1_ps(As[m]);
|
||||
// C = As * (C - Bcomp) * Bs
|
||||
__m512i vc32 = _mm512_loadu_si512(Ctmp + m * BLOCK_N + col * 16);
|
||||
vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp[col]));
|
||||
vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vas), vbs[col]);
|
||||
_mm512_storeu_ps(C + m * BLOCK_N + col * 16, vc[col]);
|
||||
};
|
||||
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
Unroll<COLS>{}(scalec, m);
|
||||
}
|
||||
#else
|
||||
TORCH_CHECK(false, "scale_C: scalar path not implemented!");
|
||||
#endif
|
||||
}
|
||||
|
||||
/// gemm for w13
|
||||
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_vnni {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B0,
|
||||
const int8_t* __restrict__ B1,
|
||||
scalar_t* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs0,
|
||||
const float* __restrict__ Bs1,
|
||||
const int32_t* __restrict__ Bcomp0,
|
||||
const int32_t* __restrict__ Bcomp1,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B0,
|
||||
const int8_t* __restrict__ B1,
|
||||
at::BFloat16* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs0,
|
||||
const float* __restrict__ Bs1,
|
||||
const int32_t* __restrict__ Bcomp0,
|
||||
const int32_t* __restrict__ Bcomp1,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int ROWS = BLOCK_M;
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
__m512i va;
|
||||
__m512i vb0[COLS];
|
||||
__m512i vb1[COLS];
|
||||
__m512i vc0[ROWS * COLS];
|
||||
__m512i vc1[ROWS * COLS];
|
||||
__m512i vcomp0[COLS];
|
||||
__m512i vcomp1[COLS];
|
||||
__m512 vas;
|
||||
__m512 vbs0[COLS];
|
||||
__m512 vbs1[COLS];
|
||||
|
||||
auto loadc = [&](auto i) {
|
||||
vc0[i] = _mm512_set1_epi32(0);
|
||||
vc1[i] = _mm512_set1_epi32(0);
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(loadc);
|
||||
|
||||
const int64_t K4 = K >> 2;
|
||||
const int64_t lda4 = lda >> 2;
|
||||
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
|
||||
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
|
||||
const int32_t* b0_ptr = reinterpret_cast<const int32_t*>(B0);
|
||||
const int32_t* b1_ptr = reinterpret_cast<const int32_t*>(B1);
|
||||
|
||||
auto compute = [&](auto i, int64_t k) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
if constexpr (col == 0) {
|
||||
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
|
||||
}
|
||||
if constexpr (row == 0) {
|
||||
vb0[col] = _mm512_loadu_si512(b0_ptr + k * ldb4 + col * 16);
|
||||
vb1[col] = _mm512_loadu_si512(b1_ptr + k * ldb4 + col * 16);
|
||||
}
|
||||
vc0[i] = _mm512_dpbusd_epi32(vc0[i], va, vb0[col]);
|
||||
vc1[i] = _mm512_dpbusd_epi32(vc1[i], va, vb1[col]);
|
||||
};
|
||||
for (int64_t k = 0; k < K4; ++k) {
|
||||
Unroll<ROWS * COLS>{}(compute, k);
|
||||
}
|
||||
|
||||
auto scalec = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
// load a scale
|
||||
if constexpr (col == 0) {
|
||||
vas = _mm512_set1_ps(As[row]);
|
||||
}
|
||||
// load b scale and vcomp
|
||||
if constexpr (row == 0) {
|
||||
vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
|
||||
vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
|
||||
vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
|
||||
vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
|
||||
}
|
||||
__m512 c0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc0[i], vcomp0[col]));
|
||||
__m512 c1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc1[i], vcomp1[col]));
|
||||
vc0[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c0, vas), vbs0[col]));
|
||||
vc1[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c1, vas), vbs1[col]));
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(scalec);
|
||||
|
||||
using Vec = at::vec::Vectorized<float>;
|
||||
const Vec one = Vec(1.f);
|
||||
auto storec = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
// for COLS = 2, 4 use 512bit store
|
||||
if constexpr (col % 2 == 0) {
|
||||
Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
|
||||
Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
|
||||
Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
|
||||
Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
|
||||
// silu
|
||||
x0 = x0 / (one + x0.neg().exp_u20());
|
||||
x1 = x1 / (one + x1.neg().exp_u20());
|
||||
// mul
|
||||
x0 = x0 * y0;
|
||||
x1 = x1 * y1;
|
||||
|
||||
_mm512_storeu_si512(
|
||||
reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
|
||||
(__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
|
||||
}
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(storec);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
|
||||
tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply( \
|
||||
A + mb_start * lda, \
|
||||
B0 + nb_start * 4, \
|
||||
B1 + nb_start * 4, \
|
||||
C + mb_start * ldc + nb_start, \
|
||||
As + mb_start, \
|
||||
Bs0 + nb_start, \
|
||||
Bs1 + nb_start, \
|
||||
Bcomp0 + nb_start, \
|
||||
Bcomp1 + nb_start, \
|
||||
K, \
|
||||
lda, \
|
||||
ldb, \
|
||||
ldc);
|
||||
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B0,
|
||||
const int8_t* __restrict__ B1,
|
||||
scalar_t* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs0,
|
||||
const float* __restrict__ Bs1,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
|
||||
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
|
||||
|
||||
// pattern: 1-(2+2)-(8+8)
|
||||
constexpr int64_t BLOCK_M = 4;
|
||||
constexpr int64_t BLOCK_N = 32;
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
for (int mb = 0; mb < MB; ++mb) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
|
||||
for (int64_t nb = 0; nb < NB; ++nb) {
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
|
||||
|
||||
switch (mb_size << 4 | nb_size >> 4) {
|
||||
case 0x12:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI(1, 32);
|
||||
break;
|
||||
case 0x22:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI(2, 32);
|
||||
break;
|
||||
case 0x32:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI(3, 32);
|
||||
break;
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI(4, 32);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// gemm for w2
|
||||
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_vnni2 {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
float* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
|
||||
}
|
||||
};
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <int BLOCK_M, int BLOCK_N>
|
||||
struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
|
||||
static inline void apply(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
float* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
const int32_t* __restrict__ Bcomp,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
constexpr int ROWS = BLOCK_M;
|
||||
constexpr int COLS = BLOCK_N / 16;
|
||||
static_assert(COLS % 2 == 0);
|
||||
|
||||
__m512i va;
|
||||
__m512i vb[COLS];
|
||||
__m512i vc[ROWS * COLS];
|
||||
__m512i vcomp[COLS];
|
||||
__m512 vas;
|
||||
__m512 vbs[COLS];
|
||||
|
||||
auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
|
||||
Unroll<ROWS * COLS>{}(loadc);
|
||||
|
||||
const int64_t K4 = K >> 2;
|
||||
const int64_t lda4 = lda >> 2;
|
||||
const int64_t ldb4 = ldb; // ldb * 4 >> 2;
|
||||
const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
|
||||
const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
|
||||
|
||||
auto compute = [&](auto i, int64_t k) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
if constexpr (col == 0) {
|
||||
va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
|
||||
}
|
||||
if constexpr (row == 0) {
|
||||
vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
|
||||
}
|
||||
vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
|
||||
};
|
||||
for (int64_t k = 0; k < K4; ++k) {
|
||||
Unroll<ROWS * COLS>{}(compute, k);
|
||||
}
|
||||
|
||||
auto storec = [&](auto i) {
|
||||
constexpr int row = i / COLS;
|
||||
constexpr int col = i % COLS;
|
||||
|
||||
// load a scale
|
||||
if constexpr (col == 0) {
|
||||
vas = _mm512_set1_ps(As[row]);
|
||||
}
|
||||
// load b scale and vcomp per 2 vectors
|
||||
// also load bias if any
|
||||
if constexpr (row == 0) {
|
||||
if constexpr (col % 2 == 0) {
|
||||
vbs[col + 0] = _mm512_loadu_ps(Bs + col * 16);
|
||||
vbs[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
|
||||
vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
|
||||
vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
|
||||
}
|
||||
}
|
||||
__m512 x = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[i], vcomp[col]));
|
||||
x = _mm512_mul_ps(_mm512_mul_ps(x, vas), vbs[col]);
|
||||
_mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x);
|
||||
};
|
||||
Unroll<ROWS * COLS>{}(storec);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
#define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE) \
|
||||
tinygemm_kernel_vnni2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
|
||||
A + mb_start * lda, \
|
||||
B + nb_start * 4, \
|
||||
C + mb_start * ldc + nb_start, \
|
||||
As + mb_start, \
|
||||
Bs + nb_start, \
|
||||
Bcomp + nb_start, \
|
||||
K, \
|
||||
lda, \
|
||||
ldb, \
|
||||
ldc);
|
||||
|
||||
template <typename scalar_t>
|
||||
void tinygemm_kernel(
|
||||
const uint8_t* __restrict__ A,
|
||||
const int8_t* __restrict__ B,
|
||||
float* __restrict__ C,
|
||||
const float* __restrict__ As,
|
||||
const float* __restrict__ Bs,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t lda,
|
||||
int64_t ldb,
|
||||
int64_t ldc) {
|
||||
// B compensation
|
||||
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
|
||||
|
||||
// pattern: 1-4-16
|
||||
constexpr int64_t BLOCK_M = 4;
|
||||
constexpr int64_t BLOCK_N = 64;
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
for (int64_t mb = 0; mb < MB; ++mb) {
|
||||
int64_t mb_start = mb * BLOCK_M;
|
||||
int64_t mb_size = std::min(BLOCK_M, M - mb_start);
|
||||
for (int64_t nb = 0; nb < NB; ++nb) {
|
||||
int64_t nb_start = nb * BLOCK_N;
|
||||
int64_t nb_size = std::min(BLOCK_N, N - nb_start);
|
||||
|
||||
switch (mb_size << 4 | nb_size >> 4) {
|
||||
case 0x12:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI2(1, 32);
|
||||
break;
|
||||
case 0x22:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI2(2, 32);
|
||||
break;
|
||||
case 0x32:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI2(3, 32);
|
||||
break;
|
||||
case 0x42:
|
||||
LAUNCH_TINYGEMM_KERNEL_VNNI2(4, 32);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
template <typename scalar_t>
|
||||
void fused_experts_int8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic1,
|
||||
scalar_t* __restrict__ ic2,
|
||||
uint8_t* __restrict__ A_tmp,
|
||||
float* __restrict__ C_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const int8_t* __restrict__ packed_w1,
|
||||
const int8_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
const float* __restrict__ topk_weights,
|
||||
const int32_t* __restrict__ sorted_ids,
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
const int32_t* __restrict__ offsets,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
int64_t E,
|
||||
int64_t topk,
|
||||
int64_t num_tokens_post_pad) {
|
||||
// handle 2 tiles per block
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
|
||||
// stage 0: quantize input to uint8, [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1: intermediate_cache1 = silu(hidden_states @ w1)
|
||||
const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
// strides for w1: [E, 2N, K]
|
||||
TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
|
||||
|
||||
// K and N are packed for int8
|
||||
const int64_t packed_K = get_row_size<int8_t>(K);
|
||||
const int64_t packed_N = get_row_size<int8_t>(N);
|
||||
|
||||
const int64_t stride_e = 2 * N * packed_K;
|
||||
const int64_t stride_n = packed_K;
|
||||
|
||||
int64_t avg_M = std::max(int64_t(1), M * topk / E);
|
||||
const bool use_brgemm = can_use_brgemm<int8_t>(avg_M);
|
||||
|
||||
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// get local pointers
|
||||
int tid = get_thread_num();
|
||||
uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
|
||||
int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
|
||||
|
||||
alignas(64) float As[BLOCK_M];
|
||||
|
||||
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
// nb_upper from top half and nb_lower from bottom half
|
||||
int64_t nb_upper = nb, nb_lower = nb + NB;
|
||||
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// B shape [K, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const int8_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
|
||||
const int8_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
|
||||
const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb_upper * BLOCK_N;
|
||||
const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb_lower * BLOCK_N;
|
||||
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
|
||||
if (nb_offset == 0) {
|
||||
// 1.a load A
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m] / topk;
|
||||
copy_stub(A + m * K, Aq_tmp + index * K, K);
|
||||
As[m] = As_tmp[index];
|
||||
}
|
||||
}
|
||||
|
||||
if (use_brgemm) {
|
||||
// 1.b gemm: C0 = A @ B0
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B0,
|
||||
/* C */ C0);
|
||||
|
||||
// 1.c gemm: C1 = A @ B1
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B1,
|
||||
/* C */ C1);
|
||||
|
||||
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
|
||||
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
|
||||
|
||||
// 1.d silu and mul
|
||||
const int64_t offset = offsets[mb];
|
||||
silu_and_mul<scalar_t, BLOCK_N>(
|
||||
ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
|
||||
} else {
|
||||
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
|
||||
const int64_t offset = offsets[mb];
|
||||
tinygemm_kernel(
|
||||
/* A */ A,
|
||||
/* B0 */ B0,
|
||||
/* B1 */ B1,
|
||||
/* C */ ic1 + offset * N + nb * BLOCK_N,
|
||||
/* As */ As,
|
||||
/* Bs0 */ Bs0,
|
||||
/* Bs1 */ Bs1,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ N);
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
|
||||
at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
|
||||
// w2 : [E, K, N] as [E, OC, IC]
|
||||
const int64_t OC = K; // rename K as OC
|
||||
const int64_t IC = N; // rename N as IC
|
||||
const int64_t MB2 = MB;
|
||||
const int64_t NB2 = div_up(OC, BLOCK_N);
|
||||
const int64_t stride_e2 = OC * packed_N;
|
||||
const int64_t stride_oc = packed_N;
|
||||
|
||||
// parallel on [MB2, NB2]
|
||||
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// get local pointers
|
||||
int tid = get_thread_num();
|
||||
float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
|
||||
|
||||
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t m_size = offsets[mb + 1] - offsets[mb];
|
||||
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// A ptr from ic1 of [M * topk, N] in sorted order
|
||||
// so as to avoid copy A to tmp buffer again
|
||||
const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * N;
|
||||
const float* __restrict__ As = As_tmp + offsets[mb];
|
||||
const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
|
||||
|
||||
// B shape [IC, n_size] in vnni format
|
||||
int32_t expert_id = expert_ids[mb];
|
||||
const int8_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
|
||||
const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N;
|
||||
|
||||
// 2.a gemm: C = A @ B
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ C32);
|
||||
|
||||
// apply scales
|
||||
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
|
||||
scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
|
||||
} else {
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ C,
|
||||
/* As */ As,
|
||||
/* Bs */ Bs,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N);
|
||||
}
|
||||
|
||||
// 2.b copy from C to ic2 in original order
|
||||
// and also mul topk_weights in float32
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
int32_t index = A_ids[m];
|
||||
float weight = topk_weights[index];
|
||||
copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 3: out = intermediate_cache2.sum(dim=1)
|
||||
// from [M, topk, K] to [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#define INSTANTIATE_MOE_INT8_TEMPLATE(TYPE) \
|
||||
template void fused_experts_int8_kernel_impl<TYPE>( \
|
||||
TYPE* __restrict__ output, \
|
||||
TYPE* __restrict__ ic1, \
|
||||
TYPE* __restrict__ ic2, \
|
||||
uint8_t* __restrict__ A_tmp, \
|
||||
float* __restrict__ C_tmp, \
|
||||
uint8_t* __restrict__ Aq_tmp, \
|
||||
float* __restrict__ As_tmp, \
|
||||
const TYPE* __restrict__ input, \
|
||||
const int8_t* __restrict__ packed_w1, \
|
||||
const int8_t* __restrict__ packed_w2, \
|
||||
const float* __restrict__ w1s, \
|
||||
const float* __restrict__ w2s, \
|
||||
const float* __restrict__ topk_weights, \
|
||||
const int32_t* __restrict__ sorted_ids, \
|
||||
const int32_t* __restrict__ expert_ids, \
|
||||
const int32_t* __restrict__ offsets, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K, \
|
||||
int64_t E, \
|
||||
int64_t topk, \
|
||||
int64_t num_tokens_post_pad)
|
||||
|
||||
INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
|
||||
|
||||
template <typename scalar_t>
|
||||
void shared_expert_int8_kernel_impl(
|
||||
scalar_t* __restrict__ output,
|
||||
scalar_t* __restrict__ ic1,
|
||||
float* __restrict__ C_tmp,
|
||||
uint8_t* __restrict__ Aq_tmp,
|
||||
float* __restrict__ As_tmp,
|
||||
const scalar_t* __restrict__ input,
|
||||
const int8_t* __restrict__ packed_w1,
|
||||
const int8_t* __restrict__ packed_w2,
|
||||
const float* __restrict__ w1s,
|
||||
const float* __restrict__ w2s,
|
||||
const scalar_t* __restrict__ fused_experts_out,
|
||||
float routed_scaling_factor,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
// handle 2 tiles per block
|
||||
constexpr int64_t BLOCK_M = block_size_m();
|
||||
constexpr int64_t BLOCK_N = block_size_n();
|
||||
|
||||
// stage 0: quantize input to uint8, [M, K]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1: intermediate_cache1 = silu(hidden_states @ w1)
|
||||
const int64_t MB = div_up(M, BLOCK_M);
|
||||
const int64_t NB = div_up(N, BLOCK_N);
|
||||
|
||||
TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
|
||||
|
||||
// K and N are packed for int8
|
||||
const int64_t packed_K = get_row_size<int8_t>(K);
|
||||
const int64_t packed_N = get_row_size<int8_t>(N);
|
||||
const int64_t stride_n = packed_K;
|
||||
|
||||
const bool use_brgemm = can_use_brgemm<int8_t>(M);
|
||||
const bool apply_scaling_factor = fused_experts_out != nullptr;
|
||||
|
||||
// here we only parallel on half of 2N to fuse silu_and_mul with gemm
|
||||
parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// get local pointers
|
||||
int tid = get_thread_num();
|
||||
int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
|
||||
|
||||
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
// nb_upper from top half and nb_lower from bottom half
|
||||
int64_t nb_upper = nb, nb_lower = nb + NB;
|
||||
int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
|
||||
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
|
||||
|
||||
// A shape [m_size, K]
|
||||
const uint8_t* A = Aq_tmp + mb * BLOCK_M * K;
|
||||
const float* As = As_tmp + mb * BLOCK_M;
|
||||
|
||||
// B shape [K, n_size] in vnni format
|
||||
const int8_t* __restrict__ B0 = packed_w1 + nb_upper * BLOCK_N * stride_n;
|
||||
const int8_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;
|
||||
const float* __restrict__ Bs0 = w1s + nb_upper * BLOCK_N;
|
||||
const float* __restrict__ Bs1 = w1s + nb_lower * BLOCK_N;
|
||||
|
||||
if (use_brgemm) {
|
||||
// 1.b gemm: C0 = A @ B0
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B0,
|
||||
/* C */ C0);
|
||||
|
||||
// 1.c gemm: C1 = A @ B1
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B1,
|
||||
/* C */ C1);
|
||||
|
||||
const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
|
||||
const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
|
||||
|
||||
// 1.d silu and mul
|
||||
silu_and_mul<scalar_t, BLOCK_N>(
|
||||
ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
|
||||
} else {
|
||||
// fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
|
||||
tinygemm_kernel(
|
||||
/* A */ A,
|
||||
/* B0 */ B0,
|
||||
/* B1 */ B1,
|
||||
/* C */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
|
||||
/* As */ As,
|
||||
/* Bs0 */ Bs0,
|
||||
/* Bs1 */ Bs1,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ K,
|
||||
/* lda */ K,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ N);
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
|
||||
// stage 1.5: quantize ic1 to uint8, [M * topk, N]
|
||||
at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
|
||||
for (int64_t m = begin; m < end; ++m) {
|
||||
quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
|
||||
}
|
||||
});
|
||||
|
||||
// stage 2: intermediate_cache2 = intermediate_cache1 @ w2
|
||||
// w2 : [K, N] as [OC, IC]
|
||||
const int64_t OC = K; // rename K as OC
|
||||
const int64_t IC = N; // rename N as IC
|
||||
const int64_t MB2 = MB;
|
||||
const int64_t NB2 = div_up(OC, BLOCK_N);
|
||||
const int64_t stride_oc = packed_N;
|
||||
|
||||
// parallel on [MB2, NB2]
|
||||
parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
|
||||
// get local pointers
|
||||
int tid = get_thread_num();
|
||||
float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
|
||||
int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
|
||||
|
||||
loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
|
||||
int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
|
||||
int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
|
||||
|
||||
// A shape [m_size, IC]
|
||||
const uint8_t* __restrict__ A = Aq_tmp + mb * BLOCK_M * N;
|
||||
const float* __restrict__ As = As_tmp + mb * BLOCK_M;
|
||||
|
||||
// B shape [IC, n_size] in vnni format
|
||||
const int8_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
|
||||
const float* __restrict__ Bs = w2s + nb * BLOCK_N;
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm(
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N,
|
||||
/* add_C */ false,
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ C32);
|
||||
|
||||
// apply scales
|
||||
const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
|
||||
scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
|
||||
} else {
|
||||
// 2.a gemm: C = A @ B
|
||||
tinygemm_kernel<scalar_t>(
|
||||
/* A */ A,
|
||||
/* B */ B,
|
||||
/* C */ C,
|
||||
/* As */ As,
|
||||
/* Bs */ Bs,
|
||||
/* M */ m_size,
|
||||
/* N */ n_size,
|
||||
/* K */ IC,
|
||||
/* lda */ IC,
|
||||
/* ldb */ n_size,
|
||||
/* ldc */ BLOCK_N);
|
||||
}
|
||||
|
||||
// 2.b copy from C to output and add fused_experts_out
|
||||
scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
|
||||
const scalar_t* __restrict__ fused_out =
|
||||
apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
|
||||
for (int64_t m = 0; m < m_size; ++m) {
|
||||
const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
|
||||
add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
|
||||
}
|
||||
});
|
||||
|
||||
if (use_brgemm) {
|
||||
at::native::cpublas::brgemm_release();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#define INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(TYPE) \
|
||||
template void shared_expert_int8_kernel_impl<TYPE>( \
|
||||
TYPE* __restrict__ output, \
|
||||
TYPE* __restrict__ ic1, \
|
||||
float* __restrict__ C_tmp, \
|
||||
uint8_t* __restrict__ Aq_tmp, \
|
||||
float* __restrict__ As_tmp, \
|
||||
const TYPE* __restrict__ input, \
|
||||
const int8_t* __restrict__ packed_w1, \
|
||||
const int8_t* __restrict__ packed_w2, \
|
||||
const float* __restrict__ w1s, \
|
||||
const float* __restrict__ w2s, \
|
||||
const TYPE* __restrict__ fused_experts_out, \
|
||||
float routed_scaling_factor, \
|
||||
int64_t M, \
|
||||
int64_t N, \
|
||||
int64_t K)
|
||||
|
||||
INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::Half);
|
||||
@@ -0,0 +1,449 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
#pragma once
|
||||
|
||||
#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
|
||||
#define CPU_CAPABILITY_AVX512
|
||||
#endif
|
||||
|
||||
#if defined(__riscv_v_min_vlen) && (__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
#define CPU_CAPABILITY_RVV
|
||||
#endif
|
||||
|
||||
#include <ATen/cpu/vec/functional.h>
|
||||
#include <ATen/cpu/vec/vec.h>
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
#include <immintrin.h>
|
||||
#endif
|
||||
|
||||
#if defined(CPU_CAPABILITY_RVV)
|
||||
#include "../cpu_types_riscv_defs.hpp"
|
||||
#endif
|
||||
namespace {
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
|
||||
inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
|
||||
return at::vec::convert_from_float<scalar_t>(a, b);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
inline void convert_from_float_and_store(scalar_t* out, const Vectorized<float>& a) {
|
||||
float out_buffer[at::vec::Vectorized<float>::size()];
|
||||
a.store(out_buffer);
|
||||
for (int i = 0; i < 16; i++) {
|
||||
out[i] = (scalar_t)out_buffer[i];
|
||||
}
|
||||
}
|
||||
|
||||
// allow f16, bf16
|
||||
template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
|
||||
inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
|
||||
using bVec = at::vec::Vectorized<scalar_t>;
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
bVec x_vec = bVec::loadu(data);
|
||||
fVec x0, x1;
|
||||
std::tie(x0, x1) = at::vec::convert_to_float(x_vec);
|
||||
return std::make_tuple(x0, x1);
|
||||
}
|
||||
|
||||
// allow f32
|
||||
inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const float* __restrict__ data) {
|
||||
using fVec = at::vec::Vectorized<float>;
|
||||
fVec x0 = fVec::loadu(data);
|
||||
fVec x1 = fVec::loadu(data + fVec::size());
|
||||
return std::make_tuple(x0, x1);
|
||||
}
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
|
||||
// `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
|
||||
// use native instruction for bfloat16->float32 conversion
|
||||
template <>
|
||||
inline Vectorized<at::BFloat16>
|
||||
convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
|
||||
return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
|
||||
}
|
||||
|
||||
template <>
|
||||
inline void convert_from_float_and_store<at::BFloat16>(at::BFloat16* out, const Vectorized<float>& a) {
|
||||
_mm256_storeu_si256((__m256i*)out, (__m256i)(_mm512_cvtneps_pbh(__m512(a))));
|
||||
}
|
||||
|
||||
#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
|
||||
|
||||
#define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
|
||||
|
||||
// this doesn't handle NaN.
|
||||
inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
|
||||
const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
|
||||
__m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780));
|
||||
combined = _mm512_slli_epi16(combined, 4);
|
||||
combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0));
|
||||
combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00));
|
||||
|
||||
const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
|
||||
return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
|
||||
}
|
||||
|
||||
inline __m512bh cvt_e4m3_bf16_intrinsic_without_denorm(__m256i fp8_vec) {
|
||||
// The following conversion is without denorm behavior, that is to say,
|
||||
// Max subnorm : S.0000.111 = 0.875 ∗ 2**(−6)
|
||||
// Min subnorm : S.0000.001 = 2**(−9)
|
||||
// 0.0019 ~ 0.0137 cannot be converted correctly.
|
||||
__m512i x = _mm512_cvtepu8_epi16(fp8_vec);
|
||||
auto mask = _mm512_cmpneq_epi16_mask(
|
||||
_mm512_and_si512(x, _mm512_set1_epi16(127)),
|
||||
_mm512_setzero_si512()); // mask = x & 0x7f
|
||||
auto mask_nan = _mm512_cmpneq_epi16_mask(
|
||||
_mm512_and_si512(x, _mm512_set1_epi16(127)),
|
||||
_mm512_set1_epi16(127)); // mask_nan = x & 0x7f
|
||||
auto mantissa = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4); // mantissa = (x & 7) << 4
|
||||
auto exponent = _mm512_add_epi16(
|
||||
_mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3),
|
||||
_mm512_set1_epi16(120)); // exponent = (((x >> 3) & 15) + 120)
|
||||
auto nonsign = _mm512_maskz_mov_epi16(mask, _mm512_or_si512(mantissa, _mm512_slli_epi16(exponent, 7)));
|
||||
nonsign = _mm512_mask_mov_epi16(_mm512_set1_epi16(0x7fff), mask_nan, nonsign); // deal with Nan
|
||||
return (__m512bh)(_mm512_or_si512(
|
||||
nonsign,
|
||||
_mm512_slli_epi16(
|
||||
_mm512_and_si512(x, _mm512_set1_epi16(128)),
|
||||
8))); // add sign (x & 128) << 8
|
||||
}
|
||||
|
||||
inline __m512bh cvt_e4m3_bf16_intrinsic_with_denorm(__m256i fp8_vec) {
|
||||
__m512i x = _mm512_cvtepu8_epi16(fp8_vec);
|
||||
__m512i lg2mant = _mm512_mask_mov_epi16(
|
||||
_mm512_mask_mov_epi16(
|
||||
_mm512_setzero_si512(), _mm512_test_epi16_mask(x, _mm512_set1_epi16(2)), _mm512_set1_epi16(1)),
|
||||
_mm512_test_epi16_mask(x, _mm512_set1_epi16(4)),
|
||||
_mm512_set1_epi16(2));
|
||||
return (__m512bh)(_mm512_or_si512(
|
||||
_mm512_maskz_mov_epi16(
|
||||
_mm512_cmpneq_epi16_mask(_mm512_and_si512(x, _mm512_set1_epi16(127)), _mm512_setzero_si512()),
|
||||
_mm512_mask_blend_epi16(
|
||||
_mm512_test_epi16_mask(x, _mm512_set1_epi16(120)),
|
||||
_mm512_or_si512(
|
||||
_mm512_and_si512(
|
||||
_mm512_sllv_epi16(
|
||||
_mm512_and_si512(x, _mm512_set1_epi16(3)), _mm512_sub_epi16(_mm512_set1_epi16(7), lg2mant)),
|
||||
_mm512_set1_epi16(0x007f)),
|
||||
_mm512_slli_epi16(_mm512_add_epi16(lg2mant, _mm512_set1_epi16(118)), 7)),
|
||||
_mm512_or_si512(
|
||||
_mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4),
|
||||
_mm512_slli_epi16(
|
||||
_mm512_add_epi16(
|
||||
_mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), _mm512_set1_epi16(120)),
|
||||
7)))),
|
||||
_mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(128)), 8)));
|
||||
}
|
||||
|
||||
inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
|
||||
#ifdef SGLANG_CPU_FP8_CVT_FTZ
|
||||
return cvt_e4m3_bf16_intrinsic_no_nan(a);
|
||||
#else
|
||||
return cvt_e4m3_bf16_intrinsic_with_denorm(a);
|
||||
#endif
|
||||
}
|
||||
// faster version of float8_e4m3fn conversion to bfloat16
|
||||
//
|
||||
// we mapped cuda implementation from below link and vectorized with avx512:
|
||||
// https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387
|
||||
//
|
||||
inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) {
|
||||
const __m512i mask0 = _mm512_set1_epi16(0x80); // sign bit
|
||||
const __m512i mask1 = _mm512_set1_epi16(0x7F); // exponent and mantissa
|
||||
const __m512i mask2 = _mm512_set1_epi16(0x4000);
|
||||
|
||||
__m512i x = _mm512_cvtepu8_epi16(a);
|
||||
__m512i vsign = _mm512_and_si512(x, mask0);
|
||||
vsign = _mm512_slli_epi16(vsign, 8);
|
||||
|
||||
__m512i vexp_and_mant = _mm512_and_si512(x, mask1);
|
||||
vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4);
|
||||
|
||||
// _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110
|
||||
return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110));
|
||||
}
|
||||
|
||||
// bias for conversion of fp8 to bf16 1/256 in float32
|
||||
#define kFP8_BIAS 0x3b800000
|
||||
|
||||
// remove warning: ignoring attributes on template argument ‘__m512bh’ [-Wignored-attributes]
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wignored-attributes"
|
||||
|
||||
#define MXFP4_VALUES \
|
||||
-6.0f, -4.0f, -3.0f, -2.0f, -1.5f, -1.0f, -0.5f, -0.0f, 6.0f, 4.0f, 3.0f, 2.0f, 1.5f, 1.0f, 0.5f, 0.0f
|
||||
|
||||
// convert 64 mxfp4 to 2x bf16 vectors, expect input 32-way packing
|
||||
inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i a, __m512i s0, __m512i s1) {
|
||||
// LUT
|
||||
const __m512 values = _mm512_set_ps(MXFP4_VALUES);
|
||||
const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values));
|
||||
|
||||
const __m512i abs_mask = _mm512_set1_epi16(0x7FFF);
|
||||
const __m512i zero = _mm512_setzero_si512();
|
||||
|
||||
// expand values to 16-bit integers
|
||||
__m512i x0 = _mm512_cvtepu8_epi16(a);
|
||||
__m512i x1 = _mm512_srli_epi32(x0, 4);
|
||||
|
||||
// LUT to convert mxfp4 values to bf16
|
||||
x0 = _mm512_permutexvar_epi16(x0, lut);
|
||||
x1 = _mm512_permutexvar_epi16(x1, lut);
|
||||
|
||||
// check for zeros
|
||||
__mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ);
|
||||
__mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ);
|
||||
|
||||
// emulate bf16 mul with scale factor
|
||||
x0 = _mm512_add_epi16(x0, s0);
|
||||
x1 = _mm512_add_epi16(x1, s1);
|
||||
|
||||
// blend with zero
|
||||
x0 = _mm512_mask_blend_epi16(mask0, x0, zero);
|
||||
x1 = _mm512_mask_blend_epi16(mask1, x1, zero);
|
||||
|
||||
return std::make_tuple(__m512bh(x0), __m512bh(x1));
|
||||
}
|
||||
|
||||
#define CVT_MXFP4_TO_BF16(a, s0, s1) cvt_mxfp4_e2m1_bf16_intrinsic_lut(a, s0, s1)
|
||||
|
||||
#pragma GCC diagnostic pop
|
||||
|
||||
#endif
|
||||
|
||||
// vector to scalar reduction
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
inline float vec_reduce_sum(const Vectorized<float>& a) {
|
||||
return _mm512_reduce_add_ps(__m512(a));
|
||||
}
|
||||
|
||||
inline float vec_reduce_max(const Vectorized<float>& a) {
|
||||
return _mm512_reduce_max_ps(__m512(a));
|
||||
}
|
||||
#else
|
||||
inline float vec_reduce_sum(const Vectorized<float>& a) {
|
||||
return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return x + y; }, a);
|
||||
}
|
||||
|
||||
inline float vec_reduce_max(const Vectorized<float>& a) {
|
||||
return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return maximum(x, y); }, a);
|
||||
}
|
||||
#endif
|
||||
|
||||
// https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
|
||||
template <typename scalar_t>
|
||||
inline void
|
||||
quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
|
||||
float amax = 0.f; // absolute max
|
||||
for (int64_t k = 0; k < K; ++k) {
|
||||
const float val = static_cast<float>(A[k]);
|
||||
amax = std::max(amax, std::abs(val));
|
||||
}
|
||||
|
||||
amax = std::max(amax, eps);
|
||||
const float scale = amax / 127;
|
||||
const float inv_scale = 127 / amax;
|
||||
|
||||
for (int64_t k = 0; k < K; ++k) {
|
||||
const float val = static_cast<float>(A[k]) * inv_scale;
|
||||
Aq[k] = static_cast<uint8_t>(static_cast<int32_t>(std::round(val)) + 128);
|
||||
}
|
||||
As = scale;
|
||||
}
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
template <>
|
||||
inline void quantize_row_int8<at::BFloat16>(
|
||||
uint8_t* __restrict__ Aq, float& As, const at::BFloat16* __restrict__ A, int64_t K, float eps) {
|
||||
const __m512 signBit = _mm512_set1_ps(-0.0f);
|
||||
const __m512i off = _mm512_set1_epi32(128);
|
||||
|
||||
// K is 32x, no remainder
|
||||
float amax = 0.f;
|
||||
__m512 vamax0 = _mm512_set1_ps(0.f);
|
||||
__m512 vamax1 = _mm512_set1_ps(0.f);
|
||||
for (int64_t k = 0; k < K; k += 32) {
|
||||
__m512i va = _mm512_loadu_si512((void*)(A + k));
|
||||
__m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
|
||||
__m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
|
||||
vamax0 = _mm512_max_ps(vamax0, _mm512_andnot_ps(signBit, va0));
|
||||
vamax1 = _mm512_max_ps(vamax1, _mm512_andnot_ps(signBit, va1));
|
||||
}
|
||||
amax = _mm512_reduce_max_ps(_mm512_max_ps(vamax0, vamax1));
|
||||
amax = std::max(amax, eps);
|
||||
const float scale = amax / 127;
|
||||
const float inv_scale = 127 / amax;
|
||||
const __m512 vd = _mm512_set1_ps(inv_scale);
|
||||
|
||||
for (int64_t k = 0; k < K; k += 32) {
|
||||
__m512i va = _mm512_loadu_si512((void*)(A + k));
|
||||
__m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
|
||||
__m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
|
||||
va0 = _mm512_mul_ps(va0, vd);
|
||||
va1 = _mm512_mul_ps(va1, vd);
|
||||
va0 = _mm512_roundscale_ps(va0, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
|
||||
va1 = _mm512_roundscale_ps(va1, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
|
||||
__m128i i0 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va0), off));
|
||||
__m128i i1 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va1), off));
|
||||
_mm256_storeu_si256(reinterpret_cast<__m256i*>(Aq + k), _mm256_set_m128i(i1, i0));
|
||||
}
|
||||
As = scale;
|
||||
}
|
||||
#endif
|
||||
|
||||
// transpose utils
|
||||
// taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
inline void transpose_16x16_32bit(__m512i* v) {
|
||||
__m512i v1[16];
|
||||
v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
|
||||
v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
|
||||
v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
|
||||
v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
|
||||
v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
|
||||
v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
|
||||
v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
|
||||
v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
|
||||
v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
|
||||
v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
|
||||
v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
|
||||
v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
|
||||
v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
|
||||
v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
|
||||
v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
|
||||
v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
|
||||
|
||||
v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
|
||||
v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
|
||||
v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
|
||||
v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
|
||||
v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
|
||||
v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
|
||||
v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
|
||||
v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
|
||||
v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
|
||||
v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
|
||||
v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
|
||||
v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
|
||||
v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
|
||||
v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
|
||||
v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
|
||||
v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
|
||||
|
||||
v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
|
||||
v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
|
||||
v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
|
||||
v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
|
||||
v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
|
||||
v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
|
||||
v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
|
||||
v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
|
||||
v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
|
||||
v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
|
||||
v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
|
||||
v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
|
||||
v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
|
||||
v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
|
||||
v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
|
||||
v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
|
||||
|
||||
v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
|
||||
v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
|
||||
v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
|
||||
v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
|
||||
v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
|
||||
v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
|
||||
v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
|
||||
v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
|
||||
v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
|
||||
v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
|
||||
v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
|
||||
v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
|
||||
v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
|
||||
v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
|
||||
v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
|
||||
v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
|
||||
}
|
||||
|
||||
// remove warning : ignoring attributes on template argument ‘__m512i’ [-Wignored-attributes]
|
||||
#pragma GCC diagnostic push
|
||||
#pragma GCC diagnostic ignored "-Wignored-attributes"
|
||||
|
||||
// transpose from [2, 32] to [32, 2]
|
||||
inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) {
|
||||
// r0: {a0, a1, ..., a31}
|
||||
// r1: {b0, b1, ..., b31}
|
||||
//
|
||||
// d0: {a0, b0, ..., a15, b15}
|
||||
// d1: {a16, b16, ..., a31, b31}
|
||||
//
|
||||
__m512i d0 = _mm512_unpacklo_epi16(r0, r1);
|
||||
__m512i d1 = _mm512_unpackhi_epi16(r0, r1);
|
||||
r0 = _mm512_shuffle_i32x4(d0, d1, 0x88);
|
||||
r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd);
|
||||
d0 = _mm512_shuffle_i32x4(r0, r1, 0x88);
|
||||
d1 = _mm512_shuffle_i32x4(r0, r1, 0xdd);
|
||||
return std::make_tuple(d0, d1);
|
||||
}
|
||||
#pragma GCC diagnostic pop
|
||||
|
||||
inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
|
||||
const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
|
||||
const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);
|
||||
const __m512 vec_c2 = _mm512_set1_ps(-0.22433836478672356);
|
||||
const __m512 vec_c3 = _mm512_set1_ps(-0.079204240219773236);
|
||||
|
||||
const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e)
|
||||
|
||||
const __m512 vec_a = _mm512_set1_ps(std::pow(2, 23) / std::log2(2));
|
||||
const __m512 vec_b = _mm512_set1_ps(std::pow(2, 23) * 127.f);
|
||||
|
||||
const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
|
||||
const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
|
||||
__m512i vec_infinity = _mm512_set1_epi32(0x7F800000);
|
||||
__m512i vec_zero = _mm512_setzero_epi32();
|
||||
|
||||
// Fast Exponential Computation on SIMD Architectures
|
||||
// A. Cristiano I. Malossi, Yves Ineichen, Costas Bekas, and Alessandro
|
||||
// Curioni exp(x) = 2**(x * log2(e))
|
||||
// = 2**xi * 2**xf - TIPS we are using the EEEE floating point
|
||||
// representation with identification to the exponent and the
|
||||
// mentissa
|
||||
// 2**xf will be approximated to a polynomial of degree 3 computed with
|
||||
// Horner method
|
||||
// mask for the boundary condition
|
||||
auto min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, _CMP_LT_OS);
|
||||
auto max_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_max, _CMP_GT_OS);
|
||||
|
||||
// transformation with log2(e)
|
||||
auto vec_src = _mm512_mul_ps(values, vec_exp_log2ef);
|
||||
auto vec_fractional = _mm512_sub_ps(vec_src, _mm512_floor_ps(vec_src));
|
||||
|
||||
// compute polynomial using Horner Scheme, for superscalar processor
|
||||
auto vec_res = _mm512_fmadd_ps(vec_fractional, vec_c3, vec_c2);
|
||||
vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c1);
|
||||
vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c0);
|
||||
|
||||
vec_src = _mm512_sub_ps(vec_src, vec_res);
|
||||
// the tips is here, headache in perspective
|
||||
auto tmp = _mm512_fmadd_ps(vec_a, vec_src, vec_b);
|
||||
// headache bis - we loose precision with the cast but it "fits", but ok
|
||||
// after f32 -> f16 later
|
||||
__m512i casted_integer = _mm512_cvttps_epi32(tmp);
|
||||
// boundary condition, lower than the min -> 0
|
||||
casted_integer = _mm512_mask_mov_epi32(casted_integer, min_mask, vec_zero);
|
||||
// boundary condition, larger than the max -> +oo
|
||||
casted_integer = _mm512_mask_mov_epi32(casted_integer, max_mask, vec_infinity);
|
||||
// final interpretation to float
|
||||
return _mm512_castsi512_ps(casted_integer);
|
||||
}
|
||||
#endif
|
||||
|
||||
} // anonymous namespace
|
||||
@@ -0,0 +1,299 @@
|
||||
// Adapted from
|
||||
// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
|
||||
|
||||
// clang-format off
|
||||
|
||||
// To use the transpose functions
|
||||
#include <ATen/native/cpu/utils.h>
|
||||
|
||||
#include "vec.h"
|
||||
|
||||
namespace {
|
||||
|
||||
using namespace at::vec;
|
||||
|
||||
template <typename index_t>
|
||||
inline index_t get_index(index_t* ind, int i) {
|
||||
return (ind == nullptr) ? (index_t)i : ind[i];
|
||||
}
|
||||
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
// key: from [N, 32] to [32/2, N, 2]
|
||||
template <typename scalar_t, typename index_t>
|
||||
inline void pack_vnni_Nx32(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int N,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
__m512i vinputs[16];
|
||||
|
||||
int n = 0;
|
||||
for (; n < N; ++n) {
|
||||
index_t index = get_index(ind, n);
|
||||
vinputs[n] = _mm512_loadu_si512(src + index * ld_src);
|
||||
}
|
||||
// padding with zero to avoid uninitialized vectors
|
||||
for (; n < 16; ++n) {
|
||||
vinputs[n] = _mm512_set1_epi32(0);
|
||||
}
|
||||
|
||||
// pack key
|
||||
transpose_16x16_32bit(vinputs);
|
||||
|
||||
const __mmask16 vmask = (1 << N) - 1;
|
||||
for (int k = 0; k < 16; ++k) {
|
||||
_mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask, vinputs[k]);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename index_t>
|
||||
inline void pack_vnni_N_remainder(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int N,
|
||||
int K,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
__m512i vinputs[16];
|
||||
|
||||
int K2 = K >> 1;
|
||||
const __mmask16 vmask = (1 << K2) - 1;
|
||||
|
||||
int n = 0;
|
||||
for (; n < N; ++n) {
|
||||
index_t index = get_index(ind, n);
|
||||
vinputs[n] = _mm512_maskz_loadu_epi32(vmask, src + index * ld_src);
|
||||
}
|
||||
// padding with zero to avoid uninitialized vectors
|
||||
for (; n < 16; ++n) {
|
||||
vinputs[n] = _mm512_set1_epi32(0);
|
||||
}
|
||||
|
||||
// pack key
|
||||
transpose_16x16_32bit(vinputs);
|
||||
|
||||
const __mmask16 vmask2 = (1 << N) - 1;
|
||||
for (int k = 0; k < K2; ++k) {
|
||||
_mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask2, vinputs[k]);
|
||||
}
|
||||
}
|
||||
|
||||
// value: from [K, 32] to [K/2, 32, 2]
|
||||
template <typename scalar_t, typename index_t>
|
||||
inline void pack_vnni_Kx32(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int K,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
__m512i vinputs[2];
|
||||
|
||||
int k = 0;
|
||||
for (; k < K; ++k) {
|
||||
index_t index = get_index(ind, k);
|
||||
vinputs[k] = _mm512_loadu_si512(src + index * ld_src);
|
||||
}
|
||||
// padding with zero to avoid uninitialized vectors
|
||||
for (; k < 2; ++k) {
|
||||
vinputs[k] = _mm512_set1_epi32(0);
|
||||
}
|
||||
|
||||
// pack value
|
||||
__m512i d0, d1;
|
||||
std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
|
||||
_mm512_storeu_si512(dst + 0 * ld_dst * 2, d0);
|
||||
_mm512_storeu_si512(dst + 0 * ld_dst * 2 + 32, d1);
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename index_t>
|
||||
inline void pack_vnni_K_remainder(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int K,
|
||||
int N,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
__m512i vinputs[2];
|
||||
|
||||
const __mmask32 vmask = (1 << N) - 1;
|
||||
|
||||
int k = 0;
|
||||
for (; k < K; ++k) {
|
||||
index_t index = get_index(ind, k);
|
||||
vinputs[k] = _mm512_maskz_loadu_epi16(vmask, src + index * ld_src);
|
||||
}
|
||||
// padding with zero to avoid uninitialized vectors
|
||||
for (; k < 2; ++k) {
|
||||
vinputs[k] = _mm512_set1_epi32(0);
|
||||
}
|
||||
|
||||
// pack value
|
||||
__m512i d0, d1;
|
||||
std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
|
||||
|
||||
if (N <= 16) {
|
||||
// 2N * 16bits: N * 32bits
|
||||
const __mmask16 vmask2 = (1 << N) - 1;
|
||||
_mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2, vmask2, d0);
|
||||
} else {
|
||||
// 2(N-16) * 16bits: (N-16) * 32bits
|
||||
const __mmask16 vmask2 = (1 << (N - 16)) - 1;
|
||||
_mm512_storeu_epi32(dst + 0 * ld_dst * 2, d0);
|
||||
_mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2 + 32, vmask2, d1);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
// convert to vnni format
|
||||
// from [N, K/2, 2] to [K/2, N, 2] for bfloat16 and float16
|
||||
template <typename scalar_t, typename index_t, bool is_indexed>
|
||||
void pack_vnni(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int N,
|
||||
int K,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
const int NB = div_up(N, 16);
|
||||
const int KB = K / 32;
|
||||
const int K_remainder = K - KB * 32;
|
||||
|
||||
for (int nb = 0; nb < NB; ++nb) {
|
||||
int nb_size = std::min(N - nb * 16, 16);
|
||||
for (int kb = 0; kb < KB; ++kb) {
|
||||
// handle 16x512bits each block
|
||||
pack_vnni_Nx32<scalar_t, index_t>(
|
||||
/* dst */ dst + ((kb * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
|
||||
/* src */ src + kb * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
|
||||
/* ind */ is_indexed ? ind + nb * 16 : nullptr,
|
||||
/* N */ nb_size,
|
||||
/* ld_src */ ld_src,
|
||||
/* ld_dst */ ld_dst);
|
||||
}
|
||||
if (K_remainder > 0) {
|
||||
pack_vnni_N_remainder<scalar_t, index_t>(
|
||||
/* dst */ dst + ((KB * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
|
||||
/* src */ src + KB * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
|
||||
/* ind */ is_indexed ? ind + nb * 16 : nullptr,
|
||||
/* N */ nb_size,
|
||||
/* K */ K_remainder,
|
||||
/* ld_src */ ld_src,
|
||||
/* ld_dst */ ld_dst);
|
||||
}
|
||||
}
|
||||
#else
|
||||
for (int n = 0; n < N; ++n) {
|
||||
index_t index = get_index(ind, n);
|
||||
for (int k = 0; k < K / 2; ++k) {
|
||||
for (int d = 0; d < 2; ++d) {
|
||||
dst[k * ld_dst * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d];
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void pack_vnni(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int N, int K, int ld_src, int ld_dst) {
|
||||
pack_vnni<scalar_t, int32_t, false>(dst, src, nullptr, N, K, ld_src, ld_dst);
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename index_t>
|
||||
void pack_vnni(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int N,
|
||||
int K,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
assert(ind != nullptr);
|
||||
pack_vnni<scalar_t, index_t, true>(dst, src, ind, N, K, ld_src, ld_dst);
|
||||
}
|
||||
|
||||
// convert to vnni format
|
||||
// from [K/2, 2, N] to [K/2, N, 2] for bfloat16 and float16
|
||||
template <typename scalar_t, typename index_t, bool is_indexed>
|
||||
void pack_vnni2(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int K,
|
||||
int N,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
#if defined(CPU_CAPABILITY_AVX512)
|
||||
const int KB = div_up(K, 2);
|
||||
const int NB = N / 32;
|
||||
const int N_remainder = N - NB * 32;
|
||||
|
||||
for (int kb = 0; kb < KB; ++kb) {
|
||||
int kb_size = std::min(K - kb * 2, 2);
|
||||
for (int nb = 0; nb < NB; ++nb) {
|
||||
// handle 2x512bits each block
|
||||
pack_vnni_Kx32<scalar_t, index_t>(
|
||||
/* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + nb * 32 * 2,
|
||||
/* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + nb * 32,
|
||||
/* ind */ is_indexed ? ind + kb * 2 : nullptr,
|
||||
/* K */ kb_size,
|
||||
/* ld_src */ ld_src,
|
||||
/* ld_dst */ ld_dst);
|
||||
}
|
||||
if (N_remainder > 0) {
|
||||
pack_vnni_K_remainder(
|
||||
/* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + NB * 32 * 2,
|
||||
/* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + NB * 32,
|
||||
/* ind */ is_indexed ? ind + kb * 2 : nullptr,
|
||||
/* K */ kb_size,
|
||||
/* N */ N_remainder,
|
||||
/* ld_src */ ld_src,
|
||||
/* ld_dst */ ld_dst);
|
||||
}
|
||||
}
|
||||
#else
|
||||
int k = 0;
|
||||
for (; k < (K >> 1) * 2; k += 2) {
|
||||
index_t index0 = get_index(ind, k + 0);
|
||||
index_t index1 = get_index(ind, k + 1);
|
||||
for (int n = 0; n < N; ++n) {
|
||||
dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = src[index0 * ld_src + n];
|
||||
dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = src[index1 * ld_src + n];
|
||||
}
|
||||
}
|
||||
if (K % 2 != 0) {
|
||||
index_t index = get_index(ind, K - 1);
|
||||
for (int n = 0; n < N; ++n) {
|
||||
dst[(K >> 1) * ld_dst * 2 + n * 2 + 0] = src[index * ld_src + n];
|
||||
dst[(K >> 1) * ld_dst * 2 + n * 2 + 1] = 0;
|
||||
}
|
||||
k += 2;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void pack_vnni2(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int K, int N, int ld_src, int ld_dst) {
|
||||
pack_vnni2<scalar_t, int32_t, false>(dst, src, nullptr, K, N, ld_src, ld_dst);
|
||||
}
|
||||
|
||||
template <typename scalar_t, typename index_t>
|
||||
void pack_vnni2(
|
||||
scalar_t* __restrict__ dst,
|
||||
const scalar_t* __restrict__ src,
|
||||
const index_t* __restrict__ ind,
|
||||
int K,
|
||||
int N,
|
||||
int ld_src,
|
||||
int ld_dst) {
|
||||
assert(ind != nullptr);
|
||||
pack_vnni2<scalar_t, index_t, true>(dst, src, ind, K, N, ld_src, ld_dst);
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
@@ -0,0 +1,869 @@
|
||||
#include "cpu/cpu_types.hpp"
|
||||
|
||||
#include <fcntl.h>
|
||||
#include <sys/mman.h>
|
||||
#include <sys/stat.h>
|
||||
#include <unistd.h>
|
||||
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
#include <atomic>
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
#define MAX_SHM_RANK_NUM 8
|
||||
#define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
|
||||
static_assert(PER_THREAD_SHM_BUFFER_BYTES % 2 == 0);
|
||||
#define PER_THREAD_SHM_BUFFER_OFFSET (PER_THREAD_SHM_BUFFER_BYTES >> 1)
|
||||
#define MIN_THREAD_PROCESS_SIZE (256)
|
||||
#define MAX_P2P_SEND_TENSOR_NUM 8
|
||||
|
||||
template <typename scalar_t>
|
||||
struct KernelVecType {
|
||||
using scalar_vec_t = void;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<float> {
|
||||
using scalar_vec_t = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<c10::BFloat16> {
|
||||
using scalar_vec_t = vec_op::BF16Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct KernelVecType<c10::Half> {
|
||||
using scalar_vec_t = vec_op::FP16Vec16;
|
||||
};
|
||||
|
||||
struct ThreadSHMContext {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
// memory model is weaker on AArch64, so we use atomic variables for
|
||||
// consumer (load-acquire) and producer (store-release) to make sure
|
||||
// that a stamp cannot be ready before the corresponding data is ready.
|
||||
std::atomic<char> _curr_thread_stamp[2];
|
||||
std::atomic<char> _ready_thread_stamp[2];
|
||||
static_assert(std::atomic<char>::is_always_lock_free);
|
||||
#else
|
||||
volatile char _curr_thread_stamp[2];
|
||||
volatile char _ready_thread_stamp[2];
|
||||
#endif // __aarch64__
|
||||
int local_stamp_buffer_idx;
|
||||
int remote_stamp_buffer_idx;
|
||||
int thread_id;
|
||||
int thread_num;
|
||||
int rank;
|
||||
int group_size;
|
||||
size_t _spinning_count;
|
||||
int swizzled_ranks[MAX_SHM_RANK_NUM];
|
||||
void* thread_shm_ptrs[MAX_SHM_RANK_NUM];
|
||||
ThreadSHMContext* shm_contexts[MAX_SHM_RANK_NUM];
|
||||
size_t _thread_buffer_mask[2];
|
||||
char _padding2[40];
|
||||
|
||||
ThreadSHMContext(const int thread_id, const int thread_num, const int rank,
|
||||
const int group_size, void* thread_shm_ptr)
|
||||
: local_stamp_buffer_idx(0),
|
||||
remote_stamp_buffer_idx(0),
|
||||
thread_id(thread_id),
|
||||
thread_num(thread_num),
|
||||
rank(rank),
|
||||
group_size(group_size),
|
||||
_spinning_count(0) {
|
||||
static_assert(sizeof(ThreadSHMContext) % 64 == 0);
|
||||
TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
|
||||
TORCH_CHECK((size_t)this % 64 == 0);
|
||||
TORCH_CHECK((size_t)thread_shm_ptr % 64 == 0);
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_curr_thread_stamp[0].store(1, std::memory_order_relaxed);
|
||||
_curr_thread_stamp[1].store(1, std::memory_order_relaxed);
|
||||
_ready_thread_stamp[0].store(0, std::memory_order_relaxed);
|
||||
_ready_thread_stamp[1].store(0, std::memory_order_relaxed);
|
||||
#else
|
||||
_curr_thread_stamp[0] = 1;
|
||||
_curr_thread_stamp[1] = 1;
|
||||
_ready_thread_stamp[0] = 0;
|
||||
_ready_thread_stamp[1] = 0;
|
||||
#endif // __aarch64__
|
||||
_thread_buffer_mask[0] = 0;
|
||||
_thread_buffer_mask[1] = 0;
|
||||
for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
|
||||
shm_contexts[i] = nullptr;
|
||||
thread_shm_ptrs[i] = nullptr;
|
||||
swizzled_ranks[i] = (i + rank) % group_size;
|
||||
}
|
||||
set_context(rank, this, thread_shm_ptr);
|
||||
}
|
||||
|
||||
void set_stamp_buffer_idx(int local, int remote) {
|
||||
local_stamp_buffer_idx = local;
|
||||
remote_stamp_buffer_idx = remote;
|
||||
}
|
||||
|
||||
void set_context(int rank, ThreadSHMContext* ptr, void* thread_shm_ptr) {
|
||||
TORCH_CHECK(rank < MAX_SHM_RANK_NUM);
|
||||
TORCH_CHECK(ptr);
|
||||
TORCH_CHECK(thread_shm_ptr);
|
||||
TORCH_CHECK_EQ(ptr->thread_num, thread_num);
|
||||
TORCH_CHECK_EQ(ptr->thread_id, thread_id);
|
||||
shm_contexts[rank] = ptr;
|
||||
thread_shm_ptrs[rank] = thread_shm_ptr;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T* get_thread_shm_ptr(int rank) {
|
||||
return reinterpret_cast<T*>(
|
||||
reinterpret_cast<int8_t*>(thread_shm_ptrs[rank]) +
|
||||
(PER_THREAD_SHM_BUFFER_OFFSET &
|
||||
_thread_buffer_mask[local_stamp_buffer_idx]));
|
||||
}
|
||||
|
||||
void next_buffer() {
|
||||
_thread_buffer_mask[local_stamp_buffer_idx] ^= 0xFFFFFFFFFFFFFFFF;
|
||||
}
|
||||
|
||||
char get_curr_stamp(int idx) const {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
return _curr_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _curr_thread_stamp[idx];
|
||||
#endif // __aarch64__
|
||||
}
|
||||
|
||||
char get_ready_stamp(int idx) const {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
return _ready_thread_stamp[idx].load(std::memory_order_acquire);
|
||||
#else
|
||||
return _ready_thread_stamp[idx];
|
||||
#endif // __aarch64__
|
||||
}
|
||||
|
||||
void next_stamp() {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].fetch_add(
|
||||
1, std::memory_order_release);
|
||||
#else
|
||||
_mm_mfence();
|
||||
_curr_thread_stamp[local_stamp_buffer_idx] += 1;
|
||||
#endif // __aarch64__
|
||||
}
|
||||
|
||||
void commit_ready_stamp() {
|
||||
#if defined(__aarch64__) || defined(__powerpc64__)
|
||||
_ready_thread_stamp[local_stamp_buffer_idx].store(
|
||||
_curr_thread_stamp[local_stamp_buffer_idx].load(
|
||||
std::memory_order_relaxed),
|
||||
std::memory_order_release);
|
||||
#else
|
||||
_mm_mfence();
|
||||
_ready_thread_stamp[local_stamp_buffer_idx] =
|
||||
_curr_thread_stamp[local_stamp_buffer_idx];
|
||||
#endif // __aarch64__
|
||||
}
|
||||
|
||||
int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
|
||||
|
||||
template <typename Cond>
|
||||
void wait_for_all(Cond&& cond) {
|
||||
for (int idx = 1; idx < group_size; ++idx) {
|
||||
int rank = get_swizzled_rank(idx);
|
||||
wait_for_one(rank, std::forward<Cond>(cond));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename Cond>
|
||||
void wait_for_one(int rank, Cond&& cond) {
|
||||
ThreadSHMContext* rank_ctx = shm_contexts[rank];
|
||||
for (;;) {
|
||||
char local_curr_stamp = get_curr_stamp(local_stamp_buffer_idx);
|
||||
char local_ready_stamp = get_ready_stamp(local_stamp_buffer_idx);
|
||||
char rank_curr_stamp = rank_ctx->get_curr_stamp(remote_stamp_buffer_idx);
|
||||
char rank_ready_stamp =
|
||||
rank_ctx->get_ready_stamp(remote_stamp_buffer_idx);
|
||||
if (cond(local_curr_stamp, local_ready_stamp, rank_curr_stamp,
|
||||
rank_ready_stamp)) {
|
||||
break;
|
||||
}
|
||||
++_spinning_count;
|
||||
#if defined(__aarch64__)
|
||||
__asm__ __volatile__("yield");
|
||||
#elif defined(__powerpc64__)
|
||||
__asm__ __volatile__("or 1,1,1");
|
||||
#else
|
||||
_mm_pause();
|
||||
#endif // __aarch64__
|
||||
}
|
||||
}
|
||||
|
||||
static bool check_no_buffer_conflict(char local_curr_stamp,
|
||||
char local_ready_stamp,
|
||||
char rank_curr_stamp,
|
||||
char rank_ready_stamp) {
|
||||
char temp = rank_curr_stamp + 2;
|
||||
return local_curr_stamp != temp;
|
||||
}
|
||||
|
||||
static bool check_stamp_ready(char local_curr_stamp, char local_ready_stamp,
|
||||
char rank_curr_stamp, char rank_ready_stamp) {
|
||||
char temp = local_curr_stamp + 1;
|
||||
return (local_curr_stamp == rank_ready_stamp) || (temp == rank_ready_stamp);
|
||||
}
|
||||
|
||||
std::string to_string() const {
|
||||
std::stringstream ss;
|
||||
ss << "SHMContext:";
|
||||
ss << "\nrank: " << rank;
|
||||
ss << "\ngroup_size: " << group_size;
|
||||
ss << "\nthread_num: " << thread_num;
|
||||
ss << "\nthread_id: " << thread_id;
|
||||
|
||||
ss << "\nshm_ctx_stat_loop_seq: [";
|
||||
for (int i = 0; i < group_size; ++i) {
|
||||
ss << swizzled_ranks[i] << ", ";
|
||||
}
|
||||
ss << "]";
|
||||
|
||||
ss << "\nshm_contexts: [";
|
||||
for (int i = 0; i < group_size; ++i) {
|
||||
if (shm_contexts[i]) {
|
||||
ss << shm_contexts[i]->rank << ", ";
|
||||
}
|
||||
}
|
||||
ss << "]";
|
||||
|
||||
return ss.str();
|
||||
}
|
||||
};
|
||||
|
||||
class SHMManager {
|
||||
public:
|
||||
explicit SHMManager(const std::string& name, const int rank,
|
||||
const int group_size, const int thread_num)
|
||||
: _rank(rank),
|
||||
_group_size(group_size),
|
||||
_thread_num(thread_num),
|
||||
_shm_names({""}),
|
||||
_shared_mem_ptrs({nullptr}),
|
||||
_shm_ctx(nullptr) {
|
||||
_shm_names[rank] = get_shm_name(name, rank);
|
||||
_shared_mem_ptrs[rank] = init_shm(rank);
|
||||
_shm_ctx = reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank]);
|
||||
|
||||
for (int i = 0; i < _thread_num; ++i) {
|
||||
ThreadSHMContext* ctx = new (_shm_ctx + i)
|
||||
ThreadSHMContext(i, _thread_num, _rank, _group_size,
|
||||
compute_thread_shm_ptr(_shm_ctx, i));
|
||||
}
|
||||
}
|
||||
|
||||
void join(const std::string& name) {
|
||||
for (int rank_idx = 0; rank_idx < _group_size; ++rank_idx) {
|
||||
if (rank_idx != _rank) {
|
||||
TORCH_CHECK(_shm_names[rank_idx].empty());
|
||||
TORCH_CHECK(_shared_mem_ptrs[rank_idx] == nullptr);
|
||||
_shm_names[rank_idx] = get_shm_name(name, rank_idx);
|
||||
_shared_mem_ptrs[rank_idx] = init_shm(rank_idx);
|
||||
ThreadSHMContext* target_ctx =
|
||||
reinterpret_cast<ThreadSHMContext*>(_shared_mem_ptrs[rank_idx]);
|
||||
for (int thread_idx = 0; thread_idx < _thread_num; ++thread_idx) {
|
||||
_shm_ctx[thread_idx].set_context(
|
||||
rank_idx, target_ctx + thread_idx,
|
||||
compute_thread_shm_ptr(target_ctx, thread_idx));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
~SHMManager() { destroy_shm(); }
|
||||
|
||||
ThreadSHMContext* get_shm_ctx() const { return _shm_ctx; }
|
||||
|
||||
static std::string get_shm_name(const std::string& name, int rank) {
|
||||
return name + "_" + std::to_string(rank);
|
||||
}
|
||||
|
||||
static int64_t create_singleton_instance(const std::string& name,
|
||||
const int group_size, const int rank,
|
||||
const int thread_num) {
|
||||
std::lock_guard<std::mutex> guard(SingletonInstancesLock);
|
||||
SingletonInstances.emplace_back(
|
||||
std::make_unique<SHMManager>(name, rank, group_size, thread_num));
|
||||
return static_cast<int64_t>(SingletonInstances.size() - 1);
|
||||
}
|
||||
|
||||
static SHMManager* get_singleton_instance(int64_t handle) {
|
||||
return SingletonInstances[handle].get();
|
||||
}
|
||||
|
||||
protected:
|
||||
static std::vector<std::unique_ptr<SHMManager>> SingletonInstances;
|
||||
static std::mutex SingletonInstancesLock;
|
||||
|
||||
private:
|
||||
static size_t round_to_alignment(size_t num) {
|
||||
return ((num + 63) / 64) * 64;
|
||||
}
|
||||
|
||||
int8_t* compute_thread_shm_ptr(ThreadSHMContext* ctx, int thread_id) {
|
||||
int8_t* thread_shm_ptr =
|
||||
reinterpret_cast<int8_t*>(ctx) +
|
||||
round_to_alignment(_thread_num * sizeof(ThreadSHMContext));
|
||||
return thread_shm_ptr +
|
||||
thread_id * round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES);
|
||||
}
|
||||
|
||||
size_t compute_shm_size() {
|
||||
const size_t rounded_rank_buffer_size =
|
||||
round_to_alignment(PER_THREAD_SHM_BUFFER_BYTES) * _thread_num;
|
||||
const size_t rounded_thread_shm_ctx_size =
|
||||
round_to_alignment(_thread_num * sizeof(ThreadSHMContext));
|
||||
const size_t shm_size =
|
||||
rounded_thread_shm_ctx_size + rounded_rank_buffer_size;
|
||||
return shm_size;
|
||||
}
|
||||
|
||||
void* init_shm(int target_rank) {
|
||||
const std::string& shm_name = _shm_names[target_rank];
|
||||
const int local_rank = _rank;
|
||||
const size_t shm_size = compute_shm_size();
|
||||
|
||||
int fd = -1;
|
||||
if (local_rank == target_rank) {
|
||||
fd = shm_open(shm_name.c_str(), O_CREAT | O_EXCL | O_RDWR,
|
||||
S_IRUSR | S_IWUSR);
|
||||
|
||||
if (fd == -1)
|
||||
TORCH_CHECK(false, "create shm in SHMManager failed. errno: " +
|
||||
std::to_string(errno));
|
||||
|
||||
if (ftruncate(fd, shm_size) == -1)
|
||||
TORCH_CHECK(false, "ftruncate in SHMManager failed. errno: " +
|
||||
std::to_string(errno));
|
||||
} else {
|
||||
fd = shm_open(shm_name.c_str(), O_RDWR, S_IRUSR | S_IWUSR);
|
||||
|
||||
if (fd == -1)
|
||||
TORCH_CHECK(false, "open shm in SHMManager failed. errno: " +
|
||||
std::to_string(errno));
|
||||
}
|
||||
|
||||
void* shm_ptr = mmap(nullptr, shm_size, PROT_READ | PROT_WRITE,
|
||||
MAP_SHARED | MAP_POPULATE, fd, 0);
|
||||
|
||||
if (shm_ptr == MAP_FAILED) {
|
||||
TORCH_CHECK(false,
|
||||
"mmap in SHMManager failed. errno: " + std::to_string(errno));
|
||||
}
|
||||
|
||||
if (close(fd) != 0) {
|
||||
TORCH_CHECK(
|
||||
false, "close in SHMManager failed. errno: " + std::to_string(errno));
|
||||
}
|
||||
|
||||
TORCH_CHECK((size_t)shm_ptr % 64 == 0);
|
||||
|
||||
return shm_ptr;
|
||||
}
|
||||
|
||||
void destroy_shm() {
|
||||
std::stringstream ss;
|
||||
ss << "local rank " << _rank << ": [";
|
||||
for (int thread_id = 0; thread_id < _thread_num; ++thread_id) {
|
||||
ss << _shm_ctx[thread_id]._spinning_count << ", ";
|
||||
}
|
||||
ss << "]\n";
|
||||
|
||||
for (int i = 0; i < MAX_SHM_RANK_NUM; ++i) {
|
||||
if (_shared_mem_ptrs[i] != nullptr) {
|
||||
munmap(_shared_mem_ptrs[i], compute_shm_size());
|
||||
}
|
||||
|
||||
if (!_shm_names[i].empty()) {
|
||||
shm_unlink(_shm_names[i].c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int _rank;
|
||||
int _group_size;
|
||||
int _thread_num;
|
||||
std::array<std::string, MAX_SHM_RANK_NUM> _shm_names;
|
||||
std::array<void*, MAX_SHM_RANK_NUM> _shared_mem_ptrs;
|
||||
ThreadSHMContext* _shm_ctx;
|
||||
};
|
||||
|
||||
namespace shm_cc_ops {
|
||||
template <typename scalar_t, typename F>
|
||||
void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
|
||||
int thread_num = ctx->thread_num;
|
||||
int64_t total_bytes = elem_num * sizeof(scalar_t);
|
||||
int64_t total_units_num =
|
||||
(total_bytes + MIN_THREAD_PROCESS_SIZE - 1) / MIN_THREAD_PROCESS_SIZE;
|
||||
int64_t per_thread_units_num =
|
||||
(total_units_num + thread_num - 1) / thread_num;
|
||||
int64_t per_unit_elem_num = MIN_THREAD_PROCESS_SIZE / sizeof(scalar_t);
|
||||
int64_t max_per_thread_iteration_elem_num =
|
||||
(PER_THREAD_SHM_BUFFER_BYTES >> 1) /
|
||||
sizeof(scalar_t); // Note: double buffer
|
||||
int64_t per_thread_elem_num = per_unit_elem_num * per_thread_units_num;
|
||||
|
||||
#pragma omp parallel for schedule(static, 1)
|
||||
for (int i = 0; i < thread_num; ++i) {
|
||||
int64_t offset = i * per_thread_elem_num;
|
||||
int64_t end = std::min(elem_num, offset + per_thread_elem_num);
|
||||
int64_t curr_elem_num =
|
||||
std::min(max_per_thread_iteration_elem_num, end - offset);
|
||||
ThreadSHMContext* thread_ctx = ctx + i;
|
||||
bool fast_mode = ((end - offset) <= max_per_thread_iteration_elem_num);
|
||||
|
||||
while (curr_elem_num > 0) {
|
||||
inner_func(thread_ctx, offset, curr_elem_num, fast_mode);
|
||||
|
||||
thread_ctx->next_stamp();
|
||||
thread_ctx->next_buffer();
|
||||
offset += max_per_thread_iteration_elem_num;
|
||||
curr_elem_num = std::min(max_per_thread_iteration_elem_num, end - offset);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void reset_threads_stamp_buffer_idx(ThreadSHMContext* ctx, int local,
|
||||
int remote) {
|
||||
int thread_num = ctx->thread_num;
|
||||
for (int i = 0; i < thread_num; ++i) {
|
||||
ThreadSHMContext* thread_ctx = ctx + i;
|
||||
thread_ctx->set_stamp_buffer_idx(local, remote);
|
||||
}
|
||||
}
|
||||
}; // namespace shm_cc_ops
|
||||
|
||||
namespace shm_cc_ops {
|
||||
|
||||
void memcpy_from_shm(void* dst, void* src, const int64_t bytes) {
|
||||
const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned
|
||||
int64_t i = 0;
|
||||
#pragma GCC unroll 4
|
||||
for (; i < aligned_bytes; i += 64) {
|
||||
vec_op::INT8Vec64 data(
|
||||
true, (int8_t*)src + i); // stream loading shm to avoid caching
|
||||
data.save((int8_t*)dst + i);
|
||||
}
|
||||
if (aligned_bytes < bytes) {
|
||||
vec_op::INT8Vec64 data(true, (int8_t*)src + aligned_bytes);
|
||||
data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
void memcpy_to_shm(void* dst, void* src, const int64_t bytes) {
|
||||
#pragma GCC unroll 4
|
||||
for (int64_t i = 0; i < bytes; i += 64) {
|
||||
vec_op::INT8Vec64 data((int8_t*)src + i);
|
||||
data.nt_save((int8_t*)dst + i);
|
||||
}
|
||||
}
|
||||
|
||||
void memcpy(void* dst, void* src, const int64_t bytes) {
|
||||
const int64_t aligned_bytes = ((bytes >> 6) << 6); // 64 bytes aligned
|
||||
int64_t i = 0;
|
||||
#pragma GCC unroll 4
|
||||
for (; i < aligned_bytes; i += 64) {
|
||||
vec_op::INT8Vec64 data((int8_t*)src + i);
|
||||
data.save((int8_t*)dst + i);
|
||||
}
|
||||
if (aligned_bytes < bytes) {
|
||||
vec_op::INT8Vec64 data((int8_t*)src + aligned_bytes);
|
||||
data.save((int8_t*)dst + aligned_bytes, bytes - aligned_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, int RANKS>
|
||||
void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
|
||||
size_t elem_num) {
|
||||
CPU_KERNEL_GUARD_IN(all_reduce_sum_impl)
|
||||
using vec_t = typename KernelVecType<scalar_t>::scalar_vec_t;
|
||||
constexpr int64_t vec_elem_num = vec_t::get_elem_num();
|
||||
const int worldsize = ctx->group_size;
|
||||
|
||||
shm_cc_ops::shm_cc_loop<scalar_t>(
|
||||
ctx, elem_num,
|
||||
[&](ThreadSHMContext* thread_ctx, int64_t data_offset,
|
||||
int64_t data_elem_num, bool fast_mode) {
|
||||
int rank = thread_ctx->rank;
|
||||
scalar_t* thread_shm_ptr =
|
||||
thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
|
||||
scalar_t* thread_data_ptr = data + data_offset;
|
||||
int64_t thread_data_elem_num = data_elem_num * sizeof(scalar_t);
|
||||
|
||||
scalar_t* remote_data_ptrs[RANKS - 1];
|
||||
vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
|
||||
remote_data_ptrs[idx] = thread_ctx->get_thread_shm_ptr<scalar_t>(
|
||||
thread_ctx->get_swizzled_rank(idx + 1));
|
||||
});
|
||||
|
||||
if (!fast_mode) {
|
||||
thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
|
||||
}
|
||||
|
||||
shm_cc_ops::memcpy_to_shm(thread_shm_ptr, thread_data_ptr,
|
||||
thread_data_elem_num);
|
||||
thread_ctx->commit_ready_stamp();
|
||||
int64_t aligned_data_elem_num =
|
||||
(data_elem_num / vec_elem_num) * vec_elem_num;
|
||||
int64_t i = 0;
|
||||
thread_ctx->wait_for_all(ThreadSHMContext::check_stamp_ready);
|
||||
#pragma GCC unroll 4
|
||||
for (; i < aligned_data_elem_num; i += vec_elem_num) {
|
||||
vec_t local_data(thread_data_ptr + i); // load from cache
|
||||
vec_op::FP32Vec16 local_data_fp32(local_data);
|
||||
vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
|
||||
vec_t remote_data(
|
||||
true, remote_data_ptrs[idx] + i); // stream load from shm
|
||||
vec_op::FP32Vec16 remote_data_fp32(remote_data);
|
||||
local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce
|
||||
});
|
||||
vec_t reduced_data(local_data_fp32);
|
||||
reduced_data.save(thread_data_ptr + i);
|
||||
}
|
||||
|
||||
if (i < data_elem_num) {
|
||||
vec_t local_data(thread_data_ptr + i); // load from cache
|
||||
vec_op::FP32Vec16 local_data_fp32(local_data);
|
||||
vec_op::unroll_loop<int, RANKS - 1>([&](int idx) {
|
||||
vec_t remote_data(
|
||||
true, remote_data_ptrs[idx] + i); // stream load from shm
|
||||
vec_op::FP32Vec16 remote_data_fp32(remote_data);
|
||||
local_data_fp32 = local_data_fp32 + remote_data_fp32; // sum reduce
|
||||
});
|
||||
vec_t reduced_data(local_data_fp32);
|
||||
reduced_data.save(thread_data_ptr + i,
|
||||
data_elem_num - aligned_data_elem_num);
|
||||
}
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
}; // namespace shm_cc_ops
|
||||
|
||||
std::vector<std::unique_ptr<SHMManager>> SHMManager::SingletonInstances = {};
|
||||
std::mutex SHMManager::SingletonInstancesLock = {};
|
||||
|
||||
template <typename scalar_t>
|
||||
void shm_allreduce_sum(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num) {
|
||||
switch (ctx->group_size) {
|
||||
case 2:
|
||||
shm_cc_ops::all_reduce_sum_impl<scalar_t, 2>(ctx, data, elem_num);
|
||||
break;
|
||||
case 3:
|
||||
shm_cc_ops::all_reduce_sum_impl<scalar_t, 3>(ctx, data, elem_num);
|
||||
break;
|
||||
case 4:
|
||||
shm_cc_ops::all_reduce_sum_impl<scalar_t, 4>(ctx, data, elem_num);
|
||||
break;
|
||||
case 8:
|
||||
shm_cc_ops::all_reduce_sum_impl<scalar_t, 8>(ctx, data, elem_num);
|
||||
break;
|
||||
default:
|
||||
TORCH_CHECK(false,
|
||||
"Invalid world size: " + std::to_string(ctx->group_size));
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
|
||||
scalar_t** outputs, const int dst) {
|
||||
CPU_KERNEL_GUARD_IN(shm_gather_impl)
|
||||
const int worldsize = ctx->group_size;
|
||||
TORCH_CHECK_LT(dst, worldsize);
|
||||
shm_cc_ops::shm_cc_loop<scalar_t>(
|
||||
ctx, elem_num,
|
||||
[&](ThreadSHMContext* thread_ctx, int64_t data_offset,
|
||||
int64_t data_elem_num, bool fast_mode) {
|
||||
int rank = thread_ctx->rank;
|
||||
scalar_t* thread_shm_ptr =
|
||||
thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
|
||||
|
||||
if (!fast_mode) {
|
||||
thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
|
||||
}
|
||||
|
||||
shm_cc_ops::memcpy(thread_shm_ptr, data + data_offset,
|
||||
data_elem_num * sizeof(scalar_t));
|
||||
thread_ctx->commit_ready_stamp();
|
||||
if (rank == dst) {
|
||||
shm_cc_ops::memcpy(outputs[rank] + data_offset, data + data_offset,
|
||||
data_elem_num * sizeof(scalar_t));
|
||||
for (int i = 1; i < worldsize; ++i) {
|
||||
int src_rank = thread_ctx->get_swizzled_rank(i);
|
||||
scalar_t* src_ptr =
|
||||
thread_ctx->get_thread_shm_ptr<scalar_t>(src_rank); // shm
|
||||
scalar_t* dst_ptr = outputs[src_rank] + data_offset;
|
||||
thread_ctx->wait_for_one(src_rank,
|
||||
ThreadSHMContext::check_stamp_ready);
|
||||
shm_cc_ops::memcpy(dst_ptr, src_ptr,
|
||||
data_elem_num * sizeof(scalar_t));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
struct MemPiece {
|
||||
void* ptr;
|
||||
int64_t size;
|
||||
|
||||
template <typename T>
|
||||
T* data_ptr() {
|
||||
return reinterpret_cast<T*>(ptr);
|
||||
}
|
||||
};
|
||||
|
||||
struct TensorListMeta {
|
||||
int64_t tensor_bytes[MAX_P2P_SEND_TENSOR_NUM];
|
||||
torch::ScalarType tensor_types[MAX_P2P_SEND_TENSOR_NUM];
|
||||
int64_t tensor_num;
|
||||
int64_t total_bytes;
|
||||
|
||||
TensorListMeta() : tensor_num(0), total_bytes(0) {
|
||||
static_assert(sizeof(TensorListMeta) % 64 == 0);
|
||||
static_assert(sizeof(TensorListMeta) <
|
||||
MIN_THREAD_PROCESS_SIZE); // To ensure the metadata always
|
||||
// hold by the thread 0
|
||||
for (int i = 0; i < MAX_P2P_SEND_TENSOR_NUM; ++i) {
|
||||
tensor_bytes[i] = 0;
|
||||
tensor_ptrs[i] = nullptr;
|
||||
tensor_types[i] = torch::ScalarType::Undefined;
|
||||
}
|
||||
}
|
||||
|
||||
// For send and recv
|
||||
void bind_tensor_list(std::vector<torch::Tensor>& tensor_list) {
|
||||
TORCH_CHECK(tensor_types[0] == torch::ScalarType::Undefined,
|
||||
"Re-bind TensorListMeta is not allowed.")
|
||||
TORCH_CHECK_LE(tensor_list.size(), MAX_P2P_SEND_TENSOR_NUM);
|
||||
tensor_num = tensor_list.size();
|
||||
int64_t bytes_sum = 0;
|
||||
for (int i = 0; i < tensor_list.size(); ++i) {
|
||||
torch::Tensor& t = tensor_list[i];
|
||||
TORCH_CHECK(t.is_contiguous());
|
||||
tensor_bytes[i] = t.nbytes();
|
||||
tensor_types[i] = t.scalar_type();
|
||||
tensor_ptrs[i] = t.data_ptr();
|
||||
bytes_sum += t.nbytes();
|
||||
}
|
||||
total_bytes = bytes_sum;
|
||||
}
|
||||
|
||||
// For recv
|
||||
std::vector<torch::Tensor> generate_tensor_list() {
|
||||
std::vector<torch::Tensor> tensor_list;
|
||||
tensor_list.reserve(tensor_num);
|
||||
|
||||
for (int i = 0; i < tensor_num; ++i) {
|
||||
int64_t bytes = tensor_bytes[i];
|
||||
auto type = tensor_types[i];
|
||||
int64_t elem_bytes = torch::elementSize(type);
|
||||
|
||||
TORCH_CHECK_EQ(bytes % elem_bytes, 0);
|
||||
int64_t elem_num = bytes / elem_bytes;
|
||||
auto options = torch::TensorOptions().dtype(type).device(torch::kCPU);
|
||||
tensor_list.emplace_back(torch::empty({elem_num}, options));
|
||||
}
|
||||
return tensor_list;
|
||||
}
|
||||
|
||||
MemPiece get_data(int64_t offset) {
|
||||
for (int i = 0; i < tensor_num; ++i) {
|
||||
if (offset < tensor_bytes[i]) {
|
||||
return {reinterpret_cast<int8_t*>(tensor_ptrs[i]) + offset,
|
||||
tensor_bytes[i] - offset};
|
||||
}
|
||||
offset -= tensor_bytes[i];
|
||||
}
|
||||
return {nullptr, 0};
|
||||
}
|
||||
|
||||
private:
|
||||
void* tensor_ptrs[MAX_P2P_SEND_TENSOR_NUM];
|
||||
int8_t _padding[40];
|
||||
};
|
||||
|
||||
void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst,
|
||||
const std::vector<torch::Tensor>& tensor_list) {
|
||||
CPU_KERNEL_GUARD_IN(shm_send_tensor_list_impl)
|
||||
std::vector<torch::Tensor> tensor_list_with_metadata;
|
||||
tensor_list_with_metadata.reserve(1 + tensor_list.size());
|
||||
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU);
|
||||
tensor_list_with_metadata.emplace_back(
|
||||
torch::empty({sizeof(TensorListMeta)}, options));
|
||||
tensor_list_with_metadata.insert(tensor_list_with_metadata.end(),
|
||||
tensor_list.begin(), tensor_list.end());
|
||||
|
||||
torch::Tensor& metadata_tensor = tensor_list_with_metadata[0];
|
||||
TORCH_CHECK_EQ(metadata_tensor.nbytes(), sizeof(TensorListMeta));
|
||||
|
||||
TensorListMeta* metadata = new (metadata_tensor.data_ptr()) TensorListMeta();
|
||||
metadata->bind_tensor_list(tensor_list_with_metadata);
|
||||
|
||||
shm_cc_ops::reset_threads_stamp_buffer_idx(ctx, 0, 1);
|
||||
shm_cc_ops::shm_cc_loop<int8_t>(
|
||||
ctx, metadata->total_bytes,
|
||||
[&](ThreadSHMContext* thread_ctx, int64_t data_offset,
|
||||
int64_t data_elem_num, bool fast_mode) {
|
||||
int rank = thread_ctx->rank;
|
||||
int64_t curr_shm_offset = 0;
|
||||
thread_ctx->wait_for_one(dst,
|
||||
ThreadSHMContext::check_no_buffer_conflict);
|
||||
while (curr_shm_offset < data_elem_num) {
|
||||
MemPiece frag = metadata->get_data(data_offset + curr_shm_offset);
|
||||
frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
|
||||
shm_cc_ops::memcpy(
|
||||
thread_ctx->get_thread_shm_ptr<int8_t>(rank) + curr_shm_offset,
|
||||
frag.ptr, frag.size);
|
||||
curr_shm_offset += frag.size;
|
||||
}
|
||||
thread_ctx->commit_ready_stamp();
|
||||
});
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
|
||||
int64_t src) {
|
||||
CPU_KERNEL_GUARD_IN(shm_recv_tensor_list_impl)
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt8).device(torch::kCPU);
|
||||
torch::Tensor metadata_tensor =
|
||||
torch::empty({sizeof(TensorListMeta)}, options);
|
||||
|
||||
shm_cc_ops::reset_threads_stamp_buffer_idx(ctx, 1, 0);
|
||||
ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
|
||||
shm_cc_ops::memcpy(metadata_tensor.data_ptr(),
|
||||
ctx->get_thread_shm_ptr<void>(src),
|
||||
sizeof(TensorListMeta));
|
||||
TensorListMeta* src_metadata =
|
||||
reinterpret_cast<TensorListMeta*>(metadata_tensor.data_ptr());
|
||||
std::vector<torch::Tensor> tensor_list_with_metadata =
|
||||
src_metadata->generate_tensor_list();
|
||||
|
||||
TensorListMeta metadata;
|
||||
metadata.bind_tensor_list(tensor_list_with_metadata);
|
||||
TORCH_CHECK_EQ(metadata.tensor_num, src_metadata->tensor_num);
|
||||
TORCH_CHECK_EQ(metadata.total_bytes, src_metadata->total_bytes);
|
||||
|
||||
shm_cc_ops::shm_cc_loop<int8_t>(
|
||||
ctx, metadata.total_bytes,
|
||||
[&](ThreadSHMContext* thread_ctx, int64_t data_offset,
|
||||
int64_t data_elem_num, bool fast_mode) {
|
||||
thread_ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
|
||||
int64_t curr_shm_offset = 0;
|
||||
while (curr_shm_offset < data_elem_num) {
|
||||
MemPiece frag = metadata.get_data(data_offset + curr_shm_offset);
|
||||
frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
|
||||
shm_cc_ops::memcpy(
|
||||
frag.ptr,
|
||||
thread_ctx->get_thread_shm_ptr<int8_t>(src) + curr_shm_offset,
|
||||
frag.size);
|
||||
curr_shm_offset += frag.size;
|
||||
}
|
||||
});
|
||||
|
||||
std::vector<torch::Tensor> tensor_list;
|
||||
tensor_list.reserve(metadata.tensor_num - 1);
|
||||
tensor_list.insert(tensor_list.begin(), tensor_list_with_metadata.begin() + 1,
|
||||
tensor_list_with_metadata.end());
|
||||
|
||||
return tensor_list;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void shm_gather(int64_t handle, torch::Tensor& data,
|
||||
const std::optional<std::vector<torch::Tensor>>& outputs,
|
||||
int64_t dst) {
|
||||
TORCH_CHECK(data.is_contiguous())
|
||||
VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_gather_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(shm_gather_impl)
|
||||
|
||||
if (outputs.has_value()) {
|
||||
TORCH_CHECK_LE(outputs->size(), MAX_SHM_RANK_NUM);
|
||||
scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr};
|
||||
for (int i = 0; i < outputs->size(); ++i) {
|
||||
output_ptrs[i] = outputs->at(i).data_ptr<scalar_t>();
|
||||
}
|
||||
shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
|
||||
data.data_ptr<scalar_t>(), data.numel(), output_ptrs,
|
||||
dst);
|
||||
} else {
|
||||
shm_gather_impl(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
|
||||
data.data_ptr<scalar_t>(), data.numel(), (scalar_t**)(0),
|
||||
dst);
|
||||
}
|
||||
|
||||
CPU_KERNEL_GUARD_OUT(shm_gather_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void shm_all_gather(int64_t handle, const torch::Tensor& data,
|
||||
torch::Tensor& output) {
|
||||
TORCH_CHECK(data.is_contiguous())
|
||||
TORCH_CHECK(output.is_contiguous())
|
||||
|
||||
const int64_t input_elem_num = data.numel();
|
||||
const int64_t output_elem_num = output.numel();
|
||||
TORCH_CHECK_EQ(output_elem_num % input_elem_num, 0);
|
||||
const int world_size = output_elem_num / input_elem_num;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_all_gather_impl", [&] {
|
||||
CPU_KERNEL_GUARD_IN(shm_all_gather_impl)
|
||||
auto ctx = SHMManager::get_singleton_instance(handle)->get_shm_ctx();
|
||||
TORCH_CHECK_EQ(ctx->group_size, world_size);
|
||||
|
||||
scalar_t* output_ptrs[MAX_SHM_RANK_NUM] = {nullptr};
|
||||
for (int i = 0; i < world_size; ++i) {
|
||||
output_ptrs[i] = output.data_ptr<scalar_t>() + i * input_elem_num;
|
||||
}
|
||||
shm_gather_impl(ctx, data.data_ptr<scalar_t>(), data.numel(), output_ptrs,
|
||||
ctx->rank);
|
||||
CPU_KERNEL_GUARD_OUT(shm_all_gather_impl)
|
||||
});
|
||||
}
|
||||
|
||||
void shm_allreduce(int64_t handle, torch::Tensor& data) {
|
||||
TORCH_CHECK(data.is_contiguous())
|
||||
VLLM_DISPATCH_FLOATING_TYPES(data.scalar_type(), "shm_allreduce_sum", [&] {
|
||||
CPU_KERNEL_GUARD_IN(shm_allreduce_sum)
|
||||
shm_allreduce_sum(SHMManager::get_singleton_instance(handle)->get_shm_ctx(),
|
||||
data.data_ptr<scalar_t>(), data.numel());
|
||||
CPU_KERNEL_GUARD_OUT(shm_allreduce_sum)
|
||||
});
|
||||
}
|
||||
|
||||
void shm_send_tensor_list(int64_t handle,
|
||||
const std::vector<torch::Tensor>& tensor_list,
|
||||
int64_t dst) {
|
||||
CPU_KERNEL_GUARD_IN(shm_send_tensor_list)
|
||||
shm_send_tensor_list_impl(
|
||||
SHMManager::get_singleton_instance(handle)->get_shm_ctx(), dst,
|
||||
tensor_list);
|
||||
CPU_KERNEL_GUARD_OUT(shm_send_tensor_list)
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src) {
|
||||
CPU_KERNEL_GUARD_IN(shm_recv_tensor_list)
|
||||
auto tensor_list = shm_recv_tensor_list_impl(
|
||||
SHMManager::get_singleton_instance(handle)->get_shm_ctx(), src);
|
||||
CPU_KERNEL_GUARD_OUT(shm_recv_tensor_list)
|
||||
return tensor_list;
|
||||
}
|
||||
|
||||
int64_t init_shm_manager(const std::string& name, const int64_t group_size,
|
||||
const int64_t rank, const int64_t thread_num) {
|
||||
return SHMManager::create_singleton_instance(name, group_size, rank,
|
||||
thread_num);
|
||||
}
|
||||
|
||||
std::string join_shm_manager(int64_t handle, const std::string& name) {
|
||||
auto shm_manager = SHMManager::get_singleton_instance(handle);
|
||||
TORCH_CHECK(shm_manager);
|
||||
shm_manager->join(name);
|
||||
return shm_manager->get_shm_ctx()->to_string();
|
||||
}
|
||||
@@ -0,0 +1,492 @@
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
namespace cpu_utils {
|
||||
|
||||
void eagle_prepare_inputs_padded_kernel_impl(
|
||||
const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& valid_sampled_tokens_count,
|
||||
const torch::Tensor& query_start_loc_gpu,
|
||||
torch::Tensor& token_indices_to_sample,
|
||||
torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) {
|
||||
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
|
||||
const int64_t* valid_count_ptr =
|
||||
valid_sampled_tokens_count.data_ptr<int64_t>();
|
||||
const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr<int32_t>();
|
||||
int32_t* indices_out_ptr = token_indices_to_sample.data_ptr<int32_t>();
|
||||
int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr<int64_t>();
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
|
||||
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
|
||||
int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx;
|
||||
int64_t num_valid_tokens = valid_count_ptr[req_idx];
|
||||
|
||||
int64_t num_rejected = 0;
|
||||
if (num_draft_tokens > 0) {
|
||||
num_rejected = num_draft_tokens + 1 - num_valid_tokens;
|
||||
}
|
||||
|
||||
int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1;
|
||||
int32_t index_to_sample = q_last_tok_idx - num_rejected;
|
||||
|
||||
indices_out_ptr[req_idx] = index_to_sample;
|
||||
rejected_out_ptr[req_idx] = num_rejected;
|
||||
}
|
||||
}
|
||||
|
||||
void eagle_prepare_next_token_padded_kernel_impl(
|
||||
const torch::Tensor& sampled_token_ids,
|
||||
const torch::Tensor& discard_request_mask,
|
||||
const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
|
||||
torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
|
||||
const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) {
|
||||
const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr<int64_t>();
|
||||
const bool* discard_mask_ptr = discard_request_mask.data_ptr<bool>();
|
||||
const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr<int64_t>();
|
||||
int64_t* next_ids_out_ptr = next_token_ids.data_ptr<int64_t>();
|
||||
int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr<int64_t>();
|
||||
|
||||
const int64_t stride = sampled_token_ids.stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
|
||||
const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride;
|
||||
int64_t valid_count = 0;
|
||||
int64_t last_valid_token = -1;
|
||||
|
||||
for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) {
|
||||
int64_t token = row_ptr[pos];
|
||||
if (token != -1 && token < vocab_size) {
|
||||
valid_count++;
|
||||
last_valid_token = token;
|
||||
}
|
||||
}
|
||||
|
||||
bool discard = discard_mask_ptr[req_idx];
|
||||
if (discard) {
|
||||
next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx];
|
||||
valid_count_out_ptr[req_idx] = 0;
|
||||
} else {
|
||||
next_ids_out_ptr[req_idx] =
|
||||
(valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx];
|
||||
valid_count_out_ptr[req_idx] = valid_count;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void eagle_step_slot_mapping_metadata_kernel_impl(
|
||||
const torch::Tensor& positions, const torch::Tensor& block_table,
|
||||
torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
|
||||
torch::Tensor& out_slot_mapping, const int64_t block_size,
|
||||
const int64_t max_model_len, const int64_t PAD_ID) {
|
||||
const int64_t batch_size = positions.size(0);
|
||||
const int64_t input_batch_size = out_slot_mapping.size(0);
|
||||
|
||||
const int64_t* pos_ptr = positions.data_ptr<int64_t>();
|
||||
const int32_t* bt_ptr = block_table.data_ptr<int32_t>();
|
||||
int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>();
|
||||
int64_t* out_clamped_ptr = out_clamped_positions.data_ptr<int64_t>();
|
||||
int64_t* out_slot_ptr = out_slot_mapping.data_ptr<int64_t>();
|
||||
|
||||
const int64_t bt_stride = block_table.stride(0);
|
||||
const int64_t n_blocks_per_req = block_table.size(1);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) {
|
||||
if (req_idx >= batch_size) {
|
||||
out_slot_ptr[req_idx] = PAD_ID;
|
||||
continue;
|
||||
}
|
||||
|
||||
int64_t position = pos_ptr[req_idx];
|
||||
int64_t new_position = position + 1;
|
||||
bool exceeds_max = new_position >= max_model_len;
|
||||
int64_t clamped_position = exceeds_max ? 0 : new_position;
|
||||
|
||||
out_clamped_ptr[req_idx] = clamped_position;
|
||||
|
||||
int64_t block_number = clamped_position / block_size;
|
||||
block_number = std::min(block_number, n_blocks_per_req - 1);
|
||||
int32_t block_id = bt_ptr[req_idx * bt_stride + block_number];
|
||||
int64_t slot_id = block_id * block_size + (clamped_position % block_size);
|
||||
out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id;
|
||||
|
||||
int32_t seq_len = seq_lens_ptr[req_idx];
|
||||
int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1);
|
||||
new_seq_len = std::min(new_seq_len, static_cast<int32_t>(max_model_len));
|
||||
seq_lens_ptr[req_idx] = new_seq_len;
|
||||
}
|
||||
}
|
||||
|
||||
void copy_and_expand_eagle_inputs_kernel_impl(
|
||||
const torch::Tensor& target_token_ids,
|
||||
const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_positions,
|
||||
torch::Tensor& out_is_rejected_token_mask,
|
||||
torch::Tensor& out_is_masked_token_mask,
|
||||
torch::Tensor& out_new_token_indices,
|
||||
torch::Tensor& out_hidden_state_mapping,
|
||||
const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
|
||||
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
|
||||
const int64_t total_input_tokens,
|
||||
const int64_t num_padding_slots_per_request, const bool shift_input_ids) {
|
||||
const int64_t num_reqs = query_end_loc.size(0);
|
||||
|
||||
const int64_t* target_ids_ptr = target_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
|
||||
const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
|
||||
const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
|
||||
const int32_t* query_end_ptr = query_end_loc.data_ptr<int32_t>();
|
||||
|
||||
int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
|
||||
int64_t* out_pos_ptr = out_positions.data_ptr<int64_t>();
|
||||
bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr<bool>();
|
||||
bool* out_mask_ptr = out_is_masked_token_mask.data_ptr<bool>();
|
||||
int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr<int32_t>();
|
||||
int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr<int32_t>();
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
|
||||
int32_t q_start = query_start_ptr[req_idx];
|
||||
int32_t next_q_start = query_start_ptr[req_idx + 1];
|
||||
int32_t q_end = query_end_ptr[req_idx];
|
||||
|
||||
int64_t num_valid_tokens =
|
||||
shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1);
|
||||
int64_t input_offset = shift_input_ids ? 1 : 0;
|
||||
|
||||
int64_t out_start = q_start + req_idx * (num_padding_slots_per_request -
|
||||
(shift_input_ids ? 1 : 0));
|
||||
int64_t num_rejected = next_q_start - q_end - 1;
|
||||
int64_t total_output_tokens =
|
||||
num_valid_tokens + num_padding_slots_per_request + num_rejected;
|
||||
|
||||
int64_t start_pos = target_pos_ptr[q_start];
|
||||
int64_t bonus_token = next_ids_ptr[req_idx];
|
||||
|
||||
for (int64_t j = 0; j < total_output_tokens; ++j) {
|
||||
int64_t out_idx = out_start + j;
|
||||
bool is_valid = j < num_valid_tokens;
|
||||
bool is_bonus = j == num_valid_tokens;
|
||||
bool is_parallel = (j > num_valid_tokens) &&
|
||||
(j < num_valid_tokens + num_padding_slots_per_request);
|
||||
bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request;
|
||||
|
||||
int64_t in_idx =
|
||||
std::min(static_cast<int64_t>(q_start + input_offset + j),
|
||||
total_input_tokens - 1);
|
||||
|
||||
int64_t token_id = padding_token_id;
|
||||
if (is_valid)
|
||||
token_id = target_ids_ptr[in_idx];
|
||||
else if (is_bonus)
|
||||
token_id = bonus_token;
|
||||
else if (is_parallel)
|
||||
token_id = parallel_drafting_token_id;
|
||||
|
||||
out_ids_ptr[out_idx] = token_id;
|
||||
out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j);
|
||||
out_rej_mask_ptr[out_idx] = is_rejected;
|
||||
out_mask_ptr[out_idx] = is_parallel;
|
||||
|
||||
if (is_bonus || is_parallel) {
|
||||
int64_t new_token_local_idx = j - num_valid_tokens;
|
||||
int64_t new_token_out_idx =
|
||||
req_idx * num_padding_slots_per_request + new_token_local_idx;
|
||||
out_new_idx_ptr[new_token_out_idx] = out_idx;
|
||||
}
|
||||
}
|
||||
|
||||
if (shift_input_ids) {
|
||||
int64_t n_input = next_q_start - q_start;
|
||||
for (int64_t j = 0; j < n_input; ++j) {
|
||||
out_hidden_map_ptr[q_start + j] = out_start + j;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void copy_and_expand_dflash_inputs_kernel_impl(
|
||||
const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
|
||||
torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
|
||||
torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
|
||||
const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
|
||||
const std::optional<torch::Tensor>& num_rejected_tokens,
|
||||
const int64_t parallel_drafting_token_id, const int64_t block_size,
|
||||
const int64_t num_query_per_req, const int64_t num_speculative_tokens,
|
||||
const int64_t total_input_tokens, const bool has_num_rejected) {
|
||||
const int64_t num_reqs = query_start_loc.size(0) - 1;
|
||||
|
||||
const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>();
|
||||
const int32_t* block_table_ptr = block_table.data_ptr<int32_t>();
|
||||
const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>();
|
||||
const int64_t* rejected_ptr =
|
||||
has_num_rejected && num_rejected_tokens.has_value()
|
||||
? num_rejected_tokens.value().data_ptr<int64_t>()
|
||||
: nullptr;
|
||||
|
||||
int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>();
|
||||
int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>();
|
||||
int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>();
|
||||
int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>();
|
||||
int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>();
|
||||
int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>();
|
||||
|
||||
const int64_t block_table_stride = block_table.stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) {
|
||||
int32_t ctx_start = query_start_ptr[req_idx];
|
||||
int32_t ctx_end = query_start_ptr[req_idx + 1];
|
||||
int64_t num_ctx = ctx_end - ctx_start;
|
||||
int64_t valid_ctx_end = ctx_end;
|
||||
if (rejected_ptr != nullptr) {
|
||||
valid_ctx_end -= rejected_ptr[req_idx];
|
||||
}
|
||||
// Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that
|
||||
// valid_ctx_end - 1 never reads before the request's context range.
|
||||
valid_ctx_end =
|
||||
std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1));
|
||||
|
||||
int64_t last_pos = target_pos_ptr[valid_ctx_end - 1];
|
||||
|
||||
for (int64_t j = 0; j < num_ctx; ++j) {
|
||||
int64_t ctx_idx = ctx_start + j;
|
||||
int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1);
|
||||
int64_t position = target_pos_ptr[ctx_pos_idx];
|
||||
int64_t block_num = position / block_size;
|
||||
block_num = std::min(block_num, block_table_stride - 1);
|
||||
int32_t block_id =
|
||||
block_table_ptr[req_idx * block_table_stride + block_num];
|
||||
int64_t slot = block_id * block_size + (position % block_size);
|
||||
|
||||
out_ctx_pos_ptr[ctx_idx] = position;
|
||||
out_ctx_slot_ptr[ctx_idx] = slot;
|
||||
}
|
||||
|
||||
for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) {
|
||||
int64_t query_out = req_idx * num_query_per_req + query_off;
|
||||
int64_t position = last_pos + 1 + query_off;
|
||||
int64_t block_num = position / block_size;
|
||||
block_num = std::min(block_num, block_table_stride - 1);
|
||||
int32_t block_id =
|
||||
block_table_ptr[req_idx * block_table_stride + block_num];
|
||||
int64_t slot = block_id * block_size + (position % block_size);
|
||||
|
||||
out_query_pos_ptr[query_out] = position;
|
||||
out_query_slot_ptr[query_out] = slot;
|
||||
out_ids_ptr[query_out] =
|
||||
query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id;
|
||||
|
||||
if (query_off > 0) {
|
||||
int64_t sample_out_idx =
|
||||
req_idx * num_speculative_tokens + (query_off - 1);
|
||||
out_token_idx_ptr[sample_out_idx] = query_out;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void rejection_greedy_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
|
||||
const torch::Tensor& bonus_token_ids,
|
||||
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len) {
|
||||
const int64_t batch_size = cu_num_draft_tokens.size(0);
|
||||
|
||||
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
|
||||
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* target_argmax_ptr = target_argmax.data_ptr<int64_t>();
|
||||
const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
|
||||
const bool* greedy_ptr =
|
||||
is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
|
||||
|
||||
const int64_t out_stride = output_token_ids.stride(0);
|
||||
const int64_t bonus_stride = bonus_token_ids.stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
|
||||
if (greedy_ptr && !greedy_ptr[req_idx]) continue;
|
||||
|
||||
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
|
||||
int64_t end_idx = cu_draft_ptr[req_idx];
|
||||
int64_t num_draft_tokens = end_idx - start_idx;
|
||||
|
||||
bool rejected = false;
|
||||
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
|
||||
int64_t target_id = target_argmax_ptr[start_idx + pos];
|
||||
out_ptr[req_idx * out_stride + pos] = target_id;
|
||||
|
||||
if (draft_ids_ptr[start_idx + pos] != target_id) {
|
||||
rejected = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!rejected) {
|
||||
out_ptr[req_idx * out_stride + num_draft_tokens] =
|
||||
bonus_ids_ptr[req_idx * bonus_stride];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void rejection_random_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids,
|
||||
const std::optional<torch::Tensor>& draft_probs,
|
||||
const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
|
||||
const torch::Tensor& recovered_token_ids,
|
||||
const torch::Tensor& uniform_probs,
|
||||
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
|
||||
const int64_t vocab_size, const bool no_draft_probs) {
|
||||
const int64_t batch_size = cu_num_draft_tokens.size(0);
|
||||
|
||||
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
|
||||
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
|
||||
const float* draft_probs_ptr =
|
||||
no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
|
||||
const float* target_probs_ptr = target_probs.data_ptr<float>();
|
||||
const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr<int64_t>();
|
||||
const float* uniform_probs_ptr = uniform_probs.data_ptr<float>();
|
||||
const bool* greedy_ptr =
|
||||
is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr;
|
||||
|
||||
const int64_t out_stride = output_token_ids.stride(0);
|
||||
const int64_t bonus_stride = bonus_token_ids.stride(0);
|
||||
const int64_t target_stride = target_probs.stride(0);
|
||||
const int64_t draft_probs_stride =
|
||||
no_draft_probs ? 0 : draft_probs.value().stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
|
||||
if (greedy_ptr && greedy_ptr[req_idx]) continue;
|
||||
|
||||
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
|
||||
int64_t end_idx = cu_draft_ptr[req_idx];
|
||||
int64_t num_draft_tokens = end_idx - start_idx;
|
||||
|
||||
bool rejected = false;
|
||||
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
|
||||
int64_t token_idx = start_idx + pos;
|
||||
int64_t draft_id = draft_ids_ptr[token_idx];
|
||||
|
||||
float p = target_probs_ptr[token_idx * target_stride + draft_id];
|
||||
float q =
|
||||
no_draft_probs
|
||||
? 1.0f
|
||||
: draft_probs_ptr[token_idx * draft_probs_stride + draft_id];
|
||||
float uniform_p = uniform_probs_ptr[token_idx];
|
||||
|
||||
float ratio = (q > 0.0f) ? (p / q) : 0.0f;
|
||||
|
||||
if (ratio >= uniform_p) {
|
||||
out_ptr[req_idx * out_stride + pos] = draft_id;
|
||||
} else {
|
||||
out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx];
|
||||
rejected = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!rejected) {
|
||||
out_ptr[req_idx * out_stride + num_draft_tokens] =
|
||||
bonus_ids_ptr[req_idx * bonus_stride];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& cu_num_tokens,
|
||||
const int64_t replace_from, const int64_t replace_to) {
|
||||
const int64_t batch_size = cu_num_tokens.size(0);
|
||||
const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr<int64_t>();
|
||||
|
||||
int64_t* out_ptr = output.data_ptr<int64_t>();
|
||||
const int64_t* in_ptr = input.data_ptr<int64_t>();
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
|
||||
int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1];
|
||||
int64_t end_idx = cu_tokens_ptr[req_idx];
|
||||
int64_t val = in_ptr[req_idx];
|
||||
|
||||
if (val == replace_from) {
|
||||
val = replace_to;
|
||||
}
|
||||
|
||||
for (int64_t i = start_idx; i < end_idx; ++i) {
|
||||
out_ptr[i] = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void sample_recovered_tokens_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids,
|
||||
const std::optional<torch::Tensor>& draft_probs,
|
||||
const torch::Tensor& target_probs, const torch::Tensor& inv_q,
|
||||
const int64_t vocab_size, const bool no_draft_probs) {
|
||||
const int64_t batch_size = cu_num_draft_tokens.size(0);
|
||||
|
||||
int64_t* out_ptr = output_token_ids.data_ptr<int64_t>();
|
||||
const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>();
|
||||
const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>();
|
||||
const float* draft_probs_ptr =
|
||||
no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>();
|
||||
const float* target_probs_ptr = target_probs.data_ptr<float>();
|
||||
const float* inv_q_ptr = inv_q.data_ptr<float>();
|
||||
|
||||
const int64_t target_stride = target_probs.stride(0);
|
||||
const int64_t draft_probs_stride =
|
||||
no_draft_probs ? 0 : draft_probs.value().stride(0);
|
||||
const int64_t inv_q_stride = inv_q.stride(0);
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) {
|
||||
int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1];
|
||||
int64_t end_idx = cu_draft_ptr[req_idx];
|
||||
int64_t num_draft_tokens = end_idx - start_idx;
|
||||
|
||||
const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride;
|
||||
|
||||
for (int64_t pos = 0; pos < num_draft_tokens; ++pos) {
|
||||
int64_t token_idx = start_idx + pos;
|
||||
int64_t draft_id = draft_ids_ptr[token_idx];
|
||||
|
||||
const float* token_target_probs =
|
||||
target_probs_ptr + token_idx * target_stride;
|
||||
const float* token_draft_probs =
|
||||
no_draft_probs ? nullptr
|
||||
: (draft_probs_ptr + token_idx * draft_probs_stride);
|
||||
|
||||
int64_t best_id = 0;
|
||||
float best_val = -1.0f;
|
||||
|
||||
for (int64_t v = 0; v < vocab_size; ++v) {
|
||||
float prob = token_target_probs[v];
|
||||
if (no_draft_probs) {
|
||||
if (v == draft_id) prob = 0.0f;
|
||||
} else {
|
||||
float diff = prob - token_draft_probs[v];
|
||||
prob = diff > 0.0f ? diff : 0.0f;
|
||||
}
|
||||
|
||||
float val = prob * req_inv_q[v];
|
||||
if (val > best_val) {
|
||||
best_val = val;
|
||||
best_id = v;
|
||||
}
|
||||
}
|
||||
out_ptr[token_idx] = best_id;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace cpu_utils
|
||||
@@ -0,0 +1,677 @@
|
||||
#include "cache.h"
|
||||
#include "ops.h"
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <torch/library.h>
|
||||
|
||||
// Note: overwrite the external definition for sharing same name between
|
||||
// libraries use different ISAs.
|
||||
#define TORCH_EXTENSION_NAME _C
|
||||
|
||||
void release_dnnl_matmul_handler(int64_t handler);
|
||||
|
||||
int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b,
|
||||
const torch::Tensor& b_scales,
|
||||
at::ScalarType output_type,
|
||||
bool dynamic_act_quant, bool use_azp,
|
||||
int64_t primitive_cache_size);
|
||||
|
||||
void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
|
||||
const torch::Tensor& a_scales,
|
||||
const std::optional<torch::Tensor>& azp,
|
||||
const std::optional<torch::Tensor>& azp_adj,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const torch::Tensor& handler_tensor);
|
||||
|
||||
int64_t create_onednn_mm_handler(const torch::Tensor& b,
|
||||
int64_t primitive_cache_size);
|
||||
|
||||
void onednn_mm(torch::Tensor& c, const torch::Tensor& a,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const torch::Tensor& handler_tensor);
|
||||
|
||||
bool is_onednn_acl_supported();
|
||||
|
||||
void mla_decode_kvcache(torch::Tensor& out, torch::Tensor& query,
|
||||
torch::Tensor& kv_cache, double scale,
|
||||
torch::Tensor& block_tables, torch::Tensor& seq_lens);
|
||||
|
||||
int64_t init_shm_manager(const std::string& name, const int64_t group_size,
|
||||
const int64_t rank, const int64_t thread_num);
|
||||
|
||||
std::string join_shm_manager(int64_t handle, const std::string& name);
|
||||
|
||||
void shm_allreduce(int64_t handle, torch::Tensor& data);
|
||||
|
||||
void shm_gather(int64_t handle, torch::Tensor& data,
|
||||
const std::optional<std::vector<torch::Tensor>>& outputs,
|
||||
int64_t dst);
|
||||
|
||||
void shm_all_gather(int64_t handle, const torch::Tensor& data,
|
||||
torch::Tensor& output);
|
||||
|
||||
void shm_send_tensor_list(int64_t handle,
|
||||
const std::vector<torch::Tensor>& tensor_list,
|
||||
int64_t dst);
|
||||
|
||||
std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
|
||||
|
||||
// SGL CPU kernels
|
||||
|
||||
at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
bool is_vnni);
|
||||
|
||||
at::Tensor convert_weight_packed(at::Tensor& weight);
|
||||
|
||||
at::Tensor convert_scale_packed(at::Tensor& scale);
|
||||
|
||||
at::Tensor fused_experts_cpu(
|
||||
at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
|
||||
at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
|
||||
int64_t moe_comp_method, const std::optional<at::Tensor>& w1_scale,
|
||||
const std::optional<at::Tensor>& w2_scale,
|
||||
const std::optional<at::Tensor>& w1_zero,
|
||||
const std::optional<at::Tensor>& w2_zero,
|
||||
const std::optional<std::vector<int64_t>> block_size,
|
||||
const std::optional<at::Tensor>& w1_bias,
|
||||
const std::optional<at::Tensor>& w2_bias,
|
||||
const std::optional<double>& alpha, const std::optional<double>& limit,
|
||||
bool is_vnni);
|
||||
|
||||
at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
|
||||
at::Tensor& scales2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni);
|
||||
|
||||
// Adapted from sglang: FP8 W8A16 kernel
|
||||
at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
|
||||
at::Tensor& scales2,
|
||||
std::vector<int64_t> block_size,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni);
|
||||
|
||||
// Adapted from sglang: INT4 W4A8 kernels
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
|
||||
at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
|
||||
at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K /
|
||||
// group_size, N / 8) , int32
|
||||
at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K /
|
||||
// group_size, N) , bfloat16
|
||||
int64_t quant_method_4bit);
|
||||
|
||||
at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
|
||||
at::Tensor& w_scales,
|
||||
std::optional<at::Tensor> bias);
|
||||
|
||||
// Adapted from sglang: GDN
|
||||
std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
|
||||
const at::Tensor& query, const at::Tensor& key, const at::Tensor& value,
|
||||
const at::Tensor& g, const at::Tensor& beta,
|
||||
const at::Tensor& initial_state, bool output_final_state,
|
||||
const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel,
|
||||
double eps = 1e-5);
|
||||
|
||||
at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
|
||||
const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
|
||||
const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,
|
||||
const at::Tensor& b, at::Tensor& initial_state_source,
|
||||
const at::Tensor& initial_state_indices, const at::Tensor& cu_seqlens,
|
||||
bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,
|
||||
double softplus_threshold = 20.0);
|
||||
|
||||
at::Tensor fused_sigmoid_gating_delta_rule_update_spec_cpu(
|
||||
const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,
|
||||
const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,
|
||||
const at::Tensor& b, at::Tensor& initial_state_source,
|
||||
const at::Tensor& spec_state_indices, const at::Tensor& num_accepted_tokens,
|
||||
const at::Tensor& cu_seqlens, bool use_qk_l2norm_in_kernel,
|
||||
double softplus_beta = 1.0, double softplus_threshold = 20.0);
|
||||
|
||||
std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(
|
||||
const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,
|
||||
const at::Tensor& dt_bias);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight);
|
||||
|
||||
at::Tensor causal_conv1d_fwd_cpu(
|
||||
const at::Tensor& x, const at::Tensor& weight,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
const std::optional<at::Tensor>& conv_states,
|
||||
const std::optional<at::Tensor>& query_start_loc,
|
||||
const std::optional<at::Tensor>& cache_indices,
|
||||
const std::optional<at::Tensor>& has_initial_state, bool silu_activation,
|
||||
int64_t pad_slot_id, bool is_vnni);
|
||||
|
||||
at::Tensor causal_conv1d_update_cpu(
|
||||
const at::Tensor& x, const at::Tensor& conv_states,
|
||||
const at::Tensor& weight, const std::optional<at::Tensor>& bias,
|
||||
bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,
|
||||
const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
|
||||
bool is_vnni);
|
||||
|
||||
void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
|
||||
const std::string& activation);
|
||||
|
||||
bool cpu_attn_has_isa(const std::string& isa);
|
||||
|
||||
torch::Tensor get_scheduler_metadata(
|
||||
const int64_t num_req, const int64_t num_heads_q,
|
||||
const int64_t num_heads_kv, const int64_t head_dim,
|
||||
const torch::Tensor& seq_lens, at::ScalarType dtype,
|
||||
const torch::Tensor& query_start_loc, const bool casual,
|
||||
const int64_t window_size, const std::string& isa_hint,
|
||||
const bool enable_kv_split,
|
||||
const std::optional<torch::Tensor>& dynamic_causal);
|
||||
|
||||
void cpu_attn_reshape_and_cache(const torch::Tensor& key,
|
||||
const torch::Tensor& value,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
const torch::Tensor& slot_mapping,
|
||||
const std::string& isa, const double k_scale,
|
||||
const double v_scale,
|
||||
const std::string& kv_cache_dtype);
|
||||
|
||||
void cpu_attention_with_kv_cache(
|
||||
const torch::Tensor& query, const torch::Tensor& key_cache,
|
||||
const torch::Tensor& value_cache, torch::Tensor& output,
|
||||
const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
|
||||
const double scale, const bool causal,
|
||||
const std::optional<torch::Tensor>& alibi_slopes,
|
||||
const int64_t sliding_window_left, const torch::Tensor& block_table,
|
||||
const double softcap, const torch::Tensor& scheduler_metadata,
|
||||
const std::optional<torch::Tensor>& s_aux,
|
||||
const std::optional<torch::Tensor>& dynamic_causal, const double k_scale,
|
||||
const double v_scale, const std::string& kv_cache_dtype);
|
||||
|
||||
// Note: just for avoiding importing errors
|
||||
void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
|
||||
|
||||
void cpu_gemm_wna16(const torch::Tensor& input, const torch::Tensor& q_weight,
|
||||
torch::Tensor& output, const torch::Tensor& scales,
|
||||
const std::optional<torch::Tensor>& zeros,
|
||||
const std::optional<torch::Tensor>& g_idx,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const int64_t pack_factor, const std::string& isa_hint);
|
||||
|
||||
void prepack_moe_weight(const torch::Tensor& weight,
|
||||
torch::Tensor& packed_weight, const std::string& isa);
|
||||
|
||||
void cpu_fused_moe(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& w13, const torch::Tensor& w2,
|
||||
const std::optional<torch::Tensor>& w13_bias,
|
||||
const std::optional<torch::Tensor>& w2_bias,
|
||||
const torch::Tensor& topk_weights,
|
||||
const torch::Tensor& topk_id, const bool skip_weighted,
|
||||
const std::string& act, const std::string& isa);
|
||||
|
||||
void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
const torch::Tensor positions,
|
||||
const torch::Tensor block_table,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
void eagle_prepare_inputs_padded_kernel_impl(
|
||||
const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& valid_sampled_tokens_count,
|
||||
const torch::Tensor& query_start_loc_gpu,
|
||||
torch::Tensor& token_indices_to_sample,
|
||||
torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs);
|
||||
void eagle_prepare_next_token_padded_kernel_impl(
|
||||
const torch::Tensor& sampled_token_ids,
|
||||
const torch::Tensor& discard_request_mask,
|
||||
const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids,
|
||||
torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size,
|
||||
const int64_t num_sampled_tokens_per_req, const int64_t num_reqs);
|
||||
void eagle_step_slot_mapping_metadata_kernel_impl(
|
||||
const torch::Tensor& positions, const torch::Tensor& block_table,
|
||||
torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions,
|
||||
torch::Tensor& out_slot_mapping, const int64_t block_size,
|
||||
const int64_t max_model_len, const int64_t PAD_ID);
|
||||
void copy_and_expand_eagle_inputs_kernel_impl(
|
||||
const torch::Tensor& target_token_ids,
|
||||
const torch::Tensor& target_positions, const torch::Tensor& next_token_ids,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_positions,
|
||||
torch::Tensor& out_is_rejected_token_mask,
|
||||
torch::Tensor& out_is_masked_token_mask,
|
||||
torch::Tensor& out_new_token_indices,
|
||||
torch::Tensor& out_hidden_state_mapping,
|
||||
const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc,
|
||||
const int64_t padding_token_id, const int64_t parallel_drafting_token_id,
|
||||
const int64_t total_input_tokens,
|
||||
const int64_t num_padding_slots_per_request, const bool shift_input_ids);
|
||||
void copy_and_expand_dflash_inputs_kernel_impl(
|
||||
const torch::Tensor& next_token_ids, const torch::Tensor& target_positions,
|
||||
torch::Tensor& out_input_ids, torch::Tensor& out_context_positions,
|
||||
torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping,
|
||||
torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices,
|
||||
const torch::Tensor& block_table, const torch::Tensor& query_start_loc,
|
||||
const std::optional<torch::Tensor>& num_rejected_tokens,
|
||||
const int64_t parallel_drafting_token_id, const int64_t block_size,
|
||||
const int64_t num_query_per_req, const int64_t num_speculative_tokens,
|
||||
const int64_t total_input_tokens, const bool has_num_rejected);
|
||||
void rejection_greedy_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax,
|
||||
const torch::Tensor& bonus_token_ids,
|
||||
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len);
|
||||
void rejection_random_sample_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids,
|
||||
const std::optional<torch::Tensor>& draft_probs,
|
||||
const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids,
|
||||
const torch::Tensor& recovered_token_ids,
|
||||
const torch::Tensor& uniform_probs,
|
||||
const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len,
|
||||
const int64_t vocab_size, const bool no_draft_probs);
|
||||
void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input,
|
||||
const torch::Tensor& cu_num_tokens,
|
||||
const int64_t replace_from, const int64_t replace_to);
|
||||
void sample_recovered_tokens_kernel_impl(
|
||||
torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens,
|
||||
const torch::Tensor& draft_token_ids,
|
||||
const std::optional<torch::Tensor>& draft_probs,
|
||||
const torch::Tensor& target_probs, const torch::Tensor& inv_q,
|
||||
const int64_t vocab_size, const bool no_draft_probs);
|
||||
} // namespace cpu_utils
|
||||
|
||||
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// vLLM custom ops
|
||||
|
||||
ops.def(
|
||||
"dynamic_4bit_int_moe("
|
||||
"Tensor x, Tensor topk_ids, Tensor topk_weights,"
|
||||
"Tensor w13_packed, Tensor w2_packed,"
|
||||
"int hidden_size, int intermediate_size,"
|
||||
"int group_size, bool apply_router_weight_on_input, int activation_kind"
|
||||
") -> Tensor");
|
||||
|
||||
ops.impl("dynamic_4bit_int_moe", torch::kCPU, &dynamic_4bit_int_moe_cpu);
|
||||
|
||||
// Activation ops
|
||||
|
||||
// Activation function used in SwiGLU.
|
||||
ops.def("silu_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCPU, &silu_and_mul);
|
||||
|
||||
// Activation function used in GeGLU with `none` approximation.
|
||||
ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_and_mul", torch::kCPU, &gelu_and_mul);
|
||||
|
||||
// Activation function used in GeGLU with `tanh` approximation.
|
||||
ops.def("gelu_tanh_and_mul(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_tanh_and_mul", torch::kCPU, &gelu_tanh_and_mul);
|
||||
|
||||
// GELU tanh implementation.
|
||||
ops.def("gelu_tanh(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_tanh", torch::kCPU, &gelu_tanh);
|
||||
|
||||
// GELU implementation used in GPT-2.
|
||||
ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_new", torch::kCPU, &gelu_new);
|
||||
|
||||
// Approximate GELU implementation.
|
||||
ops.def("gelu_fast(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_fast", torch::kCPU, &gelu_fast);
|
||||
|
||||
// Quick GELU implementation.
|
||||
ops.def("gelu_quick(Tensor! out, Tensor input) -> ()");
|
||||
ops.impl("gelu_quick", torch::kCPU, &gelu_quick);
|
||||
|
||||
#if (defined(__aarch64__) && !defined(__APPLE__))
|
||||
|
||||
ops.def(
|
||||
"activation_lut_bf16(Tensor! out, Tensor input, str activation)"
|
||||
" -> ()");
|
||||
ops.impl("activation_lut_bf16", torch::kCPU, &activation_lut_bf16);
|
||||
|
||||
#endif // (defined(__aarch64__) && !defined(__APPLE__))
|
||||
|
||||
// Layernorm
|
||||
// Apply Root Mean Square (RMS) Normalization to the input tensor.
|
||||
ops.def(
|
||||
"rms_norm(Tensor! out, Tensor input, Tensor? weight, float epsilon) -> "
|
||||
"()");
|
||||
ops.impl("rms_norm", torch::kCPU, &rms_norm);
|
||||
|
||||
// In-place fused Add and RMS Normalization.
|
||||
ops.def(
|
||||
"fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor? weight, "
|
||||
"float epsilon) -> ()");
|
||||
ops.impl("fused_add_rms_norm", torch::kCPU, &fused_add_rms_norm);
|
||||
|
||||
// Rotary embedding
|
||||
// Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
|
||||
ops.def(
|
||||
"rotary_embedding(Tensor positions, Tensor! query,"
|
||||
" Tensor!? key, int head_size,"
|
||||
" Tensor cos_sin_cache, bool is_neox, int "
|
||||
"rope_dim_offset=0, bool inverse=False) -> ()");
|
||||
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
|
||||
|
||||
// Quantization
|
||||
#if defined(__AVX512F__) || defined(__AVX2__) || \
|
||||
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__) || \
|
||||
defined(__riscv_v)
|
||||
// Helper function to release oneDNN handlers
|
||||
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
|
||||
&release_dnnl_matmul_handler);
|
||||
|
||||
// Create oneDNN GEMM handler
|
||||
ops.def(
|
||||
"create_onednn_mm_handler(Tensor b, int "
|
||||
"primitive_cache_size) -> int",
|
||||
&create_onednn_mm_handler);
|
||||
|
||||
// oneDNN GEMM
|
||||
ops.def(
|
||||
"onednn_mm(Tensor! c, Tensor a, Tensor? bias, "
|
||||
"Tensor handler_tensor) -> ()");
|
||||
ops.impl("onednn_mm", torch::kCPU, &onednn_mm);
|
||||
|
||||
// Check if oneDNN was built with ACL backend
|
||||
ops.def("is_onednn_acl_supported() -> bool", &is_onednn_acl_supported);
|
||||
|
||||
// Create oneDNN W8A8 handler
|
||||
ops.def(
|
||||
"create_onednn_scaled_mm_handler(Tensor b, Tensor b_scales, ScalarType "
|
||||
"output_type, bool dynamic_act_quant, bool use_azp, int "
|
||||
"primitive_cache_size) -> int",
|
||||
&create_onednn_scaled_mm_handler);
|
||||
|
||||
// oneDNN scaled_mm for W8A8 with static per-tensor activation quantization
|
||||
ops.def(
|
||||
"onednn_scaled_mm(Tensor! c, Tensor a, Tensor a_scales, Tensor? azp, "
|
||||
"Tensor? azp_adj, Tensor? bias, Tensor handler_tensor) -> ()");
|
||||
ops.impl("onednn_scaled_mm", torch::kCPU, &onednn_scaled_mm);
|
||||
|
||||
// Compute int8 quantized tensor for given scaling factor.
|
||||
ops.def(
|
||||
"static_scaled_int8_quant(Tensor! out, Tensor input, Tensor scale,"
|
||||
"Tensor? azp) -> ()");
|
||||
ops.impl("static_scaled_int8_quant", torch::kCPU, &static_scaled_int8_quant);
|
||||
|
||||
// Compute int8 quantized tensor and scaling factor
|
||||
ops.def(
|
||||
"dynamic_scaled_int8_quant(Tensor! out, Tensor input, Tensor! scale, "
|
||||
"Tensor!? azp) -> ()");
|
||||
ops.impl("dynamic_scaled_int8_quant", torch::kCPU,
|
||||
&dynamic_scaled_int8_quant);
|
||||
#endif
|
||||
|
||||
// SHM CCL
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
|
||||
defined(__powerpc64__)
|
||||
ops.def(
|
||||
"init_shm_manager(str name, int group_size, int rank, int thread_num) -> "
|
||||
"int",
|
||||
&init_shm_manager);
|
||||
ops.def("join_shm_manager(int handle, str name) -> str", &join_shm_manager);
|
||||
ops.def("shm_allreduce(int handle, Tensor! data) -> ()");
|
||||
ops.impl("shm_allreduce", torch::kCPU, &shm_allreduce);
|
||||
ops.def(
|
||||
"shm_gather(int handle, Tensor data, Tensor[](a!)? outputs, int dst) -> "
|
||||
"()");
|
||||
ops.impl("shm_gather", torch::kCPU, &shm_gather);
|
||||
ops.def(
|
||||
"shm_all_gather(int handle, Tensor data, Tensor! output) -> "
|
||||
"()");
|
||||
ops.impl("shm_all_gather", torch::kCPU, &shm_all_gather);
|
||||
ops.def(
|
||||
"shm_send_tensor_list(int handle, Tensor[](a) tensor_list, int dst) -> "
|
||||
"()");
|
||||
ops.impl("shm_send_tensor_list", torch::kCPU, &shm_send_tensor_list);
|
||||
ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
|
||||
&shm_recv_tensor_list);
|
||||
#endif // #if defined(__AVX512F__) || defined(__aarch64__)
|
||||
|
||||
// sgl-kernels
|
||||
#if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
|
||||
ops.def(
|
||||
"weight_packed_linear(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!)? "
|
||||
"bias, bool is_vnni) -> Tensor");
|
||||
ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
|
||||
ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
|
||||
ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
|
||||
ops.def("convert_scale_packed(Tensor! scale) -> Tensor");
|
||||
ops.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
|
||||
ops.def(
|
||||
"fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
|
||||
"topk_weights, Tensor topk_ids, bool "
|
||||
"inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
|
||||
"Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, "
|
||||
"Tensor? w1_bias, Tensor? w2_bias, float? alpha, float? limit, "
|
||||
"bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
|
||||
ops.def(
|
||||
"int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, "
|
||||
"Tensor? bias, ScalarType out_dtype, bool is_vnni) -> Tensor");
|
||||
ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
|
||||
&int8_scaled_mm_with_quant);
|
||||
|
||||
// Adapted from sglang: FP8 W8A16 kernel
|
||||
ops.def(
|
||||
"fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
|
||||
"scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
|
||||
|
||||
// Adapted from sglang: casual_conv1d kernels
|
||||
ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
|
||||
ops.impl("causal_conv1d_weight_pack", torch::kCPU,
|
||||
&causal_conv1d_weight_pack);
|
||||
ops.def(
|
||||
"causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
|
||||
"conv_states, Tensor? query_start_loc,"
|
||||
"Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
|
||||
"int pad_slot_id, bool is_vnni) -> "
|
||||
"Tensor");
|
||||
ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
|
||||
"weight, Tensor? bias, bool silu_activation,"
|
||||
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
|
||||
#endif
|
||||
|
||||
#if (defined(__AVX512BF16__) && defined(__AVX512F__) && \
|
||||
defined(__AVX512VNNI__)) || \
|
||||
defined(__riscv)
|
||||
// Adapted from sglang: INT4 W4A8 kernels
|
||||
ops.def(
|
||||
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
|
||||
"scales, int quant_method_4bit) -> (Tensor, "
|
||||
"Tensor, Tensor)");
|
||||
ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
|
||||
&convert_weight_packed_scale_zp);
|
||||
|
||||
ops.def(
|
||||
"int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
|
||||
"Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
|
||||
ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
|
||||
#endif
|
||||
|
||||
// Adapted from sglang: GDN kernels
|
||||
ops.def(
|
||||
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "
|
||||
"Tensor g, Tensor beta, "
|
||||
"Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool "
|
||||
"head_first, "
|
||||
"bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");
|
||||
ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU,
|
||||
&chunk_gated_delta_rule_cpu);
|
||||
ops.def(
|
||||
"fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor "
|
||||
"dt_bias, Tensor q, Tensor k, Tensor v, Tensor "
|
||||
"a, Tensor b, Tensor(a!) initial_state_source, Tensor "
|
||||
"initial_state_indices, Tensor cu_seqlens, bool "
|
||||
"use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "
|
||||
"softplus_threshold=20.0) -> Tensor");
|
||||
ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,
|
||||
&fused_sigmoid_gating_delta_rule_update_cpu);
|
||||
ops.def(
|
||||
"fused_sigmoid_gating_delta_rule_update_spec_cpu(Tensor A_log, Tensor "
|
||||
"dt_bias, Tensor q, Tensor k, Tensor v, Tensor a, Tensor b, "
|
||||
"Tensor(a!) initial_state_source, Tensor spec_state_indices, "
|
||||
"Tensor num_accepted_tokens, Tensor cu_seqlens, bool "
|
||||
"use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "
|
||||
"softplus_threshold=20.0) -> Tensor");
|
||||
ops.impl("fused_sigmoid_gating_delta_rule_update_spec_cpu", torch::kCPU,
|
||||
&fused_sigmoid_gating_delta_rule_update_spec_cpu);
|
||||
ops.def(
|
||||
"fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "
|
||||
"-> (Tensor, Tensor)");
|
||||
ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);
|
||||
|
||||
// CPU attention kernels
|
||||
ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa);
|
||||
ops.def(
|
||||
"get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
|
||||
"int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
|
||||
"query_start_loc, bool casual, int window_size, str isa_hint, bool "
|
||||
"enable_kv_split, Tensor? dynamic_causal) -> Tensor",
|
||||
&get_scheduler_metadata);
|
||||
ops.def(
|
||||
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
|
||||
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
|
||||
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
|
||||
"()",
|
||||
&cpu_attn_reshape_and_cache);
|
||||
ops.def(
|
||||
"cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
|
||||
"value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
|
||||
"seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
|
||||
"sliding_window_size, Tensor block_table, "
|
||||
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? "
|
||||
"dynamic_causal, "
|
||||
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
|
||||
"()",
|
||||
&cpu_attention_with_kv_cache);
|
||||
|
||||
// placeholders
|
||||
ops.def("static_scaled_fp8_quant() -> ()", placeholder_op);
|
||||
ops.def("dynamic_scaled_fp8_quant() -> ()", placeholder_op);
|
||||
ops.def("dynamic_per_token_scaled_fp8_quant() -> ()", placeholder_op);
|
||||
|
||||
// WNA16
|
||||
#if defined(__AVX512F__) || defined(__riscv_v)
|
||||
ops.def(
|
||||
"cpu_gemm_wna16(Tensor input, Tensor q_weight, Tensor(a2!) output, "
|
||||
"Tensor scales, Tensor? zeros, Tensor? g_idx, Tensor? bias, SymInt "
|
||||
"pack_factor, str isa_hint) -> ()");
|
||||
ops.impl("cpu_gemm_wna16", torch::kCPU, &cpu_gemm_wna16);
|
||||
#endif
|
||||
|
||||
// fused moe
|
||||
#if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
|
||||
ops.def(
|
||||
"prepack_moe_weight(Tensor weight, Tensor(a1!) packed_weight, str isa) "
|
||||
"-> ()");
|
||||
ops.impl("prepack_moe_weight", torch::kCPU, &prepack_moe_weight);
|
||||
ops.def(
|
||||
"cpu_fused_moe(Tensor(a0!) output, Tensor input, Tensor w13, Tensor w2, "
|
||||
"Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, Tensor topk_id, "
|
||||
"bool skip_weighted, "
|
||||
"str act, str isa) -> ()");
|
||||
ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);
|
||||
#endif // #if defined(__AVX512F__) || (defined(ARM_BF16_SUPPORT))
|
||||
ops.def(
|
||||
"mla_decode_kvcache("
|
||||
" Tensor! out, Tensor query, Tensor kv_cache,"
|
||||
" float scale, Tensor block_tables, Tensor seq_lens) -> ()");
|
||||
ops.impl("mla_decode_kvcache", torch::kCPU, &mla_decode_kvcache);
|
||||
|
||||
ops.def(
|
||||
"compute_slot_mapping_kernel_impl(Tensor query_start_loc, Tensor "
|
||||
"positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt "
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
ops.def(
|
||||
"eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
|
||||
"Tensor valid_sampled_tokens_count, Tensor query_start_loc_gpu, "
|
||||
"Tensor(a3!) token_indices_to_sample, "
|
||||
"Tensor(a4!) num_rejected_tokens_gpu, "
|
||||
"SymInt num_reqs) -> ()",
|
||||
&cpu_utils::eagle_prepare_inputs_padded_kernel_impl);
|
||||
ops.def(
|
||||
"eagle_prepare_next_token_padded_kernel_impl("
|
||||
"Tensor sampled_token_ids, Tensor discard_request_mask, "
|
||||
"Tensor backup_next_token_ids, Tensor(a3!) next_token_ids, "
|
||||
"Tensor(a4!) valid_sampled_tokens_count, SymInt vocab_size, "
|
||||
"SymInt num_sampled_tokens_per_req, SymInt num_reqs) -> ()",
|
||||
&cpu_utils::eagle_prepare_next_token_padded_kernel_impl);
|
||||
ops.def(
|
||||
"eagle_step_slot_mapping_metadata_kernel_impl("
|
||||
"Tensor positions, Tensor block_table, Tensor(a2!) seq_lens, "
|
||||
"Tensor(a3!) out_clamped_positions, Tensor(a4!) out_slot_mapping, "
|
||||
"SymInt block_size, SymInt max_model_len, SymInt PAD_ID) -> ()",
|
||||
&cpu_utils::eagle_step_slot_mapping_metadata_kernel_impl);
|
||||
ops.def(
|
||||
"copy_and_expand_eagle_inputs_kernel_impl("
|
||||
"Tensor target_token_ids, Tensor target_positions, "
|
||||
"Tensor next_token_ids, Tensor(a3!) out_input_ids, "
|
||||
"Tensor(a4!) out_positions, "
|
||||
"Tensor(a5!) out_is_rejected_token_mask, "
|
||||
"Tensor(a6!) out_is_masked_token_mask, "
|
||||
"Tensor(a7!) out_new_token_indices, "
|
||||
"Tensor(a8!) out_hidden_state_mapping, "
|
||||
"Tensor query_start_loc, Tensor query_end_loc, "
|
||||
"SymInt padding_token_id, SymInt parallel_drafting_token_id, "
|
||||
"SymInt total_input_tokens, SymInt num_padding_slots_per_request, "
|
||||
"bool shift_input_ids) -> ()",
|
||||
&cpu_utils::copy_and_expand_eagle_inputs_kernel_impl);
|
||||
ops.def(
|
||||
"copy_and_expand_dflash_inputs_kernel_impl("
|
||||
"Tensor next_token_ids, Tensor target_positions, "
|
||||
"Tensor(a2!) out_input_ids, Tensor(a3!) out_context_positions, "
|
||||
"Tensor(a4!) out_query_positions, "
|
||||
"Tensor(a5!) out_context_slot_mapping, "
|
||||
"Tensor(a6!) out_query_slot_mapping, "
|
||||
"Tensor(a7!) out_token_indices, Tensor block_table, "
|
||||
"Tensor query_start_loc, Tensor? num_rejected_tokens, "
|
||||
"SymInt parallel_drafting_token_id, SymInt block_size, "
|
||||
"SymInt num_query_per_req, SymInt num_speculative_tokens, "
|
||||
"SymInt total_input_tokens, bool has_num_rejected) -> ()",
|
||||
&cpu_utils::copy_and_expand_dflash_inputs_kernel_impl);
|
||||
ops.def(
|
||||
"rejection_greedy_sample_kernel_impl("
|
||||
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
|
||||
"Tensor draft_token_ids, Tensor target_argmax, "
|
||||
"Tensor bonus_token_ids, Tensor? is_greedy, "
|
||||
"SymInt max_spec_len) -> ()",
|
||||
&cpu_utils::rejection_greedy_sample_kernel_impl);
|
||||
ops.def(
|
||||
"rejection_random_sample_kernel_impl("
|
||||
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
|
||||
"Tensor draft_token_ids, Tensor? draft_probs, "
|
||||
"Tensor target_probs, Tensor bonus_token_ids, "
|
||||
"Tensor recovered_token_ids, Tensor uniform_probs, "
|
||||
"Tensor? is_greedy, SymInt max_spec_len, SymInt vocab_size, "
|
||||
"bool no_draft_probs) -> ()",
|
||||
&cpu_utils::rejection_random_sample_kernel_impl);
|
||||
ops.def(
|
||||
"expand_kernel_impl(Tensor(a0!) output, Tensor input, "
|
||||
"Tensor cu_num_tokens, SymInt replace_from, "
|
||||
"SymInt replace_to) -> ()",
|
||||
&cpu_utils::expand_kernel_impl);
|
||||
ops.def(
|
||||
"sample_recovered_tokens_kernel_impl("
|
||||
"Tensor(a0!) output_token_ids, Tensor cu_num_draft_tokens, "
|
||||
"Tensor draft_token_ids, Tensor? draft_probs, "
|
||||
"Tensor target_probs, Tensor inv_q, SymInt vocab_size, "
|
||||
"bool no_draft_probs) -> ()",
|
||||
&cpu_utils::sample_recovered_tokens_kernel_impl);
|
||||
}
|
||||
|
||||
REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
|
||||
@@ -0,0 +1,149 @@
|
||||
#ifndef VLLM_NUMA_DISABLED
|
||||
#include <numa.h>
|
||||
#include <unistd.h>
|
||||
#include <string>
|
||||
#include <sched.h>
|
||||
#endif
|
||||
#if __GLIBC__ == 2 && __GLIBC_MINOR__ < 30
|
||||
#include <unistd.h>
|
||||
#include <sys/syscall.h>
|
||||
#define gettid() syscall(SYS_gettid)
|
||||
#endif
|
||||
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
#ifdef VLLM_NUMA_DISABLED
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
|
||||
#else
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids) {
|
||||
// Memory node binding
|
||||
if (numa_available() != -1) {
|
||||
// Concatenate all node_ids into a single comma-separated string
|
||||
if (!node_ids.empty()) {
|
||||
std::string node_ids_str;
|
||||
for (const int node_id : node_ids) {
|
||||
if (!node_ids_str.empty()) {
|
||||
node_ids_str += ",";
|
||||
}
|
||||
node_ids_str += std::to_string(node_id);
|
||||
}
|
||||
|
||||
bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
|
||||
bitmask* src_mask = numa_get_mems_allowed();
|
||||
|
||||
int pid = getpid();
|
||||
|
||||
if (mask && src_mask) {
|
||||
// move all existing pages to the specified numa node.
|
||||
*(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
|
||||
int page_num = numa_migrate_pages(pid, src_mask, mask);
|
||||
if (page_num == -1) {
|
||||
TORCH_WARN("numa_migrate_pages failed. errno: " +
|
||||
std::to_string(errno));
|
||||
}
|
||||
|
||||
// Restrict memory allocation to the selected NUMA node(s).
|
||||
// Enhances memory locality for the threads bound to those NUMA CPUs.
|
||||
if (node_ids.size() > 1) {
|
||||
errno = 0;
|
||||
numa_set_interleave_mask(mask);
|
||||
if (errno != 0) {
|
||||
TORCH_WARN("numa_set_interleave_mask failed. errno: " +
|
||||
std::to_string(errno));
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"NUMA binding: Using INTERLEAVE policy for memory "
|
||||
"allocation across multiple NUMA nodes (nodes: " +
|
||||
node_ids_str +
|
||||
"). Memory allocations will be "
|
||||
"interleaved across the specified NUMA nodes.");
|
||||
}
|
||||
} else {
|
||||
errno = 0;
|
||||
numa_set_membind(mask);
|
||||
if (errno != 0) {
|
||||
TORCH_WARN("numa_set_membind failed. errno: " +
|
||||
std::to_string(errno));
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"NUMA binding: Using MEMBIND policy for memory "
|
||||
"allocation on the NUMA nodes (" +
|
||||
node_ids_str +
|
||||
"). Memory allocations will be "
|
||||
"strictly bound to these NUMA nodes.");
|
||||
}
|
||||
}
|
||||
|
||||
numa_set_strict(1);
|
||||
|
||||
numa_free_nodemask(mask);
|
||||
numa_free_nodemask(src_mask);
|
||||
} else {
|
||||
TORCH_WARN(
|
||||
"numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
|
||||
std::to_string(errno));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif // VLLM_NUMA_DISABLED
|
||||
|
||||
namespace cpu_utils {
|
||||
ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
|
||||
this->realloc(allocation_unit * 128);
|
||||
}
|
||||
|
||||
void ScratchPadManager::realloc(size_t new_size) {
|
||||
new_size = round(new_size);
|
||||
if (new_size > size_) {
|
||||
void* new_ptr = std::aligned_alloc(64, new_size);
|
||||
TORCH_CHECK(new_ptr != nullptr,
|
||||
"ScratchPadManager: aligned_alloc failed for size ", new_size);
|
||||
if (ptr_ != nullptr) {
|
||||
std::free(ptr_);
|
||||
}
|
||||
ptr_ = new_ptr;
|
||||
size_ = new_size;
|
||||
}
|
||||
}
|
||||
|
||||
ScratchPadManager* ScratchPadManager::get_scratchpad_manager() {
|
||||
static ScratchPadManager manager;
|
||||
return &manager;
|
||||
}
|
||||
} // namespace cpu_utils
|
||||
|
||||
void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
const torch::Tensor positions,
|
||||
const torch::Tensor block_table,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size) {
|
||||
const int32_t req_num = query_start_loc.size(0) - 1;
|
||||
const int64_t block_table_stride = block_table.stride(0);
|
||||
|
||||
const int32_t* __restrict__ query_start_loc_ptr =
|
||||
query_start_loc.data_ptr<int32_t>();
|
||||
const int64_t* __restrict__ positions_ptr = positions.data_ptr<int64_t>();
|
||||
const int32_t* __restrict__ blocktable_ptr = block_table.data_ptr<int32_t>();
|
||||
int64_t* __restrict__ slot_mapping_ptr = slot_mapping.data_ptr<int64_t>();
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int32_t req_idx = 0; req_idx < req_num; ++req_idx) {
|
||||
int32_t token_start_idx = query_start_loc_ptr[req_idx];
|
||||
int32_t token_end_idx = query_start_loc_ptr[req_idx + 1];
|
||||
int32_t token_num = token_end_idx - token_start_idx;
|
||||
const int64_t* __restrict__ curr_position_ptr =
|
||||
positions_ptr + token_start_idx;
|
||||
int64_t* __restrict__ curr_slot_mapping_ptr =
|
||||
slot_mapping_ptr + token_start_idx;
|
||||
const int32_t* __restrict__ curr_block_table_ptr =
|
||||
blocktable_ptr + req_idx * block_table_stride;
|
||||
|
||||
for (int32_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
int64_t token_position = curr_position_ptr[token_idx];
|
||||
int64_t block_id = curr_block_table_ptr[token_position / block_size];
|
||||
curr_slot_mapping_ptr[token_idx] =
|
||||
block_id * block_size + token_position % block_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
#ifndef UTILS_HPP
|
||||
#define UTILS_HPP
|
||||
|
||||
#include <atomic>
|
||||
#include <string>
|
||||
#include <unistd.h>
|
||||
#include <ATen/cpu/Utils.h>
|
||||
|
||||
#include "cpu/cpu_types.hpp"
|
||||
|
||||
namespace cpu_utils {
|
||||
enum class ISA { AMX, VEC, RVV, NEON };
|
||||
|
||||
inline ISA get_isa(const std::string& isa) {
|
||||
if (isa == "amx") {
|
||||
return ISA::AMX;
|
||||
} else if (isa == "vec") {
|
||||
return ISA::VEC;
|
||||
} else if (isa == "rvv") {
|
||||
return ISA::RVV;
|
||||
} else if (isa == "neon") {
|
||||
return ISA::NEON;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid isa type: " + isa);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct VecTypeTrait {
|
||||
using vec_t = void;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecTypeTrait<float> {
|
||||
using vec_t = vec_op::FP32Vec16;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct VecTypeTrait<c10::BFloat16> {
|
||||
using vec_t = vec_op::BF16Vec16;
|
||||
};
|
||||
|
||||
#if !defined(__powerpc__)
|
||||
template <>
|
||||
struct VecTypeTrait<c10::Half> {
|
||||
using vec_t = vec_op::FP16Vec16;
|
||||
};
|
||||
#endif
|
||||
|
||||
struct Counter {
|
||||
std::atomic<int64_t> counter;
|
||||
char _padding[56];
|
||||
|
||||
Counter() : counter(0) {}
|
||||
|
||||
void reset_counter() { counter.store(0); }
|
||||
|
||||
int64_t acquire_counter() { return counter++; }
|
||||
};
|
||||
|
||||
inline int64_t get_available_l2_size() {
|
||||
#if defined(__s390x__) || defined(__powerpc__)
|
||||
static int64_t size = []() {
|
||||
uint32_t l2_cache_size = 0;
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
auto it = caps.find("l2_cache_size");
|
||||
if (it != caps.end()) {
|
||||
l2_cache_size = static_cast<uint32_t>(it->second.toInt());
|
||||
}
|
||||
if (l2_cache_size == 0) {
|
||||
long sys_l2 = sysconf(_SC_LEVEL2_CACHE_SIZE);
|
||||
if (sys_l2 > 0) {
|
||||
l2_cache_size = static_cast<uint32_t>(sys_l2);
|
||||
}
|
||||
}
|
||||
if (l2_cache_size == 0) {
|
||||
l2_cache_size = 256 * 1024;
|
||||
}
|
||||
return static_cast<int64_t>(l2_cache_size) >> 1;
|
||||
}();
|
||||
return size;
|
||||
#else
|
||||
static int64_t size = []() {
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
|
||||
return l2_cache_size >> 1;
|
||||
}();
|
||||
return size;
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int32_t alignment_v, typename T>
|
||||
inline T round_up(T size) {
|
||||
T alignment = alignment_v;
|
||||
return (((size + alignment - 1) / alignment) * alignment);
|
||||
}
|
||||
|
||||
template <int32_t alignment_v, typename T>
|
||||
inline T round_down(T size) {
|
||||
T alignment = alignment_v;
|
||||
return (size / alignment) * alignment;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
inline void print_logits(const char* name, T* ptr, int32_t row, int32_t col,
|
||||
int32_t stride) {
|
||||
std::stringstream ss;
|
||||
ss << std::fixed << std::setprecision(5) << name << ": [\n";
|
||||
auto* curr_logits_buffer = ptr;
|
||||
for (int32_t m = 0; m < row; ++m) {
|
||||
for (int32_t n = 0; n < col; ++n) {
|
||||
ss << curr_logits_buffer[n] << ", ";
|
||||
}
|
||||
ss << "\n";
|
||||
curr_logits_buffer += stride;
|
||||
}
|
||||
ss << "]\n";
|
||||
std::printf("%s", ss.str().c_str());
|
||||
}
|
||||
|
||||
class ScratchPadManager {
|
||||
public:
|
||||
static constexpr size_t allocation_unit = 4 * 1024; // 4KB
|
||||
|
||||
static ScratchPadManager* get_scratchpad_manager();
|
||||
|
||||
ScratchPadManager();
|
||||
|
||||
template <typename T>
|
||||
T* get_data() {
|
||||
return reinterpret_cast<T*>(ptr_);
|
||||
}
|
||||
|
||||
static size_t round(size_t size) {
|
||||
return ((size + allocation_unit - 1) / allocation_unit) * allocation_unit;
|
||||
}
|
||||
|
||||
void realloc(size_t new_size);
|
||||
|
||||
private:
|
||||
size_t size_;
|
||||
void* ptr_;
|
||||
};
|
||||
} // namespace cpu_utils
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user