Update flash_attn_triton.py

#7
by Skylion007 - opened
Files changed (1) hide show
  1. flash_attn_triton.py +316 -590
flash_attn_triton.py CHANGED
@@ -1,23 +1,17 @@
1
  # Copyright 2022 MosaicML Examples authors
2
  # SPDX-License-Identifier: Apache-2.0
3
 
4
- """Triton implementation of Flash Attention.
5
-
6
- # Copyright (c) 2022, Tri Dao.
7
- #
8
- # Licensed under the Apache License, Version 2.0 (the "License");
9
- # you may not use this file except in compliance with the License.
10
- # You may obtain a copy of the License at
11
- #
12
- # http://www.apache.org/licenses/LICENSE-2.0
13
- #
14
- # Unless required by applicable law or agreed to in writing, software
15
- # distributed under the License is distributed on an "AS IS" BASIS,
16
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17
- # See the License for the specific language governing permissions and
18
- # limitations under the License.
19
 
20
  *Experimental* implementation of FlashAttention in Triton.
 
 
 
 
 
 
21
  We use the FlashAttention implementation from Phil Tillet a starting point.
22
  https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
23
 
@@ -33,87 +27,65 @@ Changes:
33
  small batch size * nheads.
34
 
35
  Caution:
 
 
 
36
  - If you plan to use headdim other than 64 and 128, you should test for race conditions
37
  (due to the Triton compiler), as done in tests/test_flash_attn.py
38
  "test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
39
  for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
40
  that there are none left for other head dimensions.
 
41
  Differences between this Triton version and the CUDA version:
42
  - Triton version doesn't support dropout.
43
- - Triton forward is generally faster than CUDA forward.
44
- - Triton backward is faster than CUDA backward when batch * nheads is small, and when headdim=64.
45
- It is slightly slower when headdim=128 and batch * nheads is large.
46
- - Triton version doesn't yet support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
 
47
  """
48
 
49
  import math
50
 
51
  import torch
52
- import triton # type: ignore (reportMissingImports)
53
- import triton.language as tl # type: ignore (reportMissingImports)
54
- from einops import repeat
55
-
56
 
57
- @triton.autotune(
58
- configs=[
59
- triton.Config({
60
- 'BLOCK_M': 128,
61
- 'BLOCK_N': 128
62
- },
63
- num_warps=8,
64
- num_stages=1),
65
- # This config has a race condition when EVEN_M == False, disabling it for now.
66
- # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
67
- ],
68
- key=[
69
- 'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
70
- 'BLOCK_HEADDIM'
71
- ])
72
- @triton.heuristics({
73
- 'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
74
- 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
75
- 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
76
- })
77
  @triton.jit
78
  def _fwd_kernel(
79
- Q,
80
- K,
81
- V,
82
- Bias,
83
- Out,
84
- Lse,
85
- TMP, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug
86
  softmax_scale,
87
- stride_qb,
88
- stride_qh,
89
- stride_qm,
90
- stride_kb,
91
- stride_kh,
92
- stride_kn,
93
- stride_vb,
94
- stride_vh,
95
- stride_vn,
96
- stride_bb,
97
- stride_bh,
98
- stride_bm,
99
- stride_ob,
100
- stride_oh,
101
- stride_om,
102
- nheads,
103
- seqlen_q,
104
- seqlen_k,
105
- seqlen_q_rounded,
106
- headdim,
107
- CACHE_KEY_SEQLEN_Q,
108
- CACHE_KEY_SEQLEN_K,
109
  BIAS_TYPE: tl.constexpr,
110
  IS_CAUSAL: tl.constexpr,
111
  BLOCK_HEADDIM: tl.constexpr,
112
- EVEN_M: tl.constexpr,
113
- EVEN_N: tl.constexpr,
114
- EVEN_HEADDIM: tl.constexpr,
115
- BLOCK_M: tl.constexpr,
116
- BLOCK_N: tl.constexpr,
117
  ):
118
  start_m = tl.program_id(0)
119
  off_hb = tl.program_id(1)
