1#pragma once

2

3#include "cute/algorithm/copy.hpp"

4

5#include "cutlass/cutlass.h"

6#include "cutlass/layout/layout.h"

7#include <cutlass/numeric_types.h>

8

9using namespace cute;

10

11template<int kHeadDim_, int kBlockM_, int kBlockN_, int kNWarps_, typename elem_type=cutlass::half_t>

12struct Flash_kernel_traits {

13

14#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800

15 using Element = elem_type;

16 static constexpr bool Has_cp_async = true;

17#else

18 using Element = cutlass::half_t;

19 static constexpr bool Has_cp_async = false;

20#endif

21

22 using ElementAccum = float;

23 using index_t = uint32_t;

24

25#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800

26 using MMA_Atom_Arch = std::conditional_t<

27 std::is_same_v<elem_type, cutlass::half_t>,

28 MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,

29 MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>

30 >;

31 using ValLayoutMNK = Layout<Shape<_1, _2, _1>>;

32#else

33 using MMA_Atom_Arch = MMA_Atom<SM75_16x8x8_F32F16F16F32_TN>;

34 using ValLayoutMNK = Layout<Shape<_1, _2, _2>>;

35#endif

36

37#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750

38 using SmemCopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, elem_type>;

39 using SmemCopyAtomTransposed = Copy_Atom<SM75_U16x8_LDSM_T, elem_type>;

40#else

41 using SmemCopyAtom = Copy_Atom<DefaultCopy, elem_type>;

42 using SmemCopyAtomTransposed = Copy_Atom<DefaultCopy, elem_type>;

43#endif

44};

45

46

47// If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true

48template<int kHeadDim_, int kBlockM_, int kBlockN_, int kNWarps_, typename elem_type=cutlass::half_t,

49 typename Base=Flash_kernel_traits<kHeadDim_, kBlockM_, kBlockN_, kNWarps_, elem_type> >

50struct Flash_fwd_kernel_traits : public Base {

51 using Element = typename Base::Element;

52 using ElementAccum = typename Base::ElementAccum;

53 using index_t = typename Base::index_t;

54 static constexpr bool Has_cp_async = Base::Has_cp_async;

55 using SmemCopyAtom = typename Base::SmemCopyAtom;

56 using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed;

57

58 // The number of threads.

59 static constexpr int kNWarps = kNWarps_;

60 static constexpr int kNThreads = kNWarps * 32;

61

62 static constexpr int kBlockM = kBlockM_;

63 static constexpr int kBlockN = kBlockN_;

64 static constexpr int kHeadDim = kHeadDim_;

65

66 // TODO: review

67 static_assert(kHeadDim % 32 == 0);

68 static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32;

69 static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32);

70 static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3;

71

72 using TiledMma = TiledMMA<

73 typename Base::MMA_Atom_Arch,

74 Layout<Shape<Int<kNWarps>,_1,_1>>, // 4x1x1 or 8x1x1 thread group

75 // NOTE: cutlass v3.3

76 // typename Base::ValLayoutMNK>; // 1x2x1 or 1x2x2 value group for 16x16x16 MMA and LDSM

77 // cutlass v3.4

78 Tile<Int<16 * kNWarps>, _16, _16>>;

79

80 using SmemLayoutAtomQ = decltype(

81 composition(Swizzle<kSwizzle, 3, 3>{},

82 // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128

83 Layout<Shape<_8, Int<kBlockKSmem>>,

84 Stride<Int<kBlockKSmem>, _1>>{}));

85 using SmemLayoutQ = decltype(tile_to_shape(

86 SmemLayoutAtomQ{},

87 Shape<Int<kBlockM>, Int<kHeadDim>>{}));

88

89 using SmemLayoutKV = decltype(tile_to_shape(

90 SmemLayoutAtomQ{},

91 Shape<Int<kBlockN>, Int<kHeadDim>>{}));

92

93 // This has to be kBlockN and not 8, otherwise we get wrong results for d=128

94 using SmemLayoutAtomVtransposedNoSwizzle = Layout<Shape<Int<kBlockKSmem>, Int<kBlockN>>,

95 Stride<_1, Int<kBlockKSmem>>>;

96 using SmemLayoutAtomVtransposed = decltype(

97 composition(Swizzle<kSwizzle, 3, 3>{}, SmemLayoutAtomVtransposedNoSwizzle{}));

98 using SmemLayoutVtransposed = decltype(tile_to_shape(

99 SmemLayoutAtomVtransposed{},

100 Shape<Int<kHeadDim>, Int<kBlockN>>{}));

101 // Maybe the VtransposeNoSwizzle just needs to have the right shape

