replicate/flash-mla

1// Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/hopper/utils.h2 3#pragma once4 5#include <assert.h>6#include <stdint.h>7#include <stdlib.h>8 9#include <cuda_bf16.h>10 11#include <cute/tensor.hpp>12 13#include <cutlass/array.h>14#include <cutlass/cutlass.h>15#include <cutlass/numeric_conversion.h>16#include <cutlass/numeric_types.h>17 18////////////////////////////////////////////////////////////////////////////////////////////////////19 20namespace flash {21 22////////////////////////////////////////////////////////////////////////////////////////////////////23 24template<typename T>25struct MaxOp {26__device__ __forceinline__ T operator()(T const & x, T const & y) { return x > y ? x : y; }27};28 29template <>30struct MaxOp<float> {31// This is slightly faster32__device__ __forceinline__ float operator()(float const &x, float const &y) { return max(x, y); }33};34 35////////////////////////////////////////////////////////////////////////////////////////////////////36 37template<typename T>38struct SumOp {39__device__ __forceinline__ T operator()(T const & x, T const & y) { return x + y; }40};41 42////////////////////////////////////////////////////////////////////////////////////////////////////43 44template<int THREADS>45struct Allreduce {46    static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);47    template<typename T, typename Operator>48    static __device__ __forceinline__ T run(T x, Operator &op) {49        constexpr int OFFSET = THREADS / 2;50        x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));51        return Allreduce<OFFSET>::run(x, op);52    }53};54 55////////////////////////////////////////////////////////////////////////////////////////////////////56 57template<>58struct Allreduce<2> {59template<typename T, typename Operator>60static __device__ __forceinline__ T run(T x, Operator &op) {61    x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));62    return x;63}64};65 66////////////////////////////////////////////////////////////////////////////////////////////////////67 68template <bool zero_init=false, int wg_wait=0, bool arrive=true, bool commit=true, typename Tensor0, typename Tensor1, typename Tensor2, typename TiledMma>69__forceinline__ __device__ void gemm(TiledMma &tiled_mma, Tensor0 const &tCrA, Tensor1 const &tCrB, Tensor2 &tCrC) {70    constexpr bool Is_RS = !cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value;71    // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take const72    if constexpr (Is_RS) { cute::warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA)); }73    warpgroup_fence_operand(tCrC);74    if constexpr (arrive) {75        warpgroup_arrive();76    }77    if constexpr (zero_init) {78        tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;79        // Unroll the K mode manually to set scale D to 180        CUTLASS_PRAGMA_UNROLL81        for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {82            cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);83            tiled_mma.accumulate_ = GMMA::ScaleOut::One;84        }85    } else {86        // cute::gemm(tiled_mma, tCrA, tCrB, tCrC);87        // Unroll the K mode manually to set scale D to 188        CUTLASS_PRAGMA_UNROLL89        for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {90            cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);91            tiled_mma.accumulate_ = GMMA::ScaleOut::One;92        }93    }94    if constexpr (commit) {95        warpgroup_commit_batch();96    }97    if constexpr (wg_wait >= 0) { warpgroup_wait<wg_wait>(); }98    warpgroup_fence_operand(tCrC);99    if constexpr (Is_RS) { warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA)); }100}101 102////////////////////////////////////////////////////////////////////////////////////////////////////103 104// For SM80, convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))105// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))106template<bool Transposed=false, typename Layout0>107__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout0 acc_layout) {108    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90109        static_assert(decltype(size<0, 0>(acc_layout))::value == 2);110        static_assert(decltype(size<0, 1>(acc_layout))::value == 2);111        static_assert(decltype(rank(acc_layout))::value == 3);112        auto l = acc_layout;113        if constexpr (!Transposed) {114            return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)));115        } else {116             return make_layout(make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)), make_layout(get<0, 1>(l), get<1>(l)));117        }118 119    } else {  // SM80120        static_assert(decltype(size<0>(acc_layout))::value == 4);121        static_assert(decltype(rank(acc_layout))::value == 3);122        auto l = logical_divide(acc_layout, Shape<_2>{});  // ((2, 2), MMA_M, MMA_N)123        if constexpr (!Transposed) {124            return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));125        } else {126            return make_layout(make_layout(get<0, 0>(l), get<2>(l)), make_layout(get<0, 1>(l), get<1>(l)));127        }128    }129};130 131////////////////////////////////////////////////////////////////////////////////////////////////////132 133// For SM80, convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)134// if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8.135// For SM90, FP16/BF16, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, MMA_N))136// For SM90, FP8, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((4, 2, 2), MMA_M, (N / 32, MMA_N))137template<typename MMA_Traits, typename Layout0>138__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout0 acc_layout) {139    using X = Underscore;140    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90141        static_assert(decltype(size<0, 0>(acc_layout))::value == 2);142        static_assert(decltype(size<0, 1>(acc_layout))::value == 2);143        static_assert(decltype(rank(acc_layout))::value == 3);144        static_assert(decltype(rank(get<0>(acc_layout)))::value == 3);145        if constexpr (sizeof(typename MMA_Traits::ValTypeA) == 2) {146            auto l = logical_divide(get<0, 2>(acc_layout), Tile<_2>{});  // ((2, N / 16))147            return make_layout(make_layout(get<0, 0>(acc_layout), get<0, 1>(acc_layout), get<0, 0>(l)), get<1>(acc_layout), coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));148        } else {149            static_assert(sizeof(typename MMA_Traits::ValTypeA) == 1);150            static_assert(decltype(stride<0, 0>(acc_layout))::value == 1);151            static_assert(decltype(stride<0, 1>(acc_layout))::value == 2);152            auto l = logical_divide(get<0, 2>(acc_layout), Tile<Layout<Shape<_2, _2>>>{});  // (((2, 2), N / 32))153            // This combines the first two modes (<0, 0> and <0, 1>) into one mode.154            // Will require register shuffling later to be correct.155            return make_layout(make_layout(Layout<_4>{}, get<0, 0, 0>(l), get<0, 0, 1>(l)),156                               get<1>(acc_layout),157                               coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));  // ((4, 2, 2), MMA_M, N / 32 * MMA_N)158            // This combination is right but doesn't work with register shuffling.159            // return make_layout(make_layout(coalesce(make_layout(get<0, 0>(acc_layout), get<0, 0, 0>(l))), get<0, 1>(acc_layout), get<0, 0, 1>(l)),160            //                    get<1>(acc_layout),161            //                    coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));162        }163    } else {  // SM80164        static_assert(decltype(size<0>(acc_layout))::value == 4);165        static_assert(decltype(rank(acc_layout))::value == 3);166        constexpr int mma_shape_K = get<2>(typename MMA_Traits::Shape_MNK{});167        static_assert(mma_shape_K == 8 || mma_shape_K == 16);168        if constexpr (mma_shape_K == 8) {169            return acc_layout;170        } else {171            auto l = logical_divide(acc_layout, Shape<X, X, _2>{});  // (4, MMA_M, (2, MMA_N / 2)))172            return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));173        }174    }175};176 177////////////////////////////////////////////////////////////////////////////////////////////////////178 179template <typename To_type, typename Engine, typename Layout>180__forceinline__ __device__ auto convert_type(Tensor<Engine, Layout> const &tensor) {181    using From_type = typename Engine::value_type;182    constexpr int numel = decltype(size(tensor))::value;183    cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;184    // HACK: this requires tensor to be "contiguous"185    auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel> *>(tensor.data()));186    return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());187}188 189////////////////////////////////////////////////////////////////////////////////////////////////////190 191// Blocks until all but N previous cp.async.commit_group operations have committed.192// This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all193// (which is equivalent to commit_group then wait_group 0).194// Instead we just call cp.async.wait_group 0, which is slightly faster.195// https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113196template <int N>197CUTE_HOST_DEVICE198void cp_async_wait() {199#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)200    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));201#endif202}203 204////////////////////////////////////////////////////////////////////////////////////////////////////205 206template <bool Is_even_MN=true, bool Is_even_K=true, bool Clear_OOB_MN=false, bool Clear_OOB_K=true,207          typename TiledCopy, typename Engine0, typename Layout0, typename Engine1, typename Layout1,208          typename Engine2, typename Layout2, typename Engine3, typename Layout3>209__forceinline__ __device__ void copy(TiledCopy tiled_copy, Tensor<Engine0, Layout0> const &S,210                            Tensor<Engine1, Layout1> &D, Tensor<Engine2, Layout2> const &identity_MN,211                            Tensor<Engine3, Layout3> const &predicate_K, const int max_MN=0) {212    CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});213    CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});214    CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D));                     // MMA215    CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D));                     // MMA_M216    CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D));                     // MMA_K217    // There's no case where !Clear_OOB_K && Clear_OOB_MN218    static_assert(!(Clear_OOB_MN && !Clear_OOB_K));219    #pragma unroll220    for (int m = 0; m < size<1>(S); ++m) {221        if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {222            #pragma unroll223            for (int k = 0; k < size<2>(S); ++k) {224                if (Is_even_K || predicate_K(k)) {225                    cute::copy(tiled_copy, S(_, m, k), D(_, m, k));226                } else if (Clear_OOB_K) {227                    cute::clear(D(_, m, k));228                }229            }230        } else if (Clear_OOB_MN) {231            cute::clear(D(_, m, _));232        }233    }234}235 236////////////////////////////////////////////////////////////////////////////////////////////////////237 238}  // namespace flash239