replicate/flash-mla

1// Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/softmax.h2 3#pragma once4 5#include <cmath>6 7#include <cute/tensor.hpp>8#include <cutlass/numeric_types.h>9 10#include "utils.h"11 12namespace flash {13 14using namespace cute;15 16////////////////////////////////////////////////////////////////////////////////////////////////////17 18template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>19__device__ __forceinline__ void thread_reduce_(Tensor<Engine0, Layout0> const &tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {20    static_assert(Layout0::rank == 2, "Only support 2D Tensor");21    static_assert(Layout1::rank == 1, "Only support 1D Tensor");22    CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));23    #pragma unroll24    for (int mi = 0; mi < size<0>(tensor); mi++) {25        summary(mi) = zero_init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));26        #pragma unroll27        for (int ni = 1; ni < size<1>(tensor); ni++) {28            summary(mi) = op(summary(mi), tensor(mi, ni));29        }30    }31}32 33template<typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>34__device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0> &dst, Tensor<Engine1, Layout1> &src, Operator &op) {35    CUTE_STATIC_ASSERT_V(size(dst) == size(src));36    #pragma unroll37    for (int i = 0; i < size(dst); i++){38        dst(i) = Allreduce<4>::run(src(i), op);39    }40}41 42template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>43__device__ __forceinline__ void reduce_(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &summary, Operator &op) {44    thread_reduce_<zero_init>(tensor, summary, op);45    quad_allreduce_(summary, summary, op);46}47 48template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>49__device__ __forceinline__ void reduce_max(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &max){50    MaxOp<float> max_op;51    reduce_<zero_init>(tensor, max, max_op);52}53 54template<bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>55__device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tensor, Tensor<Engine1, Layout1> &sum){56    SumOp<float> sum_op;57    thread_reduce_<zero_init>(tensor, sum, sum_op);58}59 60// Apply the exp to all the elements.61template <bool Scale_max=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>62__forceinline__ __device__ auto scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> const &max, const float scale) {63    static_assert(Layout0::rank == 2, "Only support 2D Tensor");64    static_assert(Layout1::rank == 1, "Only support 1D Tensor");65    CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));66    #pragma unroll67    for (int mi = 0; mi < size<0>(tensor); ++mi) {68        // If max is -inf, then all elements must have been -inf (possibly due to masking).69        // We don't want (-inf - (-inf)) since that would give NaN.70        // If we don't have float around M_LOG2E the multiplication is done in fp64.71        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));72        #pragma unroll73        for (int ni = 0; ni < size<1>(tensor); ++ni)  {74            // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -75            // max * log_2(e)) This allows the compiler to use the ffma76            // instruction instead of fadd and fmul separately.77            // The following macro will disable the use of fma.78            // See: https://github.com/pytorch/pytorch/issues/121558 for more details79            // This macro is set in PyTorch and not FlashAttention80            #ifdef UNFUSE_FMA81                tensor(mi, ni) = exp2f(__fmul_rn(tensor(mi, ni), scale) - max_scaled);82            #else83                tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);84            #endif85        }86    }87    return tensor;88}89 90// Apply the exp to all the elements.91template <bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>92__forceinline__ __device__ void max_scale_exp2_sum(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> &max, Tensor<Engine1, Layout1> &sum, const float scale) {93    static_assert(Layout0::rank == 2, "Only support 2D Tensor");94    static_assert(Layout1::rank == 1, "Only support 1D Tensor");95    CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));96    #pragma unroll97    for (int mi = 0; mi < size<0>(tensor); ++mi) {98        MaxOp<float> max_op;99        max(mi) = zero_init ? tensor(mi, 0) : max_op(max(mi), tensor(mi, 0));100        #pragma unroll101        for (int ni = 1; ni < size<1>(tensor); ni++) {102            max(mi) = max_op(max(mi), tensor(mi, ni));103        }104        max(mi) = Allreduce<4>::run(max(mi), max_op);105        // If max is -inf, then all elements must have been -inf (possibly due to masking).106        // We don't want (-inf - (-inf)) since that would give NaN.107        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * scale;108        sum(mi) = 0;109        #pragma unroll110        for (int ni = 0; ni < size<1>(tensor); ++ni)  {111            // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -112            // max * log_2(e)) This allows the compiler to use the ffma113            // instruction instead of fadd and fmul separately.114            tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);115            sum(mi) += tensor(mi, ni);116        }117        SumOp<float> sum_op;118        sum(mi) = Allreduce<4>::run(sum(mi), sum_op);119    }120}121 122template<typename Tensor0, typename Tensor1>123__forceinline__ __device__ void rescale_o(Tensor0 &acc_o, Tensor1 &scale_o) {124    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))125    Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));126    #pragma unroll127    for (int mi = 0; mi < size(scale_o); ++mi) {128        #pragma unroll129        for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scale_o(mi); }130    }131}132 133////////////////////////////////////////////////////////////////////////////////////////////////////134 135template <int kNRows>136struct Softmax {137 138    using TensorT = decltype(make_tensor<float>(Shape<Int<kNRows>>{}));139    TensorT row_max, row_sum;140 141    __forceinline__ __device__ Softmax() {};142 143    template<bool Is_first, bool Check_inf=false, typename Tensor0>144    __forceinline__ __device__ TensorT softmax(Tensor0 &acc_s, float softmax_scale_log2) {145        // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))146        Tensor scores = make_tensor(acc_s.data(), flash::convert_layout_acc_rowcol(acc_s.layout()));147        static_assert(decltype(size<0>(scores))::value == kNRows);148        TensorT scale_o;149        clear(scale_o);150        if (Is_first) {151            flash::template reduce_max</*zero_init=*/true>(scores, row_max);152            flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);153            flash::reduce_sum</*zero_init=*/true>(scores, row_sum);154        } else {155            Tensor scores_max_prev = make_fragment_like(row_max);156            cute::copy(row_max, scores_max_prev);157            flash::template reduce_max</*zero_init=*/false>(scores, row_max);158            // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))159            #pragma unroll160            for (int mi = 0; mi < size(row_max); ++mi) {161                float scores_max_cur = !Check_inf162                    ? row_max(mi)163                    : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));164                float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);165                scale_o(mi) = scores_scale;166                row_sum(mi) *= scores_scale;167            }168            flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);169            // We don't do the reduce across threads here since we don't need to use the row_sum.170            // We do that reduce at the end when we need to normalize the softmax.171            flash::reduce_sum</*zero_init=*/false>(scores, row_sum);172        }173        return scale_o;174    };175 176    template<bool Is_dropout=false, bool Split=false, typename Tensor0>177    __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0 &acc_o, float softmax_scale, float rp_dropout=1.0) {178        SumOp<float> sum_op;179        quad_allreduce_(row_sum, row_sum, sum_op);180        TensorT lse = make_fragment_like(row_sum);181        // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))182        Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));183        static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);184        #pragma unroll185        for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) {186            float sum = row_sum(mi);187            float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;188            lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum);189            float scale = !Is_dropout ? inv_sum : inv_sum * rp_dropout;190            #pragma unroll191            for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scale; }192        }193        return lse;194    };195};196 197}  // namespace flash198