@@ -130,23 +102,17 @@ def _fwd_kernel(
130
  # Adding parenthesis around indexing might use int32 math instead of int64 math?
131
  # https://github.com/openai/triton/issues/741
132
  # I'm seeing a tiny bit of difference (5-7us)
133
- q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (
134
- offs_m[:, None] * stride_qm + offs_d[None, :])
135
- k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (
136
- offs_n[:, None] * stride_kn + offs_d[None, :])
137
- v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (
138
- offs_n[:, None] * stride_vn + offs_d[None, :])
139
  if BIAS_TYPE == 'vector':
140
  b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
141
  elif BIAS_TYPE == 'matrix':
142
- b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (
143
- offs_m[:, None] * stride_bm + offs_n[None, :])
144
- else:
145
- raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
146
  # initialize pointer to m and l
147
  t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
148
- lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
149
- m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
150
  acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
151
  # load q: it will stay in SRAM throughout
152
  # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
@@ -160,13 +126,10 @@ def _fwd_kernel(
160
  if EVEN_HEADDIM:
161
  q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
162
  else:
163
- q = tl.load(q_ptrs,
164
- mask=(offs_m[:, None] < seqlen_q) &
165
- (offs_d[None, :] < headdim),
166
  other=0.0)
167
  # loop over k, v and update accumulator
168
- end_n = seqlen_k if not IS_CAUSAL else tl.minimum(
169
- (start_m + 1) * BLOCK_M, seqlen_k)
170
  for start_n in range(0, end_n, BLOCK_N):
171
  start_n = tl.multiple_of(start_n, BLOCK_N)
172
  # -- compute qk ----
@@ -174,48 +137,37 @@ def _fwd_kernel(
174
  if EVEN_HEADDIM:
175
  k = tl.load(k_ptrs + start_n * stride_kn)
176
  else:
177
- k = tl.load(k_ptrs + start_n * stride_kn,
178
- mask=offs_d[None, :] < headdim,
179
- other=0.0)
180
  else:
181
  if EVEN_HEADDIM:
182
- k = tl.load(k_ptrs + start_n * stride_kn,
183
- mask=(start_n + offs_n)[:, None] < seqlen_k,
184
  other=0.0)
185
  else:
186
  k = tl.load(k_ptrs + start_n * stride_kn,
187
- mask=((start_n + offs_n)[:, None] < seqlen_k) &
188
- (offs_d[None, :] < headdim),
189
  other=0.0)
190
  qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
191
  qk += tl.dot(q, k, trans_b=True)
192
  # Trying to combine the two masks seem to make the result wrong
193
  if not EVEN_N: # Need to mask out otherwise the softmax is wrong
194
- qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0,
195
- float('-inf'))
196
  if IS_CAUSAL:
197
- qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0,
198
- float('-inf'))
199
  if BIAS_TYPE != 'none':
200
  if BIAS_TYPE == 'vector':
201
  if EVEN_N:
202
  bias = tl.load(b_ptrs + start_n).to(tl.float32)
203
  else:
204
- bias = tl.load(b_ptrs + start_n,
205
- mask=(start_n + offs_n) < seqlen_k,
206
- other=0.0).to(tl.float32)
207
  bias = bias[None, :]
208
  elif BIAS_TYPE == 'matrix':
209
  if EVEN_M & EVEN_N:
210
  bias = tl.load(b_ptrs + start_n).to(tl.float32)
211
  else:
212
  bias = tl.load(b_ptrs + start_n,
213
- mask=(offs_m[:, None] < seqlen_q) &
214
- ((start_n + offs_n)[None, :] < seqlen_k),
215
  other=0.0).to(tl.float32)
216
- else:
217
- raise ValueError(
218
- "BIAS_TYPE must be one of {'vector', 'matrix'}")
219
  # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
220
  # can then fuse the mult and add into an fma instruction. But if we have bias we need to
221
  # to multiply with softmax_scale here.
@@ -240,18 +192,14 @@ def _fwd_kernel(
240
  if EVEN_HEADDIM:
241
  v = tl.load(v_ptrs + start_n * stride_vn)
242
  else:
243
- v = tl.load(v_ptrs + start_n * stride_vn,
244
- mask=offs_d[None, :] < headdim,
245
- other=0.0)
246
  else:
247
  if EVEN_HEADDIM:
248
- v = tl.load(v_ptrs + start_n * stride_vn,
249
- mask=(start_n + offs_n)[:, None] < seqlen_k,
250
  other=0.0)
251
  else:
252
  v = tl.load(v_ptrs + start_n * stride_vn,
253
- mask=((start_n + offs_n)[:, None] < seqlen_k) &
254
- (offs_d[None, :] < headdim),
255
  other=0.0)
256
  p = p.to(v.dtype)
257
  acc_o += tl.dot(p, v)
@@ -273,9 +221,8 @@ def _fwd_kernel(
273
  lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
274
  tl.store(lse_ptrs, lse_i)
275
  # initialize pointers to output
276
- offs_n = tl.arange(0, BLOCK_HEADDIM)
277
- out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (
278
- offs_m[:, None] * stride_om + offs_n[None, :])
279
  if EVEN_M:
280
  if EVEN_HEADDIM:
281
  tl.store(out_ptrs, acc_o)
@@ -285,29 +232,17 @@ def _fwd_kernel(
285
  if EVEN_HEADDIM:
286
  tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
287
  else:
288
- tl.store(out_ptrs,
289
- acc_o,
290
- mask=(offs_m[:, None] < seqlen_q) &
291
- (offs_d[None, :] < headdim))
292
 
293
 
294
  @triton.jit
295
  def _bwd_preprocess_do_o_dot(
296
- Out,
297
- DO,
298
- Delta,
299
- stride_ob,
300
- stride_oh,
301
- stride_om,
302
- stride_dob,
303
- stride_doh,
304
- stride_dom,
305
- nheads,
306
- seqlen_q,
307
- seqlen_q_rounded,
308
- headdim,
309
- BLOCK_M: tl.constexpr,
310
- BLOCK_HEADDIM: tl.constexpr,
311
  ):
312
  start_m = tl.program_id(0)
313
  off_hb = tl.program_id(1)
@@ -317,54 +252,54 @@ def _bwd_preprocess_do_o_dot(
317
  offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
318
  offs_d = tl.arange(0, BLOCK_HEADDIM)
319
  # load
320
- o = tl.load(Out + off_b * stride_ob + off_h * stride_oh +
321
- offs_m[:, None] * stride_om + offs_d[None, :],
322
- mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
323
- other=0.0).to(tl.float32)
324
- do = tl.load(DO + off_b * stride_dob + off_h * stride_doh +
325
- offs_m[:, None] * stride_dom + offs_d[None, :],
326
- mask=(offs_m[:, None] < seqlen_q) &
327
- (offs_d[None, :] < headdim),
328
- other=0.0).to(tl.float32)
329
  delta = tl.sum(o * do, axis=1)
330
  # write-back
331
  tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
332
 
333
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
334
  @triton.jit
335
  def _bwd_kernel_one_col_block(
336
  start_n,
337
- Q,
338
- K,
339
- V,
340
- Bias,
341
- DO,
342
- DQ,
343
- DK,
344
- DV,
345
- LSE,
346
- D,
347
  softmax_scale,
348
- stride_qm,
349
- stride_kn,
350
- stride_vn,
351
- stride_bm,
352
- stride_dom,
353
- stride_dqm,
354
- stride_dkn,
355
- stride_dvn,
356
- seqlen_q,
357
- seqlen_k,
358
- headdim,
359
  ATOMIC_ADD: tl.constexpr,
360
  BIAS_TYPE: tl.constexpr,
361
  IS_CAUSAL: tl.constexpr,
362
  BLOCK_HEADDIM: tl.constexpr,
363
- EVEN_M: tl.constexpr,
364
- EVEN_N: tl.constexpr,
365
- EVEN_HEADDIM: tl.constexpr,
366
- BLOCK_M: tl.constexpr,
367
- BLOCK_N: tl.constexpr,
368
  ):
369
  # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
370
  begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
@@ -383,11 +318,19 @@ def _bwd_kernel_one_col_block(
383
  b_ptrs = Bias + offs_n
384
  elif BIAS_TYPE == 'matrix':
385
  b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
386
- else:
387
- raise ValueError("BIAS_TYPE must be one of {'vector', 'matrix'}")
388
  # initialize dv and dk
389
  dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
390
  dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
 
 
 
 
 
 
 
 
 
 
391
  # k and v stay in SRAM throughout
392
  # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
393
  # if we just call tl.load(k_ptrs), we get the wrong output!
@@ -403,13 +346,9 @@ def _bwd_kernel_one_col_block(
403
  k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
404
  v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
405
  else:
406
- k = tl.load(k_ptrs,
407
- mask=(offs_n[:, None] < seqlen_k) &
408
- (offs_d[None, :] < headdim),
409
  other=0.0)
410
- v = tl.load(v_ptrs,
411
- mask=(offs_n[:, None] < seqlen_k) &
412
- (offs_d[None, :] < headdim),
413
  other=0.0)
414
  # loop over rows
415
  num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
@@ -422,41 +361,33 @@ def _bwd_kernel_one_col_block(
422
  q = tl.load(q_ptrs)
423
  else:
424
  if EVEN_HEADDIM:
425
- q = tl.load(q_ptrs,
426
- mask=offs_m_curr[:, None] < seqlen_q,
427
- other=0.0)
428
  else:
429
- q = tl.load(q_ptrs,
430
- mask=(offs_m_curr[:, None] < seqlen_q) &
431
- (offs_d[None, :] < headdim),
432
- other=0.0)
433
  # recompute p = softmax(qk, dim=-1).T
434
  qk = tl.dot(q, k, trans_b=True)
435
  # Trying to combine the two masks seem to make the result wrong
436
  if not EVEN_N: # Need to mask out otherwise the softmax is wrong
437
- qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf'))
438
  if IS_CAUSAL:
439
- qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk,
440
- float('-inf'))
441
  if BIAS_TYPE != 'none':
 
442
  if BIAS_TYPE == 'vector':
443
  if EVEN_N:
444
  bias = tl.load(b_ptrs).to(tl.float32)
445
  else:
446
- bias = tl.load(b_ptrs, mask=offs_n < seqlen_k,
447
- other=0.0).to(tl.float32)
448
  bias = bias[None, :]
449
  elif BIAS_TYPE == 'matrix':
450
  if EVEN_M & EVEN_N:
451
  bias = tl.load(b_ptrs).to(tl.float32)
452
  else:
453
  bias = tl.load(b_ptrs,
454
- mask=(offs_m_curr[:, None] < seqlen_q) &
455
- (offs_n[None, :] < seqlen_k),
456
  other=0.0).to(tl.float32)
457
- else:
458
- raise ValueError(
459
- "BIAS_TYPE must be one of {'vector', 'matrix'}")
460
  qk = qk * softmax_scale + bias
461
  # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
462
  # Also wrong for headdim=64.
@@ -476,10 +407,8 @@ def _bwd_kernel_one_col_block(
476
  do = tl.load(do_ptrs)
477
  else:
478
  # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
479
- do = tl.load(do_ptrs,
480
- mask=(offs_m_curr[:, None] < seqlen_q) &
481
- (offs_d[None, :] < headdim),
482
- other=0.0)
483
  # if EVEN_M:
484
  # if EVEN_HEADDIM:
485
  # do = tl.load(do_ptrs)
@@ -511,48 +440,38 @@ def _bwd_kernel_one_col_block(
511
  # compute dk = dot(ds.T, q)
512
  dk += tl.dot(ds, q, trans_a=True)
513
  # compute dq
 
 
514
  if not ATOMIC_ADD:
515
  if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
516
- dq = tl.load(dq_ptrs, eviction_policy='evict_last')
517
  dq += tl.dot(ds, k)
518
- tl.store(dq_ptrs, dq, eviction_policy='evict_last')
519
  else:
520
  if EVEN_HEADDIM:
521
- dq = tl.load(dq_ptrs,
522
- mask=offs_m_curr[:, None] < seqlen_q,
523
- other=0.0,
524
- eviction_policy='evict_last')
525
  dq += tl.dot(ds, k)
526
- tl.store(dq_ptrs,
527
- dq,
528
- mask=offs_m_curr[:, None] < seqlen_q,
529
- eviction_policy='evict_last')
530
  else:
531
  dq = tl.load(dq_ptrs,
532
- mask=(offs_m_curr[:, None] < seqlen_q) &
533
- (offs_d[None, :] < headdim),
534
- other=0.0,
535
- eviction_policy='evict_last')
536
  dq += tl.dot(ds, k)
537
- tl.store(dq_ptrs,
538
- dq,
539
- mask=(offs_m_curr[:, None] < seqlen_q) &
540
- (offs_d[None, :] < headdim),
541
- eviction_policy='evict_last')
542
  else: # If we're parallelizing across the seqlen_k dimension
543
  dq = tl.dot(ds, k)
544
  if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
545
  tl.atomic_add(dq_ptrs, dq)
546
  else:
547
  if EVEN_HEADDIM:
548
- tl.atomic_add(dq_ptrs,
549
- dq,
550
- mask=offs_m_curr[:, None] < seqlen_q)
551
  else:
552
- tl.atomic_add(dq_ptrs,
553
- dq,
554
- mask=(offs_m_curr[:, None] < seqlen_q) &
555
- (offs_d[None, :] < headdim))
556
  # increment pointers
557
  dq_ptrs += BLOCK_M * stride_dqm
558
  q_ptrs += BLOCK_M * stride_qm
@@ -562,28 +481,8 @@ def _bwd_kernel_one_col_block(
562
  # write-back
563
  dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
564
  dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
565
- # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False,
566
- # if we just call tl.store(dv_ptrs), there's a race condition
567
- if EVEN_N & EVEN_M:
568
- if EVEN_HEADDIM:
569
- tl.store(dv_ptrs, dv)
570
- tl.store(dk_ptrs, dk)
571
- else:
572
- tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
573
- tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
574
- else:
575
- if EVEN_HEADDIM:
576
- tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
577
- tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
578
- else:
579
- tl.store(dv_ptrs,
580
- dv,
581
- mask=(offs_n[:, None] < seqlen_k) &
582
- (offs_d[None, :] < headdim))
583
- tl.store(dk_ptrs,
584
- dk,
585
- mask=(offs_n[:, None] < seqlen_k) &
586
- (offs_d[None, :] < headdim))
587
 
588
 
589
  def init_to_zero(name):
@@ -592,24 +491,8 @@ def init_to_zero(name):
592
 
593
  @triton.autotune(
594
  configs=[
595
- triton.Config(
596
- {
597
- 'BLOCK_M': 128,
598
- 'BLOCK_N': 128,
599
- 'SEQUENCE_PARALLEL': False
600
- },
601
- num_warps=8,
602
- num_stages=1,
603
- pre_hook=init_to_zero('DQ')),
604
- triton.Config(
605
- {
606
- 'BLOCK_M': 128,
607
- 'BLOCK_N': 128,
608
- 'SEQUENCE_PARALLEL': True
609
- },
610
- num_warps=8,
611
- num_stages=1,
612
- pre_hook=init_to_zero('DQ')),
613
  # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now
614
  # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4*
615
  # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
@@ -617,69 +500,37 @@ def init_to_zero(name):
617
  # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
618
  # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
619
  ],
620
- key=[
621
- 'CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL',
622
- 'BLOCK_HEADDIM'
623
- ],
 
 
 
 
624
  )
625
- @triton.heuristics({
626
- 'EVEN_M': lambda args: args['seqlen_q'] % args['BLOCK_M'] == 0,
627
- 'EVEN_N': lambda args: args['seqlen_k'] % args['BLOCK_N'] == 0,
628
- 'EVEN_HEADDIM': lambda args: args['headdim'] == args['BLOCK_HEADDIM'],
629
- })
630
  @triton.jit
631
  def _bwd_kernel(
632
- Q,
633
- K,
634
- V,
635
- Bias,
636
- DO,
637
- DQ,
638
- DK,
639
- DV,
640
- LSE,
641
- D,
642
  softmax_scale,
643
- stride_qb,
644
- stride_qh,
645
- stride_qm,
646
- stride_kb,
647
- stride_kh,
648
- stride_kn,
649
- stride_vb,
650
- stride_vh,
651
- stride_vn,
652
- stride_bb,
653
- stride_bh,
654
- stride_bm,
655
- stride_dob,
656
- stride_doh,
657
- stride_dom,
658
- stride_dqb,
659
- stride_dqh,
660
- stride_dqm,
661
- stride_dkb,
662
- stride_dkh,
663
- stride_dkn,
664
- stride_dvb,
665
- stride_dvh,
666
- stride_dvn,
667
- nheads,
668
- seqlen_q,
669
- seqlen_k,
670
- seqlen_q_rounded,
671
- headdim,
672
- CACHE_KEY_SEQLEN_Q,
673
- CACHE_KEY_SEQLEN_K,
674
  BIAS_TYPE: tl.constexpr,
675
  IS_CAUSAL: tl.constexpr,
676
  BLOCK_HEADDIM: tl.constexpr,
677
  SEQUENCE_PARALLEL: tl.constexpr,
678
- EVEN_M: tl.constexpr,
679
- EVEN_N: tl.constexpr,
680
- EVEN_HEADDIM: tl.constexpr,
681
- BLOCK_M: tl.constexpr,
682
- BLOCK_N: tl.constexpr,
683
  ):
684
  off_hb = tl.program_id(1)
685
  off_b = off_hb // nheads
@@ -700,72 +551,40 @@ def _bwd_kernel(
700
  if not SEQUENCE_PARALLEL:
701
  num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
702
  for start_n in range(0, num_block_n):
703
- _bwd_kernel_one_col_block(start_n,
704
- Q,
705
- K,
706
- V,
707
- Bias,
708
- DO,
709
- DQ,
710
- DK,
711
- DV,
712
- LSE,
713
- D,
714
- softmax_scale,
715
- stride_qm,
716
- stride_kn,
717
- stride_vn,
718
- stride_bm,
719
- stride_dom,
720
- stride_dqm,
721
- stride_dkn,
722
- stride_dvn,
723
- seqlen_q,
724
- seqlen_k,
725
- headdim,
726
- ATOMIC_ADD=False,
727
- BIAS_TYPE=BIAS_TYPE,
728
- IS_CAUSAL=IS_CAUSAL,
729
- BLOCK_HEADDIM=BLOCK_HEADDIM,
730
- EVEN_M=EVEN_M,
731
- EVEN_N=EVEN_N,
732
- EVEN_HEADDIM=EVEN_HEADDIM,
733
- BLOCK_M=BLOCK_M,
734
- BLOCK_N=BLOCK_N)
735
  else:
736
  start_n = tl.program_id(0)
737
- _bwd_kernel_one_col_block(start_n,
738
- Q,
739
- K,
740
- V,
741
- Bias,
742
- DO,
743
- DQ,
744
- DK,
745
- DV,
746
- LSE,
747
- D,
748
- softmax_scale,
749
- stride_qm,
750
- stride_kn,
751
- stride_vn,
752
- stride_bm,
753
- stride_dom,
754
- stride_dqm,
755
- stride_dkn,
756
- stride_dvn,
757
- seqlen_q,
758
- seqlen_k,
759
- headdim,
760
- ATOMIC_ADD=True,
761
- BIAS_TYPE=BIAS_TYPE,
762
- IS_CAUSAL=IS_CAUSAL,
763
- BLOCK_HEADDIM=BLOCK_HEADDIM,
764
- EVEN_M=EVEN_M,
765
- EVEN_N=EVEN_N,
766
- EVEN_HEADDIM=EVEN_HEADDIM,
767
- BLOCK_M=BLOCK_M,
768
- BLOCK_N=BLOCK_N)
769
 
770
 
771
  def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
@@ -776,8 +595,7 @@ def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
776
  assert v.shape == (batch, seqlen_k, nheads, d)
777
  assert d <= 128, 'FlashAttention only support head dimensions up to 128'
778
  assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
779
- assert q.dtype in [torch.float16,
780
- torch.bfloat16], 'Only support fp16 and bf16'
781
  assert q.is_cuda and k.is_cuda and v.is_cuda
782
  softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
783
 
@@ -796,86 +614,40 @@ def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
796
  else:
797
  raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
798
  ' or (seqlen_q, seqlen_k)')
799
- if bias.shape[:2] == (1, nheads):
800
- bias = repeat(bias, '1 h ... -> b h ...', b=batch)
801
- elif bias.shape[:2] == (batch, 1):
802
- bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
803
- elif bias.shape[:2] == (1, 1):
804
- bias = repeat(bias, '1 h ... -> b h ...', b=batch)
805
- bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
806
- assert bias.shape[:2] == (
807
- batch, nheads
808
- ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
809
- assert bias is not None # for type checking
810
- bias_strides = (bias.stride(0), bias.stride(1),
811
- bias.stride(2)) if has_bias else (0, 0, 0)
812
 
813
  seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
814
- lse = torch.empty((batch, nheads, seqlen_q_rounded),
815
- device=q.device,
816
- dtype=torch.float32)
817
- tmp = torch.empty((batch, nheads, seqlen_q_rounded),
818
- device=q.device,
819
- dtype=torch.float32)
820
  o = torch.empty_like(q)
821
 
822
  BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
823
- # BLOCK = 128
824
- # num_warps = 4 if d <= 64 else 8
825
- grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
826
- _fwd_kernel[grid]( # type: ignore
827
- q,
828
- k,
829
- v,
830
- bias,
831
- o,
832
- lse,
833
- tmp,
834
  softmax_scale,
835
- q.stride(0),
836
- q.stride(2),
837
- q.stride(1),
838
- k.stride(0),
839
- k.stride(2),
840
- k.stride(1),
841
- v.stride(0),
842
- v.stride(2),
843
- v.stride(1),
844
  *bias_strides,
845
- o.stride(0),
846
- o.stride(2),
847
- o.stride(1),
848
- nheads,
849
- seqlen_q,
850
- seqlen_k,
851
- seqlen_q_rounded,
852
- d,
853
- seqlen_q // 32,
854
- seqlen_k // 32, # key for triton cache (limit number of compilations)
855
  # Can't use kwargs here because triton autotune expects key to be args, not kwargs
856
  # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
857
- bias_type,
858
- causal,
859
- BLOCK_HEADDIM,
860
- # BLOCK_M=BLOCK, BLOCK_N=BLOCK,
861
- # num_warps=num_warps,
862
- # num_stages=1,
863
  )
864
  return o, lse, softmax_scale # softmax_scale could have been updated
865
 
866
 
867
- def _flash_attn_backward(do,
868
- q,
869
- k,
870
- v,
871
- o,
872
- lse,
873
- dq,
874
- dk,
875
- dv,
876
- bias=None,
877
- causal=False,
878
- softmax_scale=None):
879
  # Make sure that the last dimension is contiguous
880
  if do.stride(-1) != 1:
881
  do = do.contiguous()
@@ -894,23 +666,13 @@ def _flash_attn_backward(do,
894
  # delta = torch.zeros_like(lse)
895
 
896
  BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
897
- grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads)
898
- _bwd_preprocess_do_o_dot[grid]( # type: ignore
899
- o,
900
- do,
901
- delta,
902
- o.stride(0),
903
- o.stride(2),
904
- o.stride(1),
905
- do.stride(0),
906
- do.stride(2),
907
- do.stride(1),
908
- nheads,
909
- seqlen_q,
910
- seqlen_q_rounded,
911
- d,
912
- BLOCK_M=128,
913
- BLOCK_HEADDIM=BLOCK_HEADDIM,
914
  )
915
 
916
  has_bias = bias is not None
@@ -927,71 +689,32 @@ def _flash_attn_backward(do,
927
  else:
928
  raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
929
  ' or (seqlen_q, seqlen_k)')
930
- if bias.shape[:2] == (1, nheads):
931
- bias = repeat(bias, '1 h ... -> b h ...', b=batch)
932
- elif bias.shape[:2] == (batch, 1):
933
- bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
934
- elif bias.shape[:2] == (1, 1):
935
- bias = repeat(bias, '1 h ... -> b h ...', b=batch)
936
- bias = repeat(bias, 'b 1 ... -> b h ...', h=nheads)
937
- assert bias.shape[:2] == (
938
- batch, nheads
939
- ), f'First 2 dimensions of bias must be broadcastible to (batch, nheads) = ({batch, nheads}). Bias has shape: {bias.shape}'
940
- assert bias is not None # type checking
941
- bias_strides = (bias.stride(0), bias.stride(1),
942
- bias.stride(2)) if has_bias else (0, 0, 0)
943
 
944
  # BLOCK_M = 128
945
  # BLOCK_N = 64
946
  # num_warps = 4
947
- grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N'])
948
- if META['SEQUENCE_PARALLEL'] else 1, batch * nheads)
949
- _bwd_kernel[grid]( # type: ignore
950
- q,
951
- k,
952
- v,
953
- bias,
954
- do,
955
- dq_accum,
956
- dk,
957
- dv,
958
- lse,
959
- delta,
960
  softmax_scale,
961
- q.stride(0),
962
- q.stride(2),
963
- q.stride(1),
964
- k.stride(0),
965
- k.stride(2),
966
- k.stride(1),
967
- v.stride(0),
968
- v.stride(2),
969
- v.stride(1),
970
  *bias_strides,
971
- do.stride(0),
972
- do.stride(2),
973
- do.stride(1),
974
- dq_accum.stride(0),
975
- dq_accum.stride(2),
976
- dq_accum.stride(1),
977
- dk.stride(0),
978
- dk.stride(2),
979
- dk.stride(1),
980
- dv.stride(0),
981
- dv.stride(2),
982
- dv.stride(1),
983
- nheads,
984
- seqlen_q,
985
- seqlen_k,
986
- seqlen_q_rounded,
987
- d,
988
- seqlen_q // 32,
989
- seqlen_k // 32, # key for triton cache (limit number of compilations)
990
  # Can't use kwargs here because triton autotune expects key to be args, not kwargs
991
  # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
992
- bias_type,
993
- causal,
994
- BLOCK_HEADDIM,
995
  # SEQUENCE_PARALLEL=False,
996
  # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
997
  # num_warps=num_warps,
@@ -1000,31 +723,23 @@ def _flash_attn_backward(do,
1000
  dq.copy_(dq_accum)
1001
 
1002
 
1003
- class _FlashAttnQKVPackedFunc(torch.autograd.Function):
1004
 
1005
  @staticmethod
1006
  def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
1007
- """Forward pass for packed FlashAttention.
1008
-
1009
- Args:
1010
- ctx: autograd context
1011
  qkv: (batch, seqlen, 3, nheads, headdim)
1012
  bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
1013
  For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
1014
  ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
1015
- causal (bool): whether to incorporate causal attention masking
1016
- softmax_scale (float, optional): scale factor for softmax
1017
  """
1018
  # Make sure that the last dimension is contiguous
1019
  if qkv.stride(-1) != 1:
1020
  qkv = qkv.contiguous()
1021
  o, lse, ctx.softmax_scale = _flash_attn_forward(
1022
- qkv[:, :, 0],
1023
- qkv[:, :, 1],
1024
- qkv[:, :, 2],
1025
- bias=bias,
1026
- causal=causal,
1027
- softmax_scale=softmax_scale)
1028
  ctx.save_for_backward(qkv, o, lse, bias)
1029
  ctx.causal = causal
1030
  return o
@@ -1032,53 +747,75 @@ class _FlashAttnQKVPackedFunc(torch.autograd.Function):
1032
  @staticmethod
1033
  def backward(ctx, do):
1034
  qkv, o, lse, bias = ctx.saved_tensors
1035
- assert not ctx.needs_input_grad[
1036
- 1], 'FlashAttention does not support bias gradient yet'
1037
  # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
1038
  # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
1039
  with torch.inference_mode():
1040
  dqkv = torch.empty_like(qkv)
1041
- _flash_attn_backward(do,
1042
- qkv[:, :, 0],
1043
- qkv[:, :, 1],
1044
- qkv[:, :, 2],
1045
- o,
1046
- lse,
1047
- dqkv[:, :, 0],
1048
- dqkv[:, :, 1],
1049
- dqkv[:, :, 2],
1050
- bias=bias,
1051
- causal=ctx.causal,
1052
- softmax_scale=ctx.softmax_scale)
1053
  return dqkv, None, None, None
1054
 
1055
 
1056
- flash_attn_qkvpacked_func = _FlashAttnQKVPackedFunc.apply
1057
 
1058
 
1059
- class _FlashAttnFunc(torch.autograd.Function):
1060
 
1061
  @staticmethod
1062
- def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
1063
- """Forward pass for FlashAttention.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1064
 
1065
- Args:
1066
- ctx: autograd context
 
 
 
 
1067
  q: (batch_size, seqlen_q, nheads, headdim)
1068
- k: (batch_size, seqlen_k, nheads, headdim)
1069
- v: (batch_size, seqlen_k, nheads, headdim)
1070
  bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
1071
  For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
1072
  ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
1073
- causal (bool): whether to incorporate causal attention masking
1074
- softmax_scale (float, optional): scale factor for softmax
1075
  """
1076
  # Make sure that the last dimension is contiguous
1077
- q, k, v = [
1078
- x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]
1079
- ]
1080
  o, lse, ctx.softmax_scale = _flash_attn_forward(
1081
- q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale)
 
1082
  ctx.save_for_backward(q, k, v, o, lse, bias)
1083
  ctx.causal = causal
1084
  return o
@@ -1086,27 +823,16 @@ class _FlashAttnFunc(torch.autograd.Function):
1086
  @staticmethod
1087
  def backward(ctx, do):
1088
  q, k, v, o, lse, bias = ctx.saved_tensors
1089
- assert not ctx.needs_input_grad[
1090
- 3], 'FlashAttention does not support bias gradient yet'
1091
  # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
1092
  # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
1093
  with torch.inference_mode():
1094
  dq = torch.empty_like(q)
1095
  dk = torch.empty_like(k)
1096
  dv = torch.empty_like(v)
1097
- _flash_attn_backward(do,
1098
- q,
1099
- k,
1100
- v,
1101
- o,
1102
- lse,
1103
- dq,
1104
- dk,
1105
- dv,
1106
- bias=bias,
1107
- causal=ctx.causal,
1108
- softmax_scale=ctx.softmax_scale)
1109
  return dq, dk, dv, None, None, None
1110
 
1111
 
1112
- flash_attn_func = _FlashAttnFunc.apply
 
1
  # Copyright 2022 MosaicML Examples authors
2
  # SPDX-License-Identifier: Apache-2.0
3
 
4
+ """
5
+ Copied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py
6
+ update imports to use 'triton_pre_mlir'
 
 
 
 
 
 
 
 
 
 
 
 
7
 
8
  *Experimental* implementation of FlashAttention in Triton.
9
+ Tested with triton==2.0.0.dev20221202.
10
+ Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions
11
+ other than 64:
12
+ https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207
13
+ We'll update this implementation with the new Triton backend once this is fixed.
14
+
15
  We use the FlashAttention implementation from Phil Tillet a starting point.
16
  https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py
17
 
 
27
  small batch size * nheads.
28
 
29
  Caution:
30
+ - This is an *experimental* implementation. The forward pass should be quite robust but
31
+ I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler).
32
+ - This implementation has only been tested on A100.
33
  - If you plan to use headdim other than 64 and 128, you should test for race conditions
34
  (due to the Triton compiler), as done in tests/test_flash_attn.py
35
  "test_flash_attn_triton_race_condition". I've tested and fixed many race conditions
36
  for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident
37
  that there are none left for other head dimensions.
38
+
39
  Differences between this Triton version and the CUDA version:
40
  - Triton version doesn't support dropout.
41
+ - Triton forward is generally faster than CUDA forward, while Triton backward is
42
+ generally slower than CUDA backward. Overall Triton forward + backward is slightly slower
43
+ than CUDA forward + backward.
44
+ - Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).
45
+ - Triton version supports attention bias, while CUDA version doesn't.
46
  """
47
 
48
  import math
49
 
50
  import torch
 
 
 
 
51
 
52
+ import triton_pre_mlir as triton
53
+ import triton_pre_mlir.language as tl
54
+
55
+
56
+ # Disabling autotune for now, set num_warps=4 if headdim=64 and num_warps=8 if headdim=128
57
+ # @triton.autotune(
58
+ # configs=[
59
+ # triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1),
60
+ # # This config has a race condition when EVEN_M == False, disabling it for now.
61
+ # # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1),
62
+ # ],
63
+ # key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM']
64
+ # )
65
+ @triton.heuristics(
66
+ {
67
+ "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
68
+ "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
69
+ "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
70
+ }
71
+ )
72
  @triton.jit
73
  def _fwd_kernel(
74
+ Q, K, V, Bias, Out,
75
+ Lse, TMP, # NOTE: TMP is a scratchpad buffer to workaround a compiler bug
 
 
 
 
 
76
  softmax_scale,
77
+ stride_qb, stride_qh, stride_qm,
78
+ stride_kb, stride_kh, stride_kn,
79
+ stride_vb, stride_vh, stride_vn,
80
+ stride_bb, stride_bh, stride_bm,
81
+ stride_ob, stride_oh, stride_om,
82
+ nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim,
83
+ CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
84
  BIAS_TYPE: tl.constexpr,
85
  IS_CAUSAL: tl.constexpr,
86
  BLOCK_HEADDIM: tl.constexpr,
87
+ EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
88
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
 
 
 
89
  ):
90
  start_m = tl.program_id(0)
91
  off_hb = tl.program_id(1)
 
102
  # Adding parenthesis around indexing might use int32 math instead of int64 math?
103
  # https://github.com/openai/triton/issues/741
104
  # I'm seeing a tiny bit of difference (5-7us)
105
+ q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :])
106
+ k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :])
107
+ v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :])
 
 
 
108
  if BIAS_TYPE == 'vector':
109
  b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n
110
  elif BIAS_TYPE == 'matrix':
111
+ b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :])
 
 
 
112
  # initialize pointer to m and l
113
  t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m
114
+ lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
115
+ m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
116
  acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32)
117
  # load q: it will stay in SRAM throughout
118
  # [2022-10-30] TD: Triton bug - in the case of EVEN_M=True and EVEN_N=False, if we just call
 
126
  if EVEN_HEADDIM:
127
  q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0)
128
  else:
129
+ q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
 
 
130
  other=0.0)
