1#include "attention_api.cuh"

2#include <cassert>

3#include <cmath>

4#include <cuda.h>

5#include <cuda_runtime.h>

6#include <iostream>

7#include <stdio.h>

8#include <cutlass/numeric_conversion.h>

9#include <cutlass/numeric_types.h>

10#include <cute/tensor.hpp>

11#include <torch/extension.h>

12#include <torch/python.h>

13#include <vector>

14

15#include "static_switch.h"

16#include "kernel_traits.h"

17#include "flash.h"

18#include "utils.h"

19

20namespace flash {

21

22using namespace cute;

23

24template <int kBlockM, int kBlockN, int kNWarps,typename Engine, typename Layout>

25inline __device__ void mask_within_nblock(Tensor<Engine, Layout> &tensor, const int m_block, const int nbi) {

26 // tensor has shape (nrow=(2, MMA_M), ncol=(2, MMA_N))

27 static_assert(Layout::rank == 2, "Only support 2D Tensor");

28 // NOTE:

29 // 确定一个MMA内的index也是一个难点

30 // (nrow=(2, MMA_M), ncol=(2, MMA_N))形如:

31 // T1.V0 T1.V1

32 // T1.V0 T1.V1

33 // 根据mma_tile的示意图来确定col和row值

34

35 // NOTE:

36 // 计算thread的处理范围, mask掉超出范围的部分

37 //

38 // NOTE:

39 // % 32表示32做组, 因为SM80_16x8x16_F32F16F16F32_TN _1_2_1中最大线程数id是32

40 // (lane_id % 4) * 2表示在哪个"颜色"的col(thread)中, *2是为了靠右(即处理的哪个value2)

41 // 因此col_idx_offset表示当前thread所处理的单个Atom中4列的哪列

42

43 // lane_id表示一个MMA tile中的"线程组"

44 const int lane_id = threadIdx.x % 32;

45 const int col_idx_offset = kBlockN * nbi + (lane_id % 4) * 2;

46

47 const int nrow_group = threadIdx.x / 32;

48 const int row_idx_offset = kBlockM * m_block + lane_id / 4 + nrow_group * 16 /* 2*8 */;

49 // (2, nrow), 2*8 for each

50 const int group_stride = kNWarps * 16;

51

52 #pragma unroll

53 for (int nj = 0; nj < size<1, 1>(tensor); ++nj) {

54 // SM80_16x8x16_F32F16F16F32_TN中的一组中, 一行4个线程处理8个value

55 const int col_idx_base = col_idx_offset + nj * 8;

56 #pragma unroll

57 for (int j = 0; j < size<1, 0>(tensor); ++j) {

58 // j用于计算value 1和value 2对应col

59 // col_idx最终表示当前thread所处理的value的列号

60 const int col_idx = col_idx_base + j;

61

62 // mask掉tCrS_rowcol中(QK后的结果)超出范围的部分

63 // 列号和行号对比

64

65 // Without the "make_coord" we get wrong results

66 // for nrow(2, MMA_M)

67 #pragma unroll

68 for (int mi = 0; mi < size<0, 0>(tensor); ++mi) {

69

70 #pragma unroll

71 for (int mj = 0; mj < size<0, 1>(tensor); ++mj) {

72 const int row_idx = row_idx_offset + mi * 8 + mj * group_stride;

73 if (col_idx > row_idx) {

74 tensor(make_coord(mi, mj), make_coord(j, nj)) = -INFINITY;

75 }

76 }

77

78 }

79

80 }

81 }

82}

83

84// NOTE: A矩阵已经在寄存器中的gemm封装

85template<typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3,

86 typename TiledMma, typename TiledCopy, typename ThrCopy>

87inline __device__ void gemm_A_in_regs(Tensor0 &tCrC, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tXsB,

88 TiledMma mma, TiledCopy s2r_copy_b,

89 ThrCopy s2r_thr_copy_b) {

90 // NOTE: 符合M N K描述: A[M, K] @ B[N, K] = C[M, N]

91 CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(tCrC)); // MMA_M

92 CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(tCrC)); // MMA_N

93 CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K

94 // NOTE: retile 成拷贝需要的大小

95 Tensor tXrB = s2r_thr_copy_b.retile_D(tCrB);

96 CUTE_STATIC_ASSERT_V(size<1>(tXsB) == size<1>(tXrB)); // N

97 cute::copy(s2r_copy_b, tXsB(_, _, _0{}), tXrB(_, _, _0{}));

98 #pragma unroll

99 for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {

100 if (k_block < size<2>(tCrA) - 1) {

101 cute::copy(s2r_copy_b, tXsB(_, _, k_block + 1), tXrB(_, _, k_block + 1));

102 }

103 cute::gemm(mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);

104 }

105}

106

107template<typename Tensor0, typename Tensor1,

108 typename Tensor2, typename Tensor3, typename Tensor4,

109 typename TiledMma, typename TiledCopyA, typename TiledCopyB,

