| 1 | // Adapted from |
| 2 | // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu |
| 3 | |
| 4 | // clang-format off |
| 5 | |
| 6 | #pragma once |
| 7 | #include "common.h" |
| 8 | #include "vec.h" |
| 9 | #include "vec_pack.h" |
| 10 | |
| 11 | template <typename scalar_t> |
| 12 | inline void fill_stub(scalar_t* __restrict__ out, float val, int size) { |
| 13 | using Vec = at::vec::Vectorized<scalar_t>; |
| 14 | constexpr int kVecSize = Vec::size(); |
| 15 | const Vec data_vec = Vec(static_cast<scalar_t>(val)); |
| 16 | int d = 0; |
| 17 | #pragma GCC unroll 4 |
| 18 | for (; d <= size - kVecSize; d += kVecSize) { |
| 19 | data_vec.store(out + d); |
| 20 | } |
| 21 | if (size - d > 0) { |
| 22 | data_vec.store(out + d, size - d); |
| 23 | } |
| 24 | } |
| 25 | |
| 26 | template <typename scalar_t, int BLOCK_N> |
| 27 | inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input) { |
| 28 | static_assert(BLOCK_N % 32 == 0); |
| 29 | using bVec = at::vec::Vectorized<scalar_t>; |
| 30 | using fVec = at::vec::Vectorized<float>; |
| 31 | |
| 32 | constexpr int COLS = BLOCK_N / 16; |
| 33 | auto store = [&](auto i) { |
| 34 | constexpr int col = i % COLS; |
| 35 | // for COLS = 2, 4 use 512bit store |
| 36 | if constexpr (col % 2 == 0) { |
| 37 | auto [a_fvec0, a_fvec1] = load_float_vec2(input + col * 16); |
| 38 | bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1); |
| 39 | out_bvec.store(out + col * 16); |
| 40 | } |
| 41 | }; |
| 42 | Unroll<COLS>{}(store); |
| 43 | } |
| 44 | |
| 45 | template <typename scalar_t> |
| 46 | inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ acc, float s, int size) { |
| 47 | using bVec = at::vec::Vectorized<scalar_t>; |
| 48 | using fVec = at::vec::Vectorized<float>; |
| 49 | constexpr int kVecSize = bVec::size(); |
| 50 | const fVec s_fvec = fVec(s); |
| 51 | int d = 0; |
| 52 | #pragma GCC unroll 4 |
| 53 | for (; d <= size - kVecSize; d += kVecSize) { |
| 54 | auto [a_fvec0, a_fvec1] = load_float_vec2(acc + d); |
| 55 | a_fvec0 = a_fvec0 * s_fvec; |
| 56 | a_fvec1 = a_fvec1 * s_fvec; |
| 57 | bVec out_bvec = convert_from_float_ext<scalar_t>(a_fvec0, a_fvec1); |
| 58 | out_bvec.store(out + d); |
| 59 | } |
| 60 | for (; d < size; ++d) { |
| 61 | out[d] = static_cast<scalar_t>(acc[d] * s); |
| 62 | } |
| 63 | } |
| 64 | |
| 65 | #if defined(CPU_CAPABILITY_AVX512) |
| 66 | template <> |
| 67 | inline void copy_stub<at::BFloat16>(at::BFloat16* __restrict__ out, const float* __restrict__ acc, float s, int size) { |
| 68 | const __m512 vscale = _mm512_set1_ps(s); |
| 69 | int d = 0; |
| 70 | #pragma GCC unroll 4 |
| 71 | for (; d <= size - 32; d += 32) { |
| 72 | __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); |
| 73 | __m512 va1 = _mm512_mul_ps(_mm512_loadu_ps(acc + d + 16), vscale); |
| 74 | __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); |
| 75 | _mm512_storeu_si512(out + d, vb); |
| 76 | } |
| 77 | int remainder = size - d; |
| 78 | if (remainder > 0) { |
| 79 | if (remainder <= 16) { |
| 80 | const __mmask16 vmask = (1ULL << remainder) - 1; |
| 81 | __m512 va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d), vscale); |
| 82 | __m256i vb = (__m256i)(_mm512_cvtneps_pbh(va)); |
| 83 | _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(out + d), vmask, vb); |
| 84 | } else { // remainder > 16 |
| 85 | const __mmask16 vmask = (1ULL << (remainder - 16)) - 1; |
| 86 | __m512 va0 = _mm512_mul_ps(_mm512_loadu_ps(acc + d), vscale); |
| 87 | __m512 va1 = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask, acc + d + 16), vscale); |
| 88 | __m512i vb = (__m512i)(_mm512_cvtne2ps_pbh(va1, va0)); |
| 89 | const __mmask32 vmask2 = (1ULL << remainder) - 1; |
| 90 | _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(out + d), vmask2, vb); |
| 91 | } |
| 92 | } |
| 93 | } |
| 94 | #endif |
| 95 | |
| 96 | template <typename scalar_t, int BLOCK_M, int BLOCK_N> |
| 97 | struct flash_attn_softmax { |
| 98 | static inline void apply( |
| 99 | float* __restrict__ s_i, |
| 100 | scalar_t* __restrict__ s_delta2, |
| 101 | float* __restrict__ v_prime, |
| 102 | float* __restrict__ s_prime, |
| 103 | float* __restrict__ m_prime, |
| 104 | int m_size, |
| 105 | int n_size, |
| 106 | int padded_n_size, |
| 107 | int head_size_v, |
| 108 | const float sm_scale, |
| 109 | int row) { |
| 110 | using Vec = at::vec::Vectorized<float>; |
| 111 | const Vec scale_vec = Vec(sm_scale); |
| 112 | float* s_delta = s_i; |
| 113 | // s_i <- s_i * scale |
| 114 | at::vec::map<float>([scale_vec](Vec x) { return x * scale_vec; }, s_i + row * BLOCK_N, s_i + row * BLOCK_N, n_size); |
| 115 | |
| 116 | // m_i: max value per row |
| 117 | float m_i = |
| 118 | at::vec::reduce_all<float>([](Vec& x, Vec& y) { return at::vec::maximum(x, y); }, s_i + row * BLOCK_N, n_size); |
| 119 | m_i = std::max(m_i, m_prime[row]); |
| 120 | |
| 121 | // m_delta <- exp(m' - m_i) |
| 122 | float m_delta = std::exp(m_prime[row] - m_i); |
| 123 | |
| 124 | // s_delta <- exp(s_i - m_i) |
| 125 | at::vec::map<float>( |
| 126 | [m_i](Vec x) { return (x - Vec(m_i)).fexp_u20(); }, s_delta + row * BLOCK_N, s_i + row * BLOCK_N, n_size); |
| 127 | |
| 128 | // s' <- s' * m_delta + sum(s_delta) |
| 129 | s_prime[row] *= m_delta; |
| 130 | s_prime[row] += at::vec::reduce_all<float>([](Vec& x, Vec& y) { return x + y; }, s_delta + row * BLOCK_N, n_size); |
| 131 | |
| 132 | m_prime[row] = m_i; |
| 133 | |
| 134 | // v' <- v' * m_delta |
| 135 | at::vec::map<float>( |
| 136 | [m_delta](Vec x) { return x * Vec(m_delta); }, |
| 137 | v_prime + row * head_size_v, |
| 138 | v_prime + row * head_size_v, |
| 139 | head_size_v); |
| 140 | |
| 141 | // Keep s_delta row-major for the following brgemm(P @ V), and only |
| 142 | // convert the columns that brgemm will consume. |
| 143 | fill_stub(s_delta + row * BLOCK_N + n_size, 0.f, padded_n_size - n_size); |
| 144 | copy_stub<scalar_t>(s_delta2 + row * BLOCK_N, s_delta + row * BLOCK_N, 1.f, padded_n_size); |
| 145 | } |
| 146 | }; |
| 147 | |
| 148 | #if defined(CPU_CAPABILITY_AVX512) |
| 149 | template <int BLOCK_M, int BLOCK_N> |
| 150 | struct flash_attn_softmax<at::BFloat16, BLOCK_M, BLOCK_N> { |
| 151 | static inline void apply( |
| 152 | float* __restrict__ s_i, |
| 153 | at::BFloat16* __restrict__ s_delta2, |
| 154 | float* __restrict__ v_prime, |
| 155 | float* __restrict__ s_prime, |
| 156 | float* __restrict__ m_prime, |
| 157 | int m_size, |
| 158 | int n_size, |
| 159 | int padded_n_size, |
| 160 | int head_size_v, |
| 161 | const float sm_scale, |
| 162 | int row) { |
| 163 | float* s_delta = s_i; |
| 164 | const __m512 vscale = _mm512_set1_ps(sm_scale); |
| 165 | |
| 166 | int n_remainder = n_size & 15; // 0xF |
| 167 | const __mmask16 vmask = (1ULL << n_remainder) - 1; |
| 168 | |
| 169 | int v_remainder = head_size_v & 15; // 0xF |
| 170 | const __mmask16 vmask1 = (1ULL << v_remainder) - 1; |
| 171 | |
| 172 | constexpr float NEG_INF = -std::numeric_limits<float>::infinity(); |
| 173 | |
| 174 | __m512 va; |
| 175 | __m256i vb; |
| 176 | __m512 vmax; |
| 177 | __m512 vsum; |
| 178 | __m512 vmdelta; |
| 179 | |
| 180 | const __m512 vneg_inf = _mm512_set1_ps(NEG_INF); |
| 181 | |
| 182 | int m = row; |
| 183 | vmax = vneg_inf; |
| 184 | |
| 185 | // s_i <- s_i * scale |
| 186 | int n = 0; |
| 187 | for (; n <= n_size - 16; n += 16) { |
| 188 | va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); |
| 189 | vmax = _mm512_max_ps(va, vmax); |
| 190 | } |
| 191 | if (n_remainder > 0) { |
| 192 | va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); |
| 193 | vmax = _mm512_max_ps(va, vmax); |
| 194 | } |
| 195 | |
| 196 | // m_i: max value per row |
| 197 | float m_i = _mm512_reduce_max_ps(vmax); |
| 198 | m_i = std::max(m_i, m_prime[m]); |
| 199 | vmax = _mm512_set1_ps(m_i); |
| 200 | |
| 201 | // m_delta <- exp(m' - m_i) |
| 202 | float m_delta = std::exp(m_prime[m] - m_i); |
| 203 | |
| 204 | // s_delta <- exp(s_i - m_i) |
| 205 | vsum = _mm512_setzero_ps(); |
| 206 | for (n = 0; n <= n_size - 16; n += 16) { |
| 207 | va = _mm512_mul_ps(_mm512_loadu_ps(s_i + m * BLOCK_N + n), vscale); |
| 208 | va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); |
| 209 | vsum = _mm512_add_ps(vsum, va); |
| 210 | |
| 211 | vb = (__m256i)(_mm512_cvtneps_pbh(va)); |
| 212 | _mm256_storeu_si256(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vb); |
| 213 | } |
| 214 | if (n_remainder > 0) { |
| 215 | va = _mm512_mul_ps(_mm512_mask_loadu_ps(vneg_inf, vmask, s_i + m * BLOCK_N + n), vscale); |
| 216 | va = _mm512_fexp_u20_ps(_mm512_sub_ps(va, vmax)); |
| 217 | vsum = _mm512_add_ps(vsum, va); |
| 218 | |
| 219 | vb = (__m256i)(_mm512_cvtneps_pbh(va)); |
| 220 | _mm256_mask_storeu_epi16(reinterpret_cast<__m256i*>(s_delta2 + m * BLOCK_N + n), vmask, vb); |
| 221 | } |
| 222 | |
| 223 | // s' <- s' * m_delta + sum(s_delta) |
| 224 | s_prime[m] *= m_delta; |
| 225 | s_prime[m] += _mm512_reduce_add_ps(vsum); |
| 226 | |
| 227 | m_prime[m] = m_i; |
| 228 | |
| 229 | // pad s_delta with 0, pad_size range from [0, 32) |
| 230 | int pad_size = padded_n_size - n_size; |
| 231 | if (pad_size > 0) { |
| 232 | const __m512i vzero = _mm512_setzero_si512(); |
| 233 | __mmask32 vmask2 = (1ULL << pad_size) - 1; |
| 234 | _mm512_mask_storeu_epi16(reinterpret_cast<__m512i*>(s_delta2 + m * BLOCK_N + n_size), vmask2, vzero); |
| 235 | } |
| 236 | |
| 237 | // v' <- v' * m_delta |
| 238 | vmdelta = _mm512_set1_ps(m_delta); |
| 239 | int k = 0; |
| 240 | for (; k <= head_size_v - 16; k += 16) { |
| 241 | va = _mm512_mul_ps(_mm512_loadu_ps(v_prime + m * head_size_v + k), vmdelta); |
| 242 | _mm512_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), va); |
| 243 | } |
| 244 | if (v_remainder > 0) { |
| 245 | va = _mm512_mul_ps(_mm512_maskz_loadu_ps(vmask1, v_prime + m * head_size_v + k), vmdelta); |
| 246 | _mm512_mask_storeu_ps(reinterpret_cast<__m512*>(v_prime + m * head_size_v + k), vmask1, va); |
| 247 | } |
| 248 | } |
| 249 | }; |
| 250 | #endif |