131
  # loop over k, v and update accumulator
132
+ end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k)
 
133
  for start_n in range(0, end_n, BLOCK_N):
134
  start_n = tl.multiple_of(start_n, BLOCK_N)
135
  # -- compute qk ----
 
137
  if EVEN_HEADDIM:
138
  k = tl.load(k_ptrs + start_n * stride_kn)
139
  else:
140
+ k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0)
 
 
141
  else:
142
  if EVEN_HEADDIM:
143
+ k = tl.load(k_ptrs + start_n * stride_kn, mask=(start_n + offs_n)[:, None] < seqlen_k,
 
144
  other=0.0)
145
  else:
146
  k = tl.load(k_ptrs + start_n * stride_kn,
147
+ mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
 
148
  other=0.0)
149
  qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
150
  qk += tl.dot(q, k, trans_b=True)
151
  # Trying to combine the two masks seem to make the result wrong
152
  if not EVEN_N: # Need to mask out otherwise the softmax is wrong
153
+ qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float("-inf"))
 
154
  if IS_CAUSAL:
155
+ qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float("-inf"))
 
156
  if BIAS_TYPE != 'none':
157
  if BIAS_TYPE == 'vector':
158
  if EVEN_N:
159
  bias = tl.load(b_ptrs + start_n).to(tl.float32)