110 typename ThrCopyA, typename ThrCopyB>

111inline __device__ void gemm_smem(Tensor0 &tCrC, Tensor1 &tCrA, Tensor2 &tCrB, Tensor3 const& tXsA,

112 Tensor4 const& tXsB, TiledMma mma,

113 TiledCopyA s2r_copy_a, TiledCopyB s2r_copy_b,

114 ThrCopyA s2r_thr_copy_a, ThrCopyB s2r_thr_copy_b) {

115 CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(tCrC)); // MMA_M

116 CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(tCrC)); // MMA_N

117 CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB)); // MMA_K

118 Tensor tXrA = s2r_thr_copy_a.retile_D(tCrA);

119 CUTE_STATIC_ASSERT_V(size<1>(tXsA) == size<1>(tXrA)); // M

120 Tensor tXrB = s2r_thr_copy_b.retile_D(tCrB);

121 CUTE_STATIC_ASSERT_V(size<1>(tXsB) == size<1>(tXrB)); // N

122 // NOTE: s -> reg

123 cute::copy(s2r_copy_a, tXsA(_, _, _0{}), tXrA(_, _, _0{}));

124 cute::copy(s2r_copy_b, tXsB(_, _, _0{}), tXrB(_, _, _0{}));

125 #pragma unroll

126 for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {

127 if (k_block < size<2>(tCrA) - 1) {

128 cute::copy(s2r_copy_a, tXsA(_, _, k_block + 1), tXrA(_, _, k_block + 1));

129 cute::copy(s2r_copy_b, tXsB(_, _, k_block + 1), tXrB(_, _, k_block + 1));

130 }

131 cute::gemm(mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);

132 }

133}

134

135// Blocks until all but N previous cp.async.commit_group operations have committed.

136// This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all

137// (which is equivalent to commit_group then wait_group 0).

138// Instead we just call cp.async.wait_group 0, which is slightly faster.

139// https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113

140template <int N>

141CUTE_HOST_DEVICE

142void cp_async_wait() {

143#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)

144 asm volatile("cp.async.wait_group %0;\n" :: "n"(N));

145#endif

146}

147

148// copy from S to D with tiled_copy

149// TODO: 需要支持causal模式的的跳过拷贝

150template <typename TiledCopy, typename Engine0, typename Layout0, typename Engine1, typename Layout1>

151inline __device__ void copy(TiledCopy tiled_copy, Tensor<Engine0, Layout0> const &S,

152 Tensor<Engine1, Layout1> &D) {

153 CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});

154 CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});

155 CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D)); // MMA

156 CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D)); // MMA_M

157 CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D)); // MMA_K

158

159 #pragma unroll

160 for (int m = 0; m < size<1>(S); ++m) {

161 // TODO: 原版处这里identity_MN是用来跳过大块的block的, predicate用于跳过block内的拷贝

162 // TODO: 添加predicate逻辑, 用于跳过无用拷贝

163 // if (get<0>(identity_MN(0, m, 0)) < max_MN)

164 #pragma unroll

165 for (int k = 0; k < size<2>(S); ++k) {

166 cute::copy(tiled_copy, S(_, m, k), D(_, m, k));

167 }

168 }

169}

170

171

172// Convert rowcol_layout from (nrow=(2, MMA_M), ncol=(2, MMA_N)) to ((2, 2, 2), MMA_M, MMA_N / 2)

173// if using m16n8k16, or to ((2, 2, 1), MMA_M, MMA_N) if using m16n8k8.

174template<typename MMA_traits, typename Layout>

175inline __device__ auto convert_layout_rowcol_Aregs(Layout rowcol_layout) {

176 using X = Underscore;

177 static_assert(decltype(size<0, 0>(rowcol_layout))::value == 2);

178 static_assert(decltype(size<1, 0>(rowcol_layout))::value == 2);

179 constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});

180 static_assert(mma_shape_K == 8 || mma_shape_K == 16);

181 constexpr int MMA_N_divisor = mma_shape_K == 8 ? 1 : 2;

182 auto l = logical_divide(rowcol_layout, Shape<X, Shape<X, Int<MMA_N_divisor>>>{}); // ((2, MMA_M), (2, (2, MMA_N / 2)))

183 // TD [2023-08-13]: Same error as above on Cutlass 3.2

184 // return make_layout(make_layout(get<1, 0>(l), get<0, 0>(l), get<1, 1, 0>(l)),

185 // get<0, 1>(l),

186 // get<1, 1, 1>(l));

187 return make_layout(make_layout(get<0>(get<1>(l)), get<0>(get<0>(l)), get<0>(get<1>(get<1>(l)))),

188 get<1>(get<0>(l)),

189 get<1>(get<1>(get<1>(l))));

190};

191

192

193// TODO: not work

194template <typename To_type, typename Engine, typename Layout>