102 // And the strides don't matter?

103 using SmemLayoutVtransposedNoSwizzle = decltype(tile_to_shape(

104 SmemLayoutAtomVtransposedNoSwizzle{},

105 Shape<Int<kHeadDim>, Int<kBlockN>>{}));

106 // using SmemLayoutVtransposedNoSwizzle = decltype(SmemLayoutVtransposed{}.layout_fn());

107

108 using SmemLayoutAtomO = decltype(

109 composition(Swizzle<kSwizzle, 3, 3>{},

110 Layout<Shape<Int<8>, Int<kBlockKSmem>>,

111 Stride<Int<kBlockKSmem>, _1>>{}));

112 using SmemLayoutO = decltype(tile_to_shape(

113 SmemLayoutAtomO{},

114 Shape<Int<kBlockM>, Int<kHeadDim>>{}));

115 using SmemCopyAtomO = Copy_Atom<DefaultCopy, Element>;

116 using SmemCopyAtomOaccum = Copy_Atom<DefaultCopy, ElementAccum>;

117

118 static constexpr int kSmemQCount = size(SmemLayoutQ{});

119 static constexpr int kSmemKVCount = size(SmemLayoutKV{}) * 2;

120 static constexpr int kSmemQSize = kSmemQCount * sizeof(Element);

121 static constexpr int kSmemKVSize = kSmemKVCount * sizeof(Element);

122 // TODO:

123 static constexpr int kSmemSize = kSmemQSize + kSmemKVSize;

124

125 static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element);

126 static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad");

127

128 // TODO: review

129

130 // Using kBlockKSmem here is 6-10% faster than kBlockKGmem for d=128 because of bank conflicts.

131 // For example, for d=128, smem is split into 2 "pages", each page takes care of columns

132 // 0-63 and 64-127. If we have 16 threads per row for gmem read, when we write to smem,

133 // thread 0 - 7 will write to the first page and thread 8 - 15 will write to the second page,

134 // to the same banks.

135 static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad;

136 static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow");

137 using GmemLayoutAtom = Layout<Shape <Int<kNThreads / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>,

138 Stride<Int<kGmemThreadsPerRow>, _1>>;

139

140 // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading

141 // from the same address by the same threadblock. This is slightly faster.

142 using Gmem_copy_struct = std::conditional_t<

143 Has_cp_async,

144 SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>,

145 DefaultCopy

146 >;

147 using GmemTiledCopyQKV = decltype(

148 make_tiled_copy(Copy_Atom<Gmem_copy_struct, Element>{},

149 GmemLayoutAtom{},

150 Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per read

151 using GmemTiledCopyO = decltype(

152 make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},

153 GmemLayoutAtom{},

154 Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per store

155 static constexpr int kGmemThreadsPerRowP = kBlockN / kGmemElemsPerLoad;

156 static_assert(kNThreads % kGmemThreadsPerRowP == 0, "kNThreads must be a multiple of kGmemThreadsPerRowP");

157 using GmemLayoutAtomP = Layout<Shape <Int<kNThreads / kGmemThreadsPerRowP>, Int<kGmemThreadsPerRowP>>,

158 Stride<Int<kGmemThreadsPerRowP>, _1>>;

159

160 using GmemTiledCopyP = decltype(

161 make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},

162 GmemLayoutAtomP{},

163 Layout<Shape<_1, _8>>{})); // Val layout, 8 vals per store

164

165 using GmemLayoutAtomOaccum = std::conditional_t<

166 kBlockKSmem == 32,

167 Layout<Shape <_16, _8>, // Thread layout, 8 threads per row

168 Stride< _8, _1>>,

169 Layout<Shape <_8, _16>, // Thread layout, 16 threads per row

170 Stride< _16, _1>>

171 >;

172 using GmemTiledCopyOaccum = decltype(

173 make_tiled_copy(Copy_Atom<DefaultCopy, ElementAccum>{},

174 GmemLayoutAtomOaccum{},

175 Layout<Shape < _1, _4>>{})); // Val layout, 4 vals per store

176 using GmemLayoutAtomRotcossin = GmemLayoutAtom;

177 using GmemTiledCopyRotcossin = decltype(

178 make_tiled_copy(Copy_Atom<UniversalCopy<uint64_t>, Element>{},

179 GmemLayoutAtomRotcossin{},

180 Layout<Shape < _1, _4>>{})); // Val layout, 4 vals per load

181 using GmemTiledCopyRotcossinCont = decltype(

182 make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},

183 GmemLayoutAtomRotcossin{},

184 Layout<Shape < _1, _8>>{})); // Val layout, 8 vals per load

185

186};

187