160
  else:
161
+ bias = tl.load(b_ptrs + start_n, mask=(start_n + offs_n) < seqlen_k, other=0.0).to(tl.float32)
 
 
162
  bias = bias[None, :]
163
  elif BIAS_TYPE == 'matrix':
164
  if EVEN_M & EVEN_N:
165
  bias = tl.load(b_ptrs + start_n).to(tl.float32)
166
  else:
167
  bias = tl.load(b_ptrs + start_n,
168
+ mask=(offs_m[:, None] < seqlen_q)
169
+ & ((start_n + offs_n)[None, :] < seqlen_k),
170
  other=0.0).to(tl.float32)
 
 
 
171
  # Slightly faster to multiply the softmax_scale in the tl.exp below since the compiler
172
  # can then fuse the mult and add into an fma instruction. But if we have bias we need to
173
  # to multiply with softmax_scale here.
 
192
  if EVEN_HEADDIM:
193
  v = tl.load(v_ptrs + start_n * stride_vn)
194
  else:
195
+ v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0)
 
 
196
  else:
197
  if EVEN_HEADDIM:
198
+ v = tl.load(v_ptrs + start_n * stride_vn, mask=(start_n + offs_n)[:, None] < seqlen_k,
 
199
  other=0.0)
200
  else:
201
  v = tl.load(v_ptrs + start_n * stride_vn,
202
+ mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
 
203
  other=0.0)