195inline __device__ auto convert_type(Tensor<Engine, Layout> const &tensor) {

196 using From_type = typename Engine::value_type;

197 constexpr int numel = decltype(size(tensor))::value;

198 cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;

199 // HACK: this requires tensor to be "contiguous"

200 auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel> *>(tensor.data()));

201 return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());

202}

203

204// TODO:

205// https://github.com/NVIDIA/cutlass/issues/802

206// TODO: convert出来后数据是否在寄存器?

207template <typename Fragment>

208inline __device__ auto convert_type_f32_to_f16(Fragment const &acc_fp32) {

209 Tensor acc_fp16 = make_tensor<cute::half_t>(shape(acc_fp32));

210 {

211 Tensor acc_fp32x2 = recast< float2>(acc_fp32);

212 Tensor acc_fp16x2 = recast<__half2>(acc_fp16);

213 for (int i = 0; i < size(acc_fp32x2); ++i) { acc_fp16x2(i) = __float22half2_rn(acc_fp32x2(i)); }

214 }

215 return acc_fp16;

216}

217

218// Apply the exp to all the elements.

219template <bool Scale_max=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>

220inline __device__ void scale_apply_exp2(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> const &max, const float scale) {

221 static_assert(Layout0::rank == 2, "Only support 2D Tensor");

222 static_assert(Layout1::rank == 1, "Only support 1D Tensor");

223 CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));

224 #pragma unroll

225 for (int mi = 0; mi < size<0>(tensor); ++mi) {

226 // If max is -inf, then all elements must have been -inf (possibly due to masking).

227 // We don't want (-inf - (-inf)) since that would give NaN.

228 // If we don't have float around M_LOG2E the multiplication is done in fp64.

229 const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));

230 #pragma unroll

231 for (int ni = 0; ni < size<1>(tensor); ++ni) {

232 // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -

233 // max * log_2(e)) This allows the compiler to use the ffma

234 // instruction instead of fadd and fmul separately.

235 tensor(mi, ni) = expf(tensor(mi, ni) * scale - max_scaled);

236 }

237 }

238}

239

240

241

242// Convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))

243// TODO: 搞清楚经过convert_layout_acc_rowcol后(nrow=(2, MMA_M), ncol=(2, MMA_N))的数学含义

244// 形象的解释是把

245// T1.V0

246// T1.V1

247// T1.V0

248// T1.V1

249// 变为

250// T1.V0 T1.V1

251// T1.V0 T1.V1

252// 这样符合MMA tile的行列直觉

253template<typename Layout>

254inline __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {

255 static_assert(decltype(size<0>(acc_layout))::value == 4);

256 static_assert(decltype(rank(acc_layout))::value == 3);

257 auto l = logical_divide(acc_layout, Shape<_2>{}); // ((2, 2), MMA_M, MMA_N)

258 // TD [2023-08-13]: Idk why but get<0, 1>(l) doesn't work for Cutlass 3.2, I'm getting

259 // "int_tuple.hpp(74): error: conversion to inaccessible base class"

260 // return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));

261 return make_layout(make_layout(get<1>(get<0>(l)), get<1>(l)), make_layout(get<0>(get<0>(l)), get<2>(l)));

262};

263

264template<bool Is_first, typename Tensor0, typename Tensor1, typename Tensor2>

265inline __device__ void softmax_rescale_o(Tensor0 &tCrS_rowcol, Tensor1 &rRowMax, Tensor1 &rRowSum,

266 Tensor2 &tCrO, float softmax_scale_log2) {

267 // tCrS_rowcol与tCrS共享存储, 从当前QK分块的分数原地更新为指数权重。

268 // rRowMax/rRowSum保存每行状态, tCrO保存需要rescale的输出分子。

269

270 if (Is_first) {

271 // NOTE: 优化, 第一次softmax不需要rescale, 只需要记录分子, max, sum

272 reduce_max</*zero_init=*/true>(tCrS_rowcol, rRowMax);

273 flash::scale_apply_exp2(tCrS_rowcol, rRowMax, softmax_scale_log2);

274 reduce_sum(tCrS_rowcol, rRowSum);

275 } else {

276 // 记录上一次的max

277 Tensor rRowMax_prev = make_fragment_like(rRowMax);

278 cute::copy(rRowMax, rRowMax_prev);

279 // TODO: reduce的实现学习一下

280 // NOTE: 计算新max到rRowMax

281 // reduce_max包含步:

282 // 1. 求当前thread内max: 遍历

283 // 2. reduce thread间的max: 使用shift技巧reduce

284 reduce_max</*zero_init=*/false>(tCrS_rowcol, rRowMax);

285 // Reshape tCrO from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))

286 // 将tCrO转换成符合2D直觉的(nrow, ncol)的形状

287 Tensor tCrO_rowcol = make_tensor(tCrO.data(), flash::convert_layout_acc_rowcol(tCrO.layout()));

288 #pragma unroll

289 for (int mi = 0; mi < size(rRowMax); ++mi) {

290 // NOTE: 辅助变量: 当前max

291 float row_max = rRowMax(mi);

292 // NOTE: 计算旧score的rescale值

293 // NOTE: 因为QK(影响max)计算时没有考虑softmax_scale, 所以这里要补上

294 float rescale = expf((rRowMax_prev(mi) - row_max) * softmax_scale_log2);

295 // NOTE: rescale旧分母部分

296 rRowSum(mi) *= rescale;

297 // NOTE: 旧分子部分rescale

298 // tCrO_rowcol.shape = (nrow, ncol)

299 #pragma unroll

300 for (int ni = 0; ni < size<1>(tCrO_rowcol); ++ni) { tCrO_rowcol(mi, ni) *= rescale; }

301 }

302 // NOTE: 计算新分子部分: 对所有tCrS_rowcol进行rescale

303 flash::scale_apply_exp2(tCrS_rowcol, rRowMax, softmax_scale_log2);

304

305 // NOTE: 累加新分母

306 Tensor rRowSum_cur = make_fragment_like(rRowSum);

307 // NOTE:利用新分子来累加新分母

308 // 1. 线程内累加: 遍历

309 // 2. 线程间累加: 使用shift技巧reduce

310 reduce_sum(tCrS_rowcol, rRowSum_cur);

311 // NOTE: 新分母累加到旧分母

312 #pragma unroll

313 for (int mi = 0; mi < size(rRowSum); ++mi) { rRowSum(mi) += rRowSum_cur(mi); }

314 }

315};

