File size: 2,554 Bytes
57c3a10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | /*
* Copyright (c) 2023 by FlashInfer team.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef FLASHINFER_STATE_CUH_
#define FLASHINFER_STATE_CUH_
#include "../math.cuh"
#include "../vec_dtypes.cuh"
namespace flashinfer {
/*!
* \brief The flashattention state.
* \tparam vec_size The size of the vector used in o.
*/
template <size_t vec_size>
struct state_t {
/* the weighted sum of v: exp(pre-softmax logit - m) * v / d */
vec_t<float, vec_size> o;
/* maximum value of pre-softmax logits */
float m;
/* sum of exp(pre-softmax logits - m) */
float d;
__device__ __forceinline__ void init() {
o.fill(0.f);
m = -math::inf;
d = 1.f;
}
__device__ __forceinline__ state_t() { init(); }
__device__ __forceinline__ float get_lse() const { return m + math::ptx_log2(d); }
/*!
* \brief Merge the state with another state.
* \param other_m The maximum value of pre-softmax logits of the other state.
* \param other_d The sum of exp(pre-softmax logits - m) of the other state.
* \param other_o The weighted sum of v of the other state.
*/
__device__ __forceinline__ void merge(const vec_t<float, vec_size>& other_o, float other_m,
float other_d) {
float m_prev = m, d_prev = d;
m = max(m_prev, other_m);
d = d_prev * math::ptx_exp2(m_prev - m) + other_d * math::ptx_exp2(other_m - m);
#pragma unroll
for (size_t i = 0; i < vec_size; ++i) {
o[i] = o[i] * math::ptx_exp2(m_prev - m) + other_o[i] * math::ptx_exp2(other_m - m);
}
}
/*!
* \brief Merge the state with another state.
* \param other The other state.
*/
__device__ __forceinline__ void merge(const state_t<vec_size>& other) {
merge(other.o, other.m, other.d);
}
__device__ __forceinline__ void normalize() {
// only normalize by d when not normalized on the fly
#pragma unroll
for (size_t i = 0; i < vec_size; ++i) {
o[i] = __fdividef(o[i], d);
}
}
};
} // namespace flashinfer
#endif // FLASHINFER_STATE_CUH_
|