204
  p = p.to(v.dtype)
205
  acc_o += tl.dot(p, v)
 
221
  lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m
222
  tl.store(lse_ptrs, lse_i)
223
  # initialize pointers to output
224
+ offs_d = tl.arange(0, BLOCK_HEADDIM)
225
+ out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:, None] * stride_om + offs_d[None, :])
 
226
  if EVEN_M:
227
  if EVEN_HEADDIM:
228
  tl.store(out_ptrs, acc_o)
 
232
  if EVEN_HEADDIM:
233
  tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q)
234
  else:
235
+ tl.store(out_ptrs, acc_o,
236
+ mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
 
 
237
 
238
 
239
  @triton.jit
240
  def _bwd_preprocess_do_o_dot(
241
+ Out, DO, Delta,
242
+ stride_ob, stride_oh, stride_om,
243
+ stride_dob, stride_doh, stride_dom,
244
+ nheads, seqlen_q, seqlen_q_rounded, headdim,
245
+ BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr,
 
 
 
 
 
 
 
 
 
 
246
  ):
247
  start_m = tl.program_id(0)
248
  off_hb = tl.program_id(1)
 
252
  offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
253
  offs_d = tl.arange(0, BLOCK_HEADDIM)
254
  # load
255
+ o = tl.load(Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :],
256
+ mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
257
+ do = tl.load(DO + off_b * stride_dob + off_h * stride_doh + offs_m[:, None] * stride_dom + offs_d[None, :],
258
+ mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32)
 
 
 
 
 
259
  delta = tl.sum(o * do, axis=1)
260
  # write-back
261
  tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta)