316

317} // namespace flash

318

319void set_params_fprop(Flash_fwd_params &params,

320

321 // device pointers

322 const torch::Tensor q,

323 const torch::Tensor k,

324 const torch::Tensor v,

325 torch::Tensor out,

326

327 void *softmax_lse_d,

328 float softmax_scale,

329 bool is_causal) {

330

331 memset(&params, 0, sizeof(params));

332

333 params.bs = q.size(0);

334 params.head = q.size(1);

335 params.q_seqlen = q.size(2);

336 params.dim = q.size(3);

337

338 params.k_head = k.size(1);

339 params.k_seqlen = k.size(2);

340

341 params.bs_stride = q.stride(0);

342 params.head_stride = q.stride(1);

343 params.seqlen_stride = q.stride(2);

344 params.dim_stride = q.stride(3);

345

346 params.softmax_scale = softmax_scale;

347 // TODO: 使用log2做scale

348 params.softmax_scale_log2 = softmax_scale * M_LOG2E;

349 params.is_causal = is_causal;

350 params.is_bf16 = q.dtype() == torch::kBFloat16;

351

352 // LogSumExp save for backward

353 params.softmax_lse_ptr = softmax_lse_d;

354

355 // TODO: get ptr

356 params.q_ptr = q.data_ptr();

357 params.k_ptr = k.data_ptr();

358 params.v_ptr = v.data_ptr();

359 params.out_ptr = out.data_ptr();

360}

361

362

363// Shared Storage with Aligned addresses.

364template <class ElementType, class SmemLayoutQ, class SmemLayoutK, class SmemLayoutV>

365struct SharedStorage {

366 // TODO: Aligned的话smem的计算是否有问题

367 cute::array_aligned<ElementType, cute::cosize_v<SmemLayoutQ>> smem_q;

368 cute::array_aligned<ElementType, cute::cosize_v<SmemLayoutK>> smem_k;

369 cute::array_aligned<ElementType, cute::cosize_v<SmemLayoutV>> smem_v;

370};

371

372template <typename Kernel_traits, bool Is_causal=false, typename Params>

