247 lines
8.4 KiB
C++
247 lines
8.4 KiB
C++
#pragma once
|
|
#include "common.h"
|
|
#include "vec.h"
|
|
#include "vec_pack.h"
|
|
|
|
template <typename scalar_t>
|
|
inline void fill_stub(scalar_t* __restrict__ out, float val, int size) {
|
|
using Vec = at::vec::Vectorized<scalar_t>;
|
|
constexpr int kVecSize = Vec::size();
|
|
const Vec data_vec = Vec(static_cast<scalar_t>(val));
|
|
int d = 0;
|
|
#pragma GCC unroll 4
|
|
for (; d <= size - kVecSize; d += kVecSize) {
|
|
data_vec.store(out + d);
|
|
}
|
|
if (size - d > 0) {
|
|
data_vec.store(out + d, size - d);
|
|
}
|
|
}
|
|
|
|
template <typename scalar_t, int BLOCK_N>
|
|
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) {
|
|
static_assert(BLOCK_N % 32 == 0);
|
|
using bVec = at::vec::Vectorized<scalar_t>;
|
|
using fVec = at::vec::Vectorized<float>;
|
|
|
|
constexpr int COLS = BLOCK_N / 16;
|
|
auto store = [&](auto i) {
|
|
constexpr int col = i % COLS;
|
|
// for COLS = 2, 4 use 512bit store
|
|
if constexpr (col % 2 == 0) {
|
|
fVec a_fvec0 = fVec::loadu(input + col * 16);
|
|
fVec a_fvec1 = fVec::loadu(input + col * 16 + 16);
|
|
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
|
|
out_bvec.store(out + col * 16);
|
|
}
|
|
};
|
|
Unroll<COLS>{}(store);
|
|
}
|
|
|
|
template <typename scalar_t>
|
|
inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) {
|
|
using bVec = at::vec::Vectorized<scalar_t>;
|
|
using fVec = at::vec::Vectorized<float>;
|
|
constexpr int kVecSize = bVec::size();
|
|
const fVec s_fvec = fVec(s);
|
|
int d = 0;
|
|
#pragma GCC unroll 4
|
|
for (; d <= size - kVecSize; d += kVecSize) {
|
|
fVec a_fvec0 = fVec::loadu(acc + d) * s_fvec;
|
|
fVec a_fvec1 = fVec::loadu(acc + d + fVec::size()) * s_fvec;
|
|
bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1);
|
|
out_bvec.store(out + d);
|
|
}
|
|
for (; d < size; ++d) {
|
|
out[d] = static_cast<scalar_t>(acc[d] * s);
|
|
}
|
|
}
|
|
|
|
#if defined(CPU_CAPABILITY_AVX512)
|
|
template <>
|
|
inline void copy_stub<at::BFloat16>(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) {
|
|
const __m512 vscale = _mm512_set1_ps(s);
|
|
int d = 0;
|
|
#pragma GCC unroll 4
|
|
for (; d <= size - 32; d += 32) {
|
|
__m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
|
|
__m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale);
|
|
__m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
|
|
_mm512_storeu_si512(out + d, vb);
|
|
}
|
|
int remainder = size - d;
|
|
if (remainder > 0) {
|
|
if (remainder <= 16) {
|
|
const __mmask16 vmask = (1ULL << remainder) - 1;
|
|
__m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale);
|
|
__m256i vb = (__m256i)(_mm512_cvtneps_pbh(va));
|
|
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb);
|
|
} else { // remainder > 16
|
|
const __mmask16 vmask = (1ULL << (remainder - 16)) - 1;
|
|
__m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale);
|
|
__m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale);
|
|
__m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0));
|
|
const __mmask32 vmask2 = (1ULL << remainder) - 1;
|
|
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb);
|
|
}
|
|
}
|
|
}
|
|
#endif
|
|
|
|
template <typename scalar_t, int BLOCK_M, int BLOCK_N>
|
|
struct flash_attn_softmax {
|
|
static inline void apply(
|
|
float* __restrict__ s_i,
|
|
scalar_t* __restrict__ s_delta2,
|
|
float* __restrict__ v_prime,
|
|
float* __restrict__ s_prime,
|
|
float* __restrict__ m_prime,
|
|
int m_size,
|
|
int n_size,
|
|
int padded_n_size,
|
|
int head_size_v,
|
|
const float sm_scale) {
|
|
using Vec = at::vec::Vectorized<float>;
|
|
const Vec scale_vec = Vec(sm_scale);
|
|
float* s_delta = s_i;
|
|
for (int row = 0; row < m_size; ++row) {
|
|
// s_i <- s_i * scale
|
|
at::vec::map<float>(
|
|
[scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
|
|
|
|
// m_i: max value per row
|
|
float m_i = at::vec::reduce_all<float>(
|
|
[](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size);
|
|
m_i = std::max(m_i, m_prime[row]);
|
|
|
|
// m_delta <- exp(m' - m_i)
|
|
float m_delta = std::exp(m_prime[row] - m_i);
|
|
|
|
// s_delta <- exp(s_i - m_i)
|
|
at::vec::map<float>(
|
|
[m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size);
|
|
|
|
// s' <- s' * m_delta + sum(s_delta)
|
|
s_prime[row] *= m_delta;
|
|
s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size);
|
|
|
|
m_prime[row] = m_i;
|
|
|
|
// v' <- v' * m_delta
|
|
at::vec::map<float>(
|
|
[m_delta](Vec x) { return x * Vec(m_delta); },
|
|
v_prime + row * head_size_v,
|
|
v_prime + row * head_size_v,
|
|
head_size_v);
|
|
|
|
// pad s_delta with 0 first and then convert to scalar_t
|
|
fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size);
|
|
copy_stub<scalar_t, BLOCK_N>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N);
|
|
}
|
|
}
|
|
};
|
|
|
|
#if defined(CPU_CAPABILITY_AVX512)
|
|
template <int BLOCK_M, int BLOCK_N>
|
|
struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> {
|
|
static inline void apply(
|
|
float* __restrict__ s_i,
|
|
at::BFloat16* __restrict__ s_delta2,
|
|
float* __restrict__ v_prime,
|
|
float* __restrict__ s_prime,
|
|
float* __restrict__ m_prime,
|
|
int m_size,
|
|
int n_size,
|
|
int padded_n_size,
|
|
int head_size_v,
|
|
const float sm_scale) {
|
|
float* s_delta = s_i;
|
|
const __m512 vscale = _mm512_set1_ps(sm_scale);
|
|
|
|
int n_remainder = n_size & 15; // 0xF
|
|
const __mmask16 vmask = (1ULL << n_remainder) - 1;
|
|
|
|
int v_remainder = head_size_v & 15; // 0xF
|
|
const __mmask16 vmask1 = (1ULL << v_remainder) - 1;
|
|
|
|
constexpr float NEG_INF = -std::numeric_limits<float>::infinity();
|
|
|
|
__m512 va;
|
|
__m256i vb;
|
|
__m512 vmax;
|
|
__m512 vsum;
|
|
__m512 vmdelta;
|
|
|
|
const __m512 vneg_inf = _mm512_set1_ps(NEG_INF);
|
|
|
|
for (int m = 0; m < m_size; ++m) {
|
|
vmax = vneg_inf;
|
|
|
|
// s_i <- s_i * scale
|
|
int n = 0;
|
|
for (; n <= n_size - 16; n += 16) {
|
|
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
|
|
vmax = _mm512_max_ps(va, vmax);
|
|
}
|
|
if (n_remainder > 0) {
|
|
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
|
|
vmax = _mm512_max_ps(va, vmax);
|
|
}
|
|
|
|
// m_i: max value per row
|
|
float m_i = _mm512_reduce_max_ps(vmax);
|
|
m_i = std::max(m_i, m_prime[m]);
|
|
vmax = _mm512_set1_ps(m_i);
|
|
|
|
// m_delta <- exp(m' - m_i)
|
|
float m_delta = std::exp(m_prime[m] - m_i);
|
|
|
|
// s_delta <- exp(s_i - m_i)
|
|
vsum = _mm512_setzero_ps();
|
|
for (n = 0; n <= n_size - 16; n += 16) {
|
|
va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale);
|
|
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
|
|
vsum = _mm512_add_ps(vsum, va);
|
|
|
|
vb = (__m256i)(_mm512_cvtneps_pbh(va));
|
|
_mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb);
|
|
}
|
|
if (n_remainder > 0) {
|
|
va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale);
|
|
va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax));
|
|
vsum = _mm512_add_ps(vsum, va);
|
|
|
|
vb = (__m256i)(_mm512_cvtneps_pbh(va));
|
|
_mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb);
|
|
}
|
|
|
|
// s' <- s' * m_delta + sum(s_delta)
|
|
s_prime[m] *= m_delta;
|
|
s_prime[m] += _mm512_reduce_add_ps(vsum);
|
|
|
|
m_prime[m] = m_i;
|
|
|
|
// pad s_delta with 0, pad_size range from [0, 32)
|
|
int pad_size = padded_n_size - n_size;
|
|
if (pad_size > 0) {
|
|
const __m512i vzero = _mm512_setzero_si512();
|
|
__mmask32 vmask2 = (1ULL << pad_size) - 1;
|
|
_mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero);
|
|
}
|
|
|
|
// v' <- v' * m_delta
|
|
vmdelta = _mm512_set1_ps(m_delta);
|
|
int k = 0;
|
|
for (; k <= head_size_v - 16; k += 16) {
|
|
va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta);
|
|
_mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va);
|
|
}
|
|
if (v_remainder > 0) {
|
|
va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta);
|
|
_mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
#endif
|