262
 
263
 
264
+ @triton.jit
265
+ def _bwd_store_dk_dv(
266
+ dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
267
+ EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
268
+ ):
269
+ # [2022-11-01] TD: Same bug. In the case of EVEN_N=True and EVEN_M=False,
270
+ # if we just call tl.store(dv_ptrs), there's a race condition
271
+ if EVEN_N & EVEN_M:
272
+ if EVEN_HEADDIM:
273
+ tl.store(dv_ptrs, dv)
274
+ tl.store(dk_ptrs, dk)
275
+ else:
276
+ tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim)
277
+ tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim)
278
+ else:
279
+ if EVEN_HEADDIM:
280
+ tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k)
281
+ tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k)
282
+ else:
283
+ tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
284
+ tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim))
285
+
286
+
287
  @triton.jit
288
  def _bwd_kernel_one_col_block(
289
  start_n,
290
+ Q, K, V, Bias,
291
+ DO, DQ, DK, DV,
292
+ LSE, D,
 
 
 
 
 
 
 
293
  softmax_scale,
294
+ stride_qm, stride_kn, stride_vn, stride_bm,
295
+ stride_dom, stride_dqm, stride_dkn, stride_dvn,
296
+ seqlen_q, seqlen_k, headdim,
 
 
 
 
 
 
 
 
297
  ATOMIC_ADD: tl.constexpr,
298
  BIAS_TYPE: tl.constexpr,
299
  IS_CAUSAL: tl.constexpr,
300
  BLOCK_HEADDIM: tl.constexpr,
301
+ EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
302
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
 
 
 
303
  ):
304
  # We need to make sure begin_m is a multiple of BLOCK_M (not BLOCK_N)
305
  begin_m = 0 if not IS_CAUSAL else ((start_n * BLOCK_N) // BLOCK_M) * BLOCK_M
 
318
  b_ptrs = Bias + offs_n
319
  elif BIAS_TYPE == 'matrix':
320
  b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :])
 
 
321
  # initialize dv and dk
322
  dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
323
  dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
324
+ # There seems to be some problem with Triton pipelining that makes results wrong for
325
+ # headdim=64, seqlen=(113, 255), bias_type='matrix'. In this case the for loop
326
+ # may have zero step, and pipelining with the bias matrix could screw it up.
327
+ # So we just exit early.
328
+ if begin_m >= seqlen_q:
329
+ dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
330
+ dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
331
+ _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
332
+ EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
333
+ return
334
  # k and v stay in SRAM throughout
335
  # [2022-10-30] TD: Same bug as the fwd. In the case of EVEN_N=True and EVEN_M=False,
336
  # if we just call tl.load(k_ptrs), we get the wrong output!
 
346
  k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
347
  v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0)
348
  else:
349
+ k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
 
 
350
  other=0.0)
351
+ v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim),
 
 
352
  other=0.0)
353
  # loop over rows
354
  num_block_m = tl.cdiv(seqlen_q, BLOCK_M)
 
361
  q = tl.load(q_ptrs)
362
  else:
363
  if EVEN_HEADDIM:
364
+ q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0)
 
 
365
  else:
366
+ q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
367
+ & (offs_d[None, :] < headdim), other=0.0)
 
 
368
  # recompute p = softmax(qk, dim=-1).T
369
  qk = tl.dot(q, k, trans_b=True)
370
  # Trying to combine the two masks seem to make the result wrong
371
  if not EVEN_N: # Need to mask out otherwise the softmax is wrong
372
+ qk = tl.where(offs_n[None, :] < seqlen_k, qk, float("-inf"))
373
  if IS_CAUSAL:
374
+ qk = tl.where(offs_m_curr[:, None] >= (offs_n[None, :]), qk, float("-inf"))
 
375
  if BIAS_TYPE != 'none':
376
+ tl.debug_barrier() # Race condition otherwise
377
  if BIAS_TYPE == 'vector':