373__global__ void flash_attention_v2_cutlass_kernel(const Params params) {

374

375 using namespace cute;

376

377 // m block index

378 const int m_block = blockIdx.x;

379

380 // bs * head

381 const int base_id = blockIdx.y;

382

383 // TODO: 传入泛型

384 // NOTE: 小技巧

385 using Element = typename Kernel_traits::Element;

386 using ElementAccum = typename Kernel_traits::ElementAccum;

387 // using TiledMma = typename Kernel_traits::MMA;

388 using TiledMma = typename Kernel_traits::TiledMma;

389 using index_t = typename Kernel_traits::index_t;

390 using SmemLayoutQ = typename Kernel_traits::SmemLayoutQ;

391 using SmemLayoutK = typename Kernel_traits::SmemLayoutKV;

392 using SmemLayoutV = typename Kernel_traits::SmemLayoutKV;

393 using SmemLayoutVt = typename Kernel_traits::SmemLayoutVtransposed;

394 using SmemLayoutVtNoSwizzle = typename Kernel_traits::SmemLayoutVtransposedNoSwizzle;

395

396 constexpr int kNWarps = Kernel_traits::kNWarps;

397 constexpr int kBlockM = Kernel_traits::kBlockM;

398 constexpr int kBlockN = Kernel_traits::kBlockN;

399 constexpr int kHeadDim = Kernel_traits::kHeadDim;

400

401 // Shared memory.

402 extern __shared__ char smem_[];

403 using SharedStorage = SharedStorage<Element, SmemLayoutQ, SmemLayoutK, SmemLayoutV>;

404 SharedStorage &shared_storage = *reinterpret_cast<SharedStorage *>(smem_);

405

406 const int bs_head_offset = base_id * params.head_stride;

407

408 // TODO: base offset for MHA

409 // NOTE: convert C pointer to Tensor for convenience

410 Tensor mQ = make_tensor(

411 make_gmem_ptr(reinterpret_cast<Element *>(params.q_ptr) + bs_head_offset),

412 make_shape(params.q_seqlen, Int<kHeadDim>{}),

413 make_stride(Int<kHeadDim>{}, Int<1>{}));

414 Tensor mK = make_tensor(

415 make_gmem_ptr(reinterpret_cast<Element *>(params.k_ptr) + bs_head_offset),

416 make_shape(params.k_seqlen, Int<kHeadDim>{}),

417 make_stride(Int<kHeadDim>{}, Int<1>{}));

418 Tensor mV = make_tensor(

419 make_gmem_ptr(reinterpret_cast<Element *>(params.v_ptr) + bs_head_offset),

420 make_shape(params.k_seqlen, Int<kHeadDim>{}),

421 make_stride(Int<kHeadDim>{}, Int<1>{}));

422 Tensor mO = make_tensor(

423 make_gmem_ptr(reinterpret_cast<Element *>(params.out_ptr) + bs_head_offset),

424 make_shape(params.q_seqlen, Int<kHeadDim>{}),

425 make_stride(Int<kHeadDim>{}, Int<1>{}));

426 // TODO:

427 Tensor mLSE = make_tensor(

428 make_gmem_ptr(reinterpret_cast<ElementAccum *>(params.softmax_lse_ptr) + base_id * params.q_seqlen),

429 // Shape<Int<kBlockM>, Stride<_1>{}>{},

430 make_shape(params.q_seqlen),

431 make_stride(Int<1>{}));

432

433

434 // 加载Q, K, V分块

435 // (kBlockM, kHeadDim, num_tile_n)

436 Tensor gQ = local_tile(mQ, make_tile(Int<kBlockM>{}, Int<kHeadDim>{}), make_coord(m_block, _));

437

438 // (kBlockN, kHeadDim, num_tile_n)

439 // NOTE: loading流水线, 初次加载所需K, V

440 Tensor gK = local_tile(mK, make_tile(Int<kBlockN>{}, Int<kHeadDim>{}), make_coord(0, _));

441 Tensor gV = local_tile(mV, make_tile(Int<kBlockN>{}, Int<kHeadDim>{}), make_coord(0, _));

442

443 // 获取MMA抽象

444 TiledMma mma;

445 auto thr_mma = mma.get_slice(threadIdx.x);

446

447 // Construct SMEM tensors.

448 Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});

449 Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), SmemLayoutK{});

450 Tensor sV = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutV{});

451

452 // Tensor for V Transpose; used in GEMM-II.

453 Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVt{});

454 Tensor sVtNoSwizzle = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), SmemLayoutVtNoSwizzle{});

455

456 // NOTE: copy抽象

457 // NOTE: QKV gmem -> smem拷贝的抽象

458 typename Kernel_traits::GmemTiledCopyQKV copy_qkv;

459 auto thr_copy_qkv = copy_qkv.get_thread_slice(threadIdx.x);

460

461 // NOTE: 定义gmem -> smem拷贝的src, dst

462 Tensor tQgQ = thr_copy_qkv.partition_S(gQ(_, _, 0));

463 Tensor tQsQ = thr_copy_qkv.partition_D(sQ);

464 Tensor tKgK = thr_copy_qkv.partition_S(gK(_, _, 0));

465 Tensor tKsK = thr_copy_qkv.partition_D(sK);

466 Tensor tVgV = thr_copy_qkv.partition_S(gV(_, _, 0));

467 Tensor tVsV = thr_copy_qkv.partition_D(sV);

468

469

470 // NOTE: 定义smem -> reg拷贝的dst

471 // partition_fragment与partition类似, 只是返回的是寄存器表示

472 Tensor tCrQ = thr_mma.partition_fragment_A(sQ); // (MMA,MMA_M,MMA_K)

473 Tensor tCrK = thr_mma.partition_fragment_B(sK); // (MMA,MMA_N,MMA_K)

474 Tensor tCrVt = thr_mma.partition_fragment_B(sVtNoSwizzle); // (MMA, MMA_K,MMA_N)

475

476 //

477 // Copy Atom retiling

478 //

479

480 // TODO: 理解这里的atom retiling

481

482 // NOTE: 准备从smem加载Q, K, V到寄存器的Copy对象

483 auto s2r_copy_q = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtom{}, mma);

484 auto s2r_thr_copy_q = s2r_copy_q.get_thread_slice(threadIdx.x);