378
  if EVEN_N:
379
  bias = tl.load(b_ptrs).to(tl.float32)
380
  else:
381
+ bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32)
 
382
  bias = bias[None, :]
383
  elif BIAS_TYPE == 'matrix':
384
  if EVEN_M & EVEN_N:
385
  bias = tl.load(b_ptrs).to(tl.float32)
386
  else:
387
  bias = tl.load(b_ptrs,
388
+ mask=(offs_m_curr[:, None] < seqlen_q)
389
+ & (offs_n[None, :] < seqlen_k),
390
  other=0.0).to(tl.float32)
 
 
 
391
  qk = qk * softmax_scale + bias
392
  # There seems to be a race condition when headdim=48/96, and dq, dk, dv are wrong.
393
  # Also wrong for headdim=64.
 
407
  do = tl.load(do_ptrs)
408
  else:
409
  # [2022-11-01] TD: Triton bug, there's a race condition if we just use m_mask and not d_mask.
410
+ do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q)
411
+ & (offs_d[None, :] < headdim), other=0.0)
 
 
412
  # if EVEN_M:
413
  # if EVEN_HEADDIM:
414
  # do = tl.load(do_ptrs)
 
440
  # compute dk = dot(ds.T, q)
441
  dk += tl.dot(ds, q, trans_a=True)
442
  # compute dq
443
+ if not (EVEN_M & EVEN_HEADDIM): # Otherewise there's a race condition when BIAS_TYPE='matrix'
444
+ tl.debug_barrier()
445
  if not ATOMIC_ADD:
446
  if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
447
+ dq = tl.load(dq_ptrs, eviction_policy="evict_last")
448
  dq += tl.dot(ds, k)
449
+ tl.store(dq_ptrs, dq, eviction_policy="evict_last")
450
  else:
451
  if EVEN_HEADDIM:
452
+ dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0,
453
+ eviction_policy="evict_last")
 
 
454
  dq += tl.dot(ds, k)
455
+ tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q,
456
+ eviction_policy="evict_last")
 
 
457
  else:
458
  dq = tl.load(dq_ptrs,
459
+ mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
460
+ other=0.0, eviction_policy="evict_last")
 
 
461
  dq += tl.dot(ds, k)
462
+ tl.store(dq_ptrs, dq,
463
+ mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim),
464
+ eviction_policy="evict_last")
 
 
465
  else: # If we're parallelizing across the seqlen_k dimension
466
  dq = tl.dot(ds, k)
467
  if EVEN_M & EVEN_HEADDIM: # Race condition if we just do EVEN_M
468
  tl.atomic_add(dq_ptrs, dq)
469
  else:
470
  if EVEN_HEADDIM:
471
+ tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q)
 
 
472
  else:
473
+ tl.atomic_add(dq_ptrs, dq,
474
+ mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim))
 
 
475
  # increment pointers
476
  dq_ptrs += BLOCK_M * stride_dqm
477
  q_ptrs += BLOCK_M * stride_qm
 
481
  # write-back
482
  dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :])
483
  dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :])
484
+ _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim,
485
+ EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
486
 
487
 
488
  def init_to_zero(name):
 
491
 
492
  @triton.autotune(
493
  configs=[
494
+ triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
495
+ triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "SEQUENCE_PARALLEL": True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
496
  # Other configs seem to give wrong results when seqlen_q % 128 != 0, disabling them for now
497
  # # Kernel is buggy (give wrong result) if we set BLOCK_m=128, BLOCK_n=64, num_warps=*4*
498
  # triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')),
 
500
  # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": False}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
501
  # triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "SEQUENCE_PARALLEL": True}, num_warps=4, num_stages=1, pre_hook=init_to_zero('DQ')),
502
  ],
503
+ key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'],
504
+ )
505
+ @triton.heuristics(
506
+ {
507
+ "EVEN_M": lambda args: args["seqlen_q"] % args["BLOCK_M"] == 0,
508
+ "EVEN_N": lambda args: args["seqlen_k"] % args["BLOCK_N"] == 0,
509
+ "EVEN_HEADDIM": lambda args: args["headdim"] == args["BLOCK_HEADDIM"],
510
+ }
511
  )
 
 
 
 
 
512
  @triton.jit
513
  def _bwd_kernel(
514
+ Q, K, V, Bias,
515
+ DO, DQ, DK, DV,
516
+ LSE, D,
 
 
 
 
 
 
 
517
  softmax_scale,
518
+ stride_qb, stride_qh, stride_qm,
519
+ stride_kb, stride_kh, stride_kn,
520
+ stride_vb, stride_vh, stride_vn,
521
+ stride_bb, stride_bh, stride_bm,
522
+ stride_dob, stride_doh, stride_dom,
523
+ stride_dqb, stride_dqh, stride_dqm,
524
+ stride_dkb, stride_dkh, stride_dkn,
525
+ stride_dvb, stride_dvh, stride_dvn,
526
+ nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim,
527
+ CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K,
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
528
  BIAS_TYPE: tl.constexpr,
529
  IS_CAUSAL: tl.constexpr,
530
  BLOCK_HEADDIM: tl.constexpr,
531
  SEQUENCE_PARALLEL: tl.constexpr,
532
+ EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr,
533
+ BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
 
 
 
534
  ):
535
  off_hb = tl.program_id(1)
536
  off_b = off_hb // nheads
 
551
  if not SEQUENCE_PARALLEL:
552
  num_block_n = tl.cdiv(seqlen_k, BLOCK_N)
553
  for start_n in range(0, num_block_n):
554
+ _bwd_kernel_one_col_block(
555
+ start_n,
556
+ Q, K, V, Bias,
557
+ DO, DQ, DK, DV,
558
+ LSE, D,
559
+ softmax_scale,
560
+ stride_qm, stride_kn, stride_vn, stride_bm,
561
+ stride_dom, stride_dqm, stride_dkn, stride_dvn,
562
+ seqlen_q, seqlen_k, headdim,
563
+ ATOMIC_ADD=False,
564
+ BIAS_TYPE=BIAS_TYPE,
565
+ IS_CAUSAL=IS_CAUSAL,
566
+ BLOCK_HEADDIM=BLOCK_HEADDIM,
567
+ EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM,
568
+ BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N
569
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
570
  else:
571
  start_n = tl.program_id(0)
572
+ _bwd_kernel_one_col_block(
573
+ start_n,
574
+ Q, K, V, Bias,
575
+ DO, DQ, DK, DV,
576
+ LSE, D,
577
+ softmax_scale,
578
+ stride_qm, stride_kn, stride_vn, stride_bm,
579
+ stride_dom, stride_dqm, stride_dkn, stride_dvn,
580
+ seqlen_q, seqlen_k, headdim,
581
+ ATOMIC_ADD=True,
582
+ BIAS_TYPE=BIAS_TYPE,
583
+ IS_CAUSAL=IS_CAUSAL,
584
+ BLOCK_HEADDIM=BLOCK_HEADDIM,
585
+ EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM,
586
+ BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N
587
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
588
 
589
 
590
  def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None):
 
595
  assert v.shape == (batch, seqlen_k, nheads, d)
596
  assert d <= 128, 'FlashAttention only support head dimensions up to 128'
597
  assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type'
598
+ assert q.dtype in [torch.float16, torch.bfloat16], 'Only support fp16 and bf16'
 
599
  assert q.is_cuda and k.is_cuda and v.is_cuda
600
  softmax_scale = softmax_scale or 1.0 / math.sqrt(d)
601
 
 
614
  else:
615
  raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
616
  ' or (seqlen_q, seqlen_k)')
617
+ bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
618
+ bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
 
 
 
 
 
 
 
 
 
 
 
619
 
620
  seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128
621
+ lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
622
+ tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32)
 
 
 
 
623
  o = torch.empty_like(q)
624
 
625
  BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