485 Tensor tXsQ = s2r_thr_copy_q.partition_S(sQ);

486

487 auto s2r_copy_k = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtom{}, mma);

488 auto s2r_thr_copy_k = s2r_copy_k.get_thread_slice(threadIdx.x);

489 Tensor tXsK = s2r_thr_copy_k.partition_S(sK);

490

491 // TODO: 拷贝时转置

492 // NOTE: smem->reg拷贝Vt

493 auto s2r_copy_v = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed{}, mma);

494 auto s2r_thr_copy_v = s2r_copy_v.get_thread_slice(threadIdx.x);

495 Tensor tXsVt = s2r_thr_copy_v.partition_S(sVt);

496

497 // NOTE: tQ/tK/tV/tC/tX表示线程分工, g/s/r表示gmem/smem/register

498 // retile只建立寄存器的Copy视图, 真正的数据加载由copy执行

499 // tKgK表示gmem中的K, 用作gmem->smem的src

500 // tKsK表示smem中的K, 用作gmem->smem的dst

501 // tXsK表示smem中的K, 用作smem->reg的src

502

503

504 // NOTE: make_identity_tensor创建只有形状的tensor用于拷贝

505 // 在copy时用于跳过整块的block

506

507 // // TODO: cQ等用在causal模式, 暂时无用

508 // // Construct identity layout for sQ and sK

509 // Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_M,BLK_K) -> (blk_m,blk_k)

510 // Tensor cKV = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK))); // (BLK_N,BLK_K) -> (blk_n,blk_k)

511

512 // // Repeat the partitioning with identity layouts

513 // Tensor tQcQ = thr_copy_qkv.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)

514 // Tensor tKVcKV = thr_copy_qkv.partition_S(cKV); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)

515

516 // 流水线加载初始Q, K

517 // 加载Q到smem

518 flash::copy(copy_qkv, tQgQ, tQsQ);

519 // 加载K到smem

520 flash::copy(copy_qkv, tKgK, tKsK);

521 // 开始执行异步拷贝

522 cute::cp_async_fence();

523

524 Tensor tCrO = partition_fragment_C(mma, Shape<Int<kBlockM>, Int<kHeadDim>>{});

525

526 // step1: slice-k compute QK block

527 // Q[BLOCK_M, BLOCK_N] @ K[BLOCK_M, BLOCK_N].T = O[BLOCK_M, BLOCK_M]

528 //

529 // step2:

530 // advance K, V

531

532 // NOTE: K, V分块的数量: 处理的区间

533 const int n_block_min = 0;

534 // NOTE: 1. mask between N BLOCKs if is causal mode

535 int seqlen_start = m_block * kBlockM;

536 int seqlen_end = (m_block + 1) * kBlockM;

537 int n_block_max = Is_causal ? cute::ceil_div(seqlen_end, kBlockN) : cute::ceil_div(params.k_seqlen, kBlockN);

538

539 // NOTE: 需要记录的max

540 Tensor rRowMax = make_tensor<ElementAccum>(Shape<Int<2 * size<1>(tCrO)>>{});

541 // NOTE: 需要记录的denom

542 Tensor rRowSum = make_fragment_like(rRowMax);

543

544 clear(tCrO);

545

546 for (int nbi = n_block_min; nbi < n_block_max; nbi++) {

547 auto tCrS = partition_fragment_C(mma, make_shape(Int<kBlockM>{}, Int<kBlockN>{}));

548

549 clear(tCrS);

550

551 // 等待Q, K的gmem -> smem拷贝完成, 即Q, K就绪

552 // wait<0>表示等待还剩0个未完成

553 flash::cp_async_wait<0>();

554 __syncthreads();

555

556 // gemm的同时异步加载V

557 gV = local_tile(mV, make_tile(Int<kBlockN>{}, Int<kHeadDim>{}), make_coord(nbi, _));

558 tVgV = thr_copy_qkv.partition_S(gV(_, _, 0));

559 // 异步加载V到smem

560 flash::copy(copy_qkv, tVgV, tVsV);

561 // 发起异步拷贝

562 cute::cp_async_fence();

563

564 // O = Q@K.T

565 // NOTE: 加载smem中的数据到reg再做gemm, **加载期间执行retile**

566 flash::gemm_smem(tCrS, tCrQ, tCrK, tXsQ, tXsK, mma, s2r_copy_q, s2r_copy_k,

567 s2r_thr_copy_q, s2r_thr_copy_k

568 );

569

570 Tensor tCrS_rowcol = make_tensor(tCrS.data(), flash::convert_layout_acc_rowcol(tCrS.layout()));

571

572 // NOTE: 2. mask within N BLOCKs

573 if (Is_causal == true && nbi * kBlockN >= seqlen_start) {

574 flash::mask_within_nblock<kBlockM, kBlockN, kNWarps>(tCrS_rowcol, m_block, nbi);

575 }

576

577 // NOTE: 等待V加载完成, 为下个K加载准备初始状态

578 flash::cp_async_wait<0>();

579 __syncthreads();

580

581 // advance K

582 if (nbi != n_block_max - 1) {

583 gK = local_tile(mK, make_tile(Int<kBlockN>{}, Int<kHeadDim>{}), make_coord(nbi + 1, _));

584 tKgK = thr_copy_qkv.partition_S(gK(_, _, 0));

585 flash::copy(copy_qkv, tKgK, tKsK);

586 cute::cp_async_fence();

587 }

588

589 // 计算softmax

590 // NOTE: tCrO记录softmax后所有的分子

591 nbi == 0 ? flash::softmax_rescale_o</*Is_first=*/true>(tCrS_rowcol, rRowMax, rRowSum, tCrO, params.softmax_scale) :

592 flash::softmax_rescale_o</*Is_first=*/false>(tCrS_rowcol, rRowMax, rRowSum, tCrO, params.softmax_scale);

593

594 // 实际执行QK @ V

595 // (score AKA tCrS): QK[M, N] @ V[N, dim]

596 // NOTE: DABC: F32F16F16F32, convert D type(F32) to A type(F16)

597 // TODO: convert_type目前写死

598 Tensor rP = flash::convert_type_f32_to_f16(tCrS);

599 // NOTE: Convert from layout C to layout A

600 Tensor tCrP = make_tensor(rP.data(), flash::convert_layout_rowcol_Aregs<TiledMma>(tCrS_rowcol.layout()));

601

602 flash::gemm_A_in_regs(tCrO, tCrP, tCrVt, tXsVt, mma, s2r_copy_v, s2r_thr_copy_v);

603 }

604

605 // Epilogue

606

607 // NOTE: 最后统一除上分母部分

608 // Reshape tCrO from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))

609 // AKA reshape to (nrow, ncol) but with specific MMA layout

610 Tensor tCrO_rowcol = make_tensor(tCrO.data(), flash::convert_layout_acc_rowcol(tCrO.layout()));

611 // NOTE: 保存lse给bwd

612 Tensor rLSE = make_fragment_like(rRowSum);

613 // for row

614 #pragma unroll

615 for (int mi = 0; mi < size<0>(tCrO_rowcol); ++mi) {

616 float sum = rRowSum(mi);

617 float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;

618 // compute rLSE

619 // NOTE: here we use max * scale

620 rLSE(mi) = (sum == 0.f || sum != sum) ? INFINITY : rRowMax(mi) * params.softmax_scale + __logf(sum);

621 float scale = inv_sum;

622 // for col

623 #pragma unroll

624 for (int ni = 0; ni < size<1>(tCrO_rowcol); ++ni) {

625 tCrO_rowcol(mi, ni) *= scale;

626 }

627 }

628

629 // Convert tCrO from fp32 to fp16.

630 Tensor tCrO_fp16 = flash::convert_type_f32_to_f16(tCrO);

631 // 复用sQ的smem做sO的拷出

632 Tensor sO = make_tensor(sQ.data(), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N)

633

634 // Partition sO to match the accumulator partitioning

635 // TODO: review

636 auto r2s_copy_o = make_tiled_copy_C(typename Kernel_traits::SmemCopyAtomO{}, mma);

637 auto r2s_thr_copy_o = r2s_copy_o.get_thread_slice(threadIdx.x);

638 Tensor tXrO = r2s_thr_copy_o.retile_S(tCrO_fp16); // ((Atom,AtomNum), MMA_M, MMA_N)

639 Tensor tXsO = r2s_thr_copy_o.partition_D(sO); // ((Atom,AtomNum),PIPE_M,PIPE_N)

640

641 // NOTE: 先拷贝到smem

642 cute::copy(r2s_copy_o, tXrO, tXsO);

643

644 Tensor gO = local_tile(mO, make_tile(Int<kBlockM>{}, Int<kHeadDim>{}), make_coord(m_block, _));

645

646 // 创建到smem -> gmem的拷贝

647 typename Kernel_traits::GmemTiledCopyO copy_o;

648 auto thr_copy_o = copy_o.get_thread_slice(threadIdx.x);

649 Tensor tOsO = thr_copy_o.partition_S(sO); // ((Atom,AtomNum),ATOM_M,ATOM_N)

650 Tensor tOgO = thr_copy_o.partition_D(gO(_, _, 0));

651

652 __syncthreads();

653

654 // NOTE:: 再拷贝到gmem

655

656 // TODO: review, 这里两个copy的作用

657 Tensor tOrO = make_tensor<Element>(shape(tOgO));

658 cute::copy(copy_o, tOsO, tOrO);

659

660 flash::copy(copy_o, tOrO, tOgO);

661

662

663 // NOTE: 写回lse

664 Tensor gLSE = local_tile(mLSE, make_tile(Int<kBlockM>{}), make_coord(m_block));

665 Tensor cO = make_identity_tensor(Shape<Int<kBlockM>, Int<kHeadDim>>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k)