626
+ BLOCK = 128
627
+ num_warps = 4 if d <= 64 else 8
628
+ grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads)
629
+ _fwd_kernel[grid](
630
+ q, k, v, bias, o,
631
+ lse, tmp,
 
 
 
 
 
632
  softmax_scale,
633
+ q.stride(0), q.stride(2), q.stride(1),
634
+ k.stride(0), k.stride(2), k.stride(1),
635
+ v.stride(0), v.stride(2), v.stride(1),
 
 
 
 
 
 
636
  *bias_strides,
637
+ o.stride(0), o.stride(2), o.stride(1),
638
+ nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d,
639
+ seqlen_q // 32, seqlen_k // 32, # key for triton cache (limit number of compilations)
 
 
 
 
 
 
 
640
  # Can't use kwargs here because triton autotune expects key to be args, not kwargs
641
  # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
642
+ bias_type, causal, BLOCK_HEADDIM,
643
+ BLOCK_M=BLOCK, BLOCK_N=BLOCK,
644
+ num_warps=num_warps,
645
+ num_stages=1,
 
 
646
  )
647
  return o, lse, softmax_scale # softmax_scale could have been updated
648
 
649
 
650
+ def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None):
 
 
 
 
 
 
 
 
 
 
 
651
  # Make sure that the last dimension is contiguous
652
  if do.stride(-1) != 1:
653
  do = do.contiguous()
 
666
  # delta = torch.zeros_like(lse)
667
 
668
  BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16)
669
+ grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads)
670
+ _bwd_preprocess_do_o_dot[grid](
671
+ o, do, delta,
672
+ o.stride(0), o.stride(2), o.stride(1),
673
+ do.stride(0), do.stride(2), do.stride(1),
674
+ nheads, seqlen_q, seqlen_q_rounded, d,
675
+ BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM,
 
 
 
 
 
 
 
 
 
 
676
  )
677
 
678
  has_bias = bias is not None
 
689
  else:
690
  raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k)'
691
  ' or (seqlen_q, seqlen_k)')
692
+ bias = bias.expand(batch, nheads, seqlen_q, seqlen_k)
693
+ bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0)
 
 
 
 
 
 
 
 
 
 
 
694
 
695
  # BLOCK_M = 128
696
  # BLOCK_N = 64
697
  # num_warps = 4
698
+ grid = lambda META: (triton.cdiv(seqlen_k, META["BLOCK_N"]) if META["SEQUENCE_PARALLEL"] else 1,
699
+ batch * nheads)
700
+ _bwd_kernel[grid](
701
+ q, k, v, bias,
702
+ do, dq_accum, dk, dv,
703
+ lse, delta,
 
 
 
 
 
 
 
704
  softmax_scale,
705
+ q.stride(0), q.stride(2), q.stride(1),
706
+ k.stride(0), k.stride(2), k.stride(1),
707
+ v.stride(0), v.stride(2), v.stride(1),
 
 
 
 
 
 
708
  *bias_strides,
709
+ do.stride(0), do.stride(2), do.stride(1),
710
+ dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1),
711
+ dk.stride(0), dk.stride(2), dk.stride(1),
712
+ dv.stride(0), dv.stride(2), dv.stride(1),
713
+ nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d,
714
+ seqlen_q // 32, seqlen_k // 32, # key for triton cache (limit number of compilations)
 
 
 
 
 
 
 
 
 
 
 
 
 
715
  # Can't use kwargs here because triton autotune expects key to be args, not kwargs
716
  # IS_CAUSAL=causal, BLOCK_HEADDIM=d,
717
+ bias_type, causal, BLOCK_HEADDIM,
 
 
718
  # SEQUENCE_PARALLEL=False,
719
  # BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
720
  # num_warps=num_warps,
 
723
  dq.copy_(dq_accum)
724
 
725
 
726
+ class FlashAttnQKVPackedFunc(torch.autograd.Function):
727
 
728
  @staticmethod
729
  def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
730
+ """
 
 
 
731
  qkv: (batch, seqlen, 3, nheads, headdim)
732
  bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).
733
  For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).
734
  ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)
 
 
735
  """
736
  # Make sure that the last dimension is contiguous
737
  if qkv.stride(-1) != 1:
738
  qkv = qkv.contiguous()
739
  o, lse, ctx.softmax_scale = _flash_attn_forward(
740
+ qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal,
741
+ softmax_scale=softmax_scale
742
+ )
 
 
 
743
  ctx.save_for_backward(qkv, o, lse, bias)
744
  ctx.causal = causal
745
  return o
 
747
  @staticmethod
748
  def backward(ctx, do):
749
  qkv, o, lse, bias = ctx.saved_tensors
750
+ assert not ctx.needs_input_grad[1], 'FlashAttention does not support bias gradient yet'
 
751
  # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
752
  # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
753
  with torch.inference_mode():
754
  dqkv = torch.empty_like(qkv)
755
+ _flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse,
756
+ dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2],
757
+ bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
 
 
 
 
 
 
 
 
 
758
  return dqkv, None, None, None
759
 
760
 
761
+ flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply
762
 
763
 
764
+ class FlashAttnKVPackedFunc(torch.autograd.Function):
765
 
766
  @staticmethod
767
+ def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
768
+ """
769
+ q: (batch, seqlen_q, nheads, headdim)
770
+ kv: (batch, seqlen_k, 2, nheads, headdim)
771
+ bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
772
+ For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
773
+ ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
774
+ """
775
+ # Make sure that the last dimension is contiguous
776
+ q, kv = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]]
777
+ o, lse, ctx.softmax_scale = _flash_attn_forward(
778
+ q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale
779
+ )
780
+ ctx.save_for_backward(q, kv, o, lse, bias)
781
+ ctx.causal = causal
782
+ return o
783
+
784
+ @staticmethod
785
+ def backward(ctx, do):
786
+ q, kv, o, lse, bias = ctx.saved_tensors
787
+ if len(ctx.needs_input_grad) >= 3:
788
+ assert not ctx.needs_input_grad[2], 'FlashAttention does not support bias gradient yet'
789
+ # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
790
+ # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
791
+ with torch.inference_mode():
792
+ dq = torch.empty_like(q)
793
+ dkv = torch.empty_like(kv)
794
+ _flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse,
795
+ dq, dkv[:, :, 0], dkv[:, :, 1],
796
+ bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
797
+ return dq, dkv, None, None, None
798
+
799
+
800
+ flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply
801
 
802
+
803
+ class FlashAttnFunc(torch.autograd.Function):
804
+
805
+ @staticmethod
806
+ def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
807
+ """
808
  q: (batch_size, seqlen_q, nheads, headdim)
809
+ k, v: (batch_size, seqlen_k, nheads, headdim)
 
810
  bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).
811
  For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).
812
  ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)
 
 
813
  """
814
  # Make sure that the last dimension is contiguous
815
+ q, k, v = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]]
 
 
816
  o, lse, ctx.softmax_scale = _flash_attn_forward(
817
+ q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale
818
+ )
819
  ctx.save_for_backward(q, k, v, o, lse, bias)
820
  ctx.causal = causal
821
  return o
 
823
  @staticmethod
824
  def backward(ctx, do):
825
  q, k, v, o, lse, bias = ctx.saved_tensors
826
+ assert not ctx.needs_input_grad[3], 'FlashAttention does not support bias gradient yet'
 
827
  # Triton's autotune causes the Tensor._version to change, and so Pytorch autograd
828
  # does a memcpy. To avoid this we run in inference_mode, which doesn't track the version.
829
  with torch.inference_mode():
830
  dq = torch.empty_like(q)
831
  dk = torch.empty_like(k)
832
  dv = torch.empty_like(v)
833
+ _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv,
834
+ bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale)
 
 
 
 
 
 
 
 
 
 
835
  return dq, dk, dv, None, None, None
836
 
837
 
838
+ flash_attn_func = FlashAttnFunc.apply