666 Tensor tCcO = thr_mma.partition_C(cO); // (MMA,MMA_M,MMA_K)

667 static_assert(decltype(size<0>(tCcO))::value == 4);

668 // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices.

669 // TODO: review this shape

670 Tensor tCcO_row = logical_divide(tCcO, Shape<_2>{})(make_coord(0, _), _, 0);

671 CUTE_STATIC_ASSERT_V(size(rLSE) == size(tCcO_row)); // MMA_M

672 // TODO: 搞清楚这里的逻辑

673 if (get<1>(tCcO_row(0)) == 0) {

674 #pragma unroll

675 for (int mi = 0; mi < size(rLSE); ++mi) {

676 const int row = get<0>(tCcO_row(mi));

677 // if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSE(row) = rLSE(mi); }

678 gLSE(row) = rLSE(mi);

679 }

680 }

681

682}

683

684template<typename Kernel_traits, bool Is_causal>

685void run_flash_fwd(Flash_fwd_params &params, cudaStream_t stream) {

686 // TODO: check if works: default stream = 0

687 using Element = typename Kernel_traits::Element;

688 using SmemLayoutQ = typename Kernel_traits::SmemLayoutQ;

689 using SmemLayoutK = typename Kernel_traits::SmemLayoutKV;

690 using SmemLayoutV = typename Kernel_traits::SmemLayoutKV;

691

692 const int num_m_block =

693 (params.q_seqlen + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;

694

695 dim3 grid(num_m_block, params.bs * params.head, 1);

696 dim3 block(Kernel_traits::kNThreads);

697

698 int smem_size = int(sizeof(SharedStorage<Element, SmemLayoutQ, SmemLayoutK, SmemLayoutV>));

699

700 auto kernel = &flash_attention_v2_cutlass_kernel<Kernel_traits, Is_causal, Flash_fwd_params>;

701 // NOTE: smem过大时需要设置

702 if (smem_size >= 48 * 1024) {

703 CUDA_ERROR_CHECK(cudaFuncSetAttribute(

704 kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));

705 }

706

707 // TODO: stream

708 kernel<<<grid, block, smem_size>>>(params);

709}

710

711template<typename T, int Headdim>

712void run_flash_fwd_(Flash_fwd_params &params, cudaStream_t stream);

713

714// TODO: 挨个写出特化, 目前使用通用模板

715// 如, run_flash_fwd_hdim32用于特化hdim=32

716// 这样做可以根据实际情况微调kBlockN和kBlockM的组合, 也可以加速编译

717template<typename T, int Headdim>

718void run_flash_fwd_(Flash_fwd_params &params, cudaStream_t stream) {

719 BOOL_SWITCH(params.is_causal, Is_causal, [&] {

720 // run_flash_fwd<Flash_fwd_kernel_traits<Headdim, /*kBlockM_=*/128, /*kBlockN_=*/128, /*kNWarps_=*/4, T>, Is_causal>(params, stream);

721

722 // TODO: kBlockM, kBlockN的组合

723 run_flash_fwd<Flash_fwd_kernel_traits<Headdim, /*kBlockM_=*/64, /*kBlockN_=*/64, /*kNWarps_=*/4, T>, Is_causal>(params, stream);

724 });

725}

726

727// entry point of flash attention

728void run_flash_attn_cutlass(Flash_fwd_params &params, cudaStream_t stream) {

729 // FP16_SWITCH yield elem_type namespace

730 FP16_SWITCH(!params.is_bf16, [&] {

731 // FWD_HEADDIM_SWITCH yield kHeadDim constexpr

732 FWD_HEADDIM_SWITCH(params.dim, [&] {

733 run_flash_fwd_<elem_type, kHeadDim>(params, stream);

734 });

735 });

736}

737

738std::vector<torch::Tensor> flash_attention_v2_cutlass(torch::Tensor q, torch::Tensor k,

739 torch::Tensor v, bool is_causal = false, float softmax_scale=1) {

740

741 CHECK_INPUT(q);

742 CHECK_INPUT(k);

743 CHECK_INPUT(v);

744

745 // batch size

746 int bs = q.size(0);

747 // head number

748 int head = q.size(1);

749 // seqlen

750 int seqlen = q.size(2);

751 // dim

752 int dim = q.size(3);

753 auto out = torch::empty_like(q);

754

755 auto opts = q.options();

756 auto softmax_lse = torch::empty({bs, head, seqlen}, opts.dtype(torch::kFloat32));

757

758 Flash_fwd_params params;

759 set_params_fprop(params, q, k, v, out,

760 softmax_lse.data_ptr(), softmax_scale, is_causal);

761

762 run_flash_attn_cutlass(params, 0);

763

764 // Wait until kernel finish.

765 cudaDeviceSynchronize();

766 CUDA_ERROR_CHECK(cudaGetLastError());

767

768 return {out, softmax_lse};

769}

770

771

772

773