Spaces:
Running
on
A10G
Running
on
A10G
""" | |
Copied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py | |
update imports to use 'triton_pre_mlir' | |
*Experimental* implementation of FlashAttention in Triton. | |
Tested with triton==2.0.0.dev20221202. | |
Triton 2.0 has a new backend (MLIR) but seems like it doesn't yet work for head dimensions | |
other than 64: | |
https://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207 | |
We'll update this implementation with the new Triton backend once this is fixed. | |
We use the FlashAttention implementation from Phil Tillet a starting point. | |
https://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py | |
Changes: | |
- Implement both causal and non-causal attention. | |
- Implement both self-attention and cross-attention. | |
- Support arbitrary seqlens (not just multiples of 128), for both forward and backward. | |
- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward. | |
- Support attention bias. | |
- Speed up the forward pass a bit, and only store the LSE instead of m and l. | |
- Make the backward for d=128 much faster by reducing register spilling. | |
- Optionally parallelize the backward pass across seqlen_k, to deal with the case of | |
small batch size * nheads. | |
Caution: | |
- This is an *experimental* implementation. The forward pass should be quite robust but | |
I'm not 100% sure that the backward pass doesn't have race conditions (due to the Triton compiler). | |
- This implementation has only been tested on A100. | |
- If you plan to use headdim other than 64 and 128, you should test for race conditions | |
(due to the Triton compiler), as done in tests/test_flash_attn.py | |
"test_flash_attn_triton_race_condition". I've tested and fixed many race conditions | |
for different head dimensions (40, 48, 64, 128, 80, 88, 96), but I'm still not 100% confident | |
that there are none left for other head dimensions. | |
Differences between this Triton version and the CUDA version: | |
- Triton version doesn't support dropout. | |
- Triton forward is generally faster than CUDA forward, while Triton backward is | |
generally slower than CUDA backward. Overall Triton forward + backward is slightly slower | |
than CUDA forward + backward. | |
- Triton version doesn't support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor). | |
- Triton version supports attention bias, while CUDA version doesn't. | |
""" | |
import math | |
import torch | |
import triton_pre_mlir as triton | |
import triton_pre_mlir.language as tl | |
def _fwd_kernel(Q, K, V, Bias, Out, Lse, TMP, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_ob, stride_oh, stride_om, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): | |
start_m = tl.program_id(0) | |
off_hb = tl.program_id(1) | |
off_b = off_hb // nheads | |
off_h = off_hb % nheads | |
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) | |
offs_n = tl.arange(0, BLOCK_N) | |
offs_d = tl.arange(0, BLOCK_HEADDIM) | |
q_ptrs = Q + off_b * stride_qb + off_h * stride_qh + (offs_m[:, None] * stride_qm + offs_d[None, :]) | |
k_ptrs = K + off_b * stride_kb + off_h * stride_kh + (offs_n[:, None] * stride_kn + offs_d[None, :]) | |
v_ptrs = V + off_b * stride_vb + off_h * stride_vh + (offs_n[:, None] * stride_vn + offs_d[None, :]) | |
if BIAS_TYPE == 'vector': | |
b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + offs_n | |
elif BIAS_TYPE == 'matrix': | |
b_ptrs = Bias + off_b * stride_bb + off_h * stride_bh + (offs_m[:, None] * stride_bm + offs_n[None, :]) | |
t_ptrs = TMP + off_hb * seqlen_q_rounded + offs_m | |
lse_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf') | |
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf') | |
acc_o = tl.zeros([BLOCK_M, BLOCK_HEADDIM], dtype=tl.float32) | |
if EVEN_M & EVEN_N: | |
if EVEN_HEADDIM: | |
q = tl.load(q_ptrs) | |
else: | |
q = tl.load(q_ptrs, mask=offs_d[None, :] < headdim, other=0.0) | |
elif EVEN_HEADDIM: | |
q = tl.load(q_ptrs, mask=offs_m[:, None] < seqlen_q, other=0.0) | |
else: | |
q = tl.load(q_ptrs, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0) | |
end_n = seqlen_k if not IS_CAUSAL else tl.minimum((start_m + 1) * BLOCK_M, seqlen_k) | |
for start_n in range(0, end_n, BLOCK_N): | |
start_n = tl.multiple_of(start_n, BLOCK_N) | |
if EVEN_N & EVEN_M: | |
if EVEN_HEADDIM: | |
k = tl.load(k_ptrs + start_n * stride_kn) | |
else: | |
k = tl.load(k_ptrs + start_n * stride_kn, mask=offs_d[None, :] < headdim, other=0.0) | |
elif EVEN_HEADDIM: | |
k = tl.load(k_ptrs + start_n * stride_kn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0) | |
else: | |
k = tl.load(k_ptrs + start_n * stride_kn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0) | |
qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32) | |
qk += tl.dot(q, k, trans_b=True) | |
if not EVEN_N: | |
qk += tl.where((start_n + offs_n)[None, :] < seqlen_k, 0, float('-inf')) | |
if IS_CAUSAL: | |
qk += tl.where(offs_m[:, None] >= (start_n + offs_n)[None, :], 0, float('-inf')) | |
if BIAS_TYPE != 'none': | |
if BIAS_TYPE == 'vector': | |
if EVEN_N: | |
bias = tl.load(b_ptrs + start_n).to(tl.float32) | |
else: | |
bias = tl.load(b_ptrs + start_n, mask=start_n + offs_n < seqlen_k, other=0.0).to(tl.float32) | |
bias = bias[None, :] | |
elif BIAS_TYPE == 'matrix': | |
if EVEN_M & EVEN_N: | |
bias = tl.load(b_ptrs + start_n).to(tl.float32) | |
else: | |
bias = tl.load(b_ptrs + start_n, mask=(offs_m[:, None] < seqlen_q) & ((start_n + offs_n)[None, :] < seqlen_k), other=0.0).to(tl.float32) | |
qk = qk * softmax_scale + bias | |
m_ij = tl.maximum(tl.max(qk, 1), lse_i) | |
p = tl.exp(qk - m_ij[:, None]) | |
else: | |
m_ij = tl.maximum(tl.max(qk, 1) * softmax_scale, lse_i) | |
p = tl.exp(qk * softmax_scale - m_ij[:, None]) | |
l_ij = tl.sum(p, 1) | |
acc_o_scale = tl.exp(m_i - m_ij) | |
tl.store(t_ptrs, acc_o_scale) | |
acc_o_scale = tl.load(t_ptrs) | |
acc_o = acc_o * acc_o_scale[:, None] | |
if EVEN_N & EVEN_M: | |
if EVEN_HEADDIM: | |
v = tl.load(v_ptrs + start_n * stride_vn) | |
else: | |
v = tl.load(v_ptrs + start_n * stride_vn, mask=offs_d[None, :] < headdim, other=0.0) | |
elif EVEN_HEADDIM: | |
v = tl.load(v_ptrs + start_n * stride_vn, mask=(start_n + offs_n)[:, None] < seqlen_k, other=0.0) | |
else: | |
v = tl.load(v_ptrs + start_n * stride_vn, mask=((start_n + offs_n)[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0) | |
p = p.to(v.dtype) | |
acc_o += tl.dot(p, v) | |
m_i = m_ij | |
l_i_new = tl.exp(lse_i - m_ij) + l_ij | |
lse_i = m_ij + tl.log(l_i_new) | |
o_scale = tl.exp(m_i - lse_i) | |
tl.store(t_ptrs, o_scale) | |
o_scale = tl.load(t_ptrs) | |
acc_o = acc_o * o_scale[:, None] | |
start_m = tl.program_id(0) | |
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) | |
lse_ptrs = Lse + off_hb * seqlen_q_rounded + offs_m | |
tl.store(lse_ptrs, lse_i) | |
offs_d = tl.arange(0, BLOCK_HEADDIM) | |
out_ptrs = Out + off_b * stride_ob + off_h * stride_oh + (offs_m[:, None] * stride_om + offs_d[None, :]) | |
if EVEN_M: | |
if EVEN_HEADDIM: | |
tl.store(out_ptrs, acc_o) | |
else: | |
tl.store(out_ptrs, acc_o, mask=offs_d[None, :] < headdim) | |
elif EVEN_HEADDIM: | |
tl.store(out_ptrs, acc_o, mask=offs_m[:, None] < seqlen_q) | |
else: | |
tl.store(out_ptrs, acc_o, mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim)) | |
def _bwd_preprocess_do_o_dot(Out, DO, Delta, stride_ob, stride_oh, stride_om, stride_dob, stride_doh, stride_dom, nheads, seqlen_q, seqlen_q_rounded, headdim, BLOCK_M: tl.constexpr, BLOCK_HEADDIM: tl.constexpr): | |
start_m = tl.program_id(0) | |
off_hb = tl.program_id(1) | |
off_b = off_hb // nheads | |
off_h = off_hb % nheads | |
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) | |
offs_d = tl.arange(0, BLOCK_HEADDIM) | |
o = tl.load(Out + off_b * stride_ob + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32) | |
do = tl.load(DO + off_b * stride_dob + off_h * stride_doh + offs_m[:, None] * stride_dom + offs_d[None, :], mask=(offs_m[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0).to(tl.float32) | |
delta = tl.sum(o * do, axis=1) | |
tl.store(Delta + off_hb * seqlen_q_rounded + offs_m, delta) | |
def _bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr): | |
if EVEN_N & EVEN_M: | |
if EVEN_HEADDIM: | |
tl.store(dv_ptrs, dv) | |
tl.store(dk_ptrs, dk) | |
else: | |
tl.store(dv_ptrs, dv, mask=offs_d[None, :] < headdim) | |
tl.store(dk_ptrs, dk, mask=offs_d[None, :] < headdim) | |
elif EVEN_HEADDIM: | |
tl.store(dv_ptrs, dv, mask=offs_n[:, None] < seqlen_k) | |
tl.store(dk_ptrs, dk, mask=offs_n[:, None] < seqlen_k) | |
else: | |
tl.store(dv_ptrs, dv, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)) | |
tl.store(dk_ptrs, dk, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim)) | |
def _bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD: tl.constexpr, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): | |
begin_m = 0 if not IS_CAUSAL else start_n * BLOCK_N // BLOCK_M * BLOCK_M | |
offs_qm = begin_m + tl.arange(0, BLOCK_M) | |
offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N) | |
offs_m = tl.arange(0, BLOCK_M) | |
offs_d = tl.arange(0, BLOCK_HEADDIM) | |
q_ptrs = Q + (offs_qm[:, None] * stride_qm + offs_d[None, :]) | |
k_ptrs = K + (offs_n[:, None] * stride_kn + offs_d[None, :]) | |
v_ptrs = V + (offs_n[:, None] * stride_vn + offs_d[None, :]) | |
do_ptrs = DO + (offs_qm[:, None] * stride_dom + offs_d[None, :]) | |
dq_ptrs = DQ + (offs_qm[:, None] * stride_dqm + offs_d[None, :]) | |
if BIAS_TYPE == 'vector': | |
b_ptrs = Bias + offs_n | |
elif BIAS_TYPE == 'matrix': | |
b_ptrs = Bias + (offs_qm[:, None] * stride_bm + offs_n[None, :]) | |
dv = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32) | |
dk = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32) | |
if begin_m >= seqlen_q: | |
dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :]) | |
dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :]) | |
_bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM) | |
return | |
if EVEN_N & EVEN_M: | |
if EVEN_HEADDIM: | |
k = tl.load(k_ptrs) | |
v = tl.load(v_ptrs) | |
else: | |
k = tl.load(k_ptrs, mask=offs_d[None, :] < headdim, other=0.0) | |
v = tl.load(v_ptrs, mask=offs_d[None, :] < headdim, other=0.0) | |
elif EVEN_HEADDIM: | |
k = tl.load(k_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0) | |
v = tl.load(v_ptrs, mask=offs_n[:, None] < seqlen_k, other=0.0) | |
else: | |
k = tl.load(k_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0) | |
v = tl.load(v_ptrs, mask=(offs_n[:, None] < seqlen_k) & (offs_d[None, :] < headdim), other=0.0) | |
num_block_m = tl.cdiv(seqlen_q, BLOCK_M) | |
for start_m in range(begin_m, num_block_m * BLOCK_M, BLOCK_M): | |
start_m = tl.multiple_of(start_m, BLOCK_M) | |
offs_m_curr = start_m + offs_m | |
if EVEN_M & EVEN_HEADDIM: | |
q = tl.load(q_ptrs) | |
elif EVEN_HEADDIM: | |
q = tl.load(q_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0) | |
else: | |
q = tl.load(q_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0) | |
qk = tl.dot(q, k, trans_b=True) | |
if not EVEN_N: | |
qk = tl.where(offs_n[None, :] < seqlen_k, qk, float('-inf')) | |
if IS_CAUSAL: | |
qk = tl.where(offs_m_curr[:, None] >= offs_n[None, :], qk, float('-inf')) | |
if BIAS_TYPE != 'none': | |
tl.debug_barrier() | |
if BIAS_TYPE == 'vector': | |
if EVEN_N: | |
bias = tl.load(b_ptrs).to(tl.float32) | |
else: | |
bias = tl.load(b_ptrs, mask=offs_n < seqlen_k, other=0.0).to(tl.float32) | |
bias = bias[None, :] | |
elif BIAS_TYPE == 'matrix': | |
if EVEN_M & EVEN_N: | |
bias = tl.load(b_ptrs).to(tl.float32) | |
else: | |
bias = tl.load(b_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_n[None, :] < seqlen_k), other=0.0).to(tl.float32) | |
qk = qk * softmax_scale + bias | |
if not EVEN_M & EVEN_HEADDIM: | |
tl.debug_barrier() | |
lse_i = tl.load(LSE + offs_m_curr) | |
if BIAS_TYPE == 'none': | |
p = tl.exp(qk * softmax_scale - lse_i[:, None]) | |
else: | |
p = tl.exp(qk - lse_i[:, None]) | |
if EVEN_M & EVEN_HEADDIM: | |
do = tl.load(do_ptrs) | |
else: | |
do = tl.load(do_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0) | |
dv += tl.dot(p.to(do.dtype), do, trans_a=True) | |
if not EVEN_M & EVEN_HEADDIM: | |
tl.debug_barrier() | |
dp = tl.dot(do, v, trans_b=True) | |
if not EVEN_HEADDIM: | |
tl.debug_barrier() | |
Di = tl.load(D + offs_m_curr) | |
ds = (p * (dp - Di[:, None]) * softmax_scale).to(q.dtype) | |
dk += tl.dot(ds, q, trans_a=True) | |
if not EVEN_M & EVEN_HEADDIM: | |
tl.debug_barrier() | |
if not ATOMIC_ADD: | |
if EVEN_M & EVEN_HEADDIM: | |
dq = tl.load(dq_ptrs, eviction_policy='evict_last') | |
dq += tl.dot(ds, k) | |
tl.store(dq_ptrs, dq, eviction_policy='evict_last') | |
elif EVEN_HEADDIM: | |
dq = tl.load(dq_ptrs, mask=offs_m_curr[:, None] < seqlen_q, other=0.0, eviction_policy='evict_last') | |
dq += tl.dot(ds, k) | |
tl.store(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q, eviction_policy='evict_last') | |
else: | |
dq = tl.load(dq_ptrs, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), other=0.0, eviction_policy='evict_last') | |
dq += tl.dot(ds, k) | |
tl.store(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim), eviction_policy='evict_last') | |
else: | |
dq = tl.dot(ds, k) | |
if EVEN_M & EVEN_HEADDIM: | |
tl.atomic_add(dq_ptrs, dq) | |
elif EVEN_HEADDIM: | |
tl.atomic_add(dq_ptrs, dq, mask=offs_m_curr[:, None] < seqlen_q) | |
else: | |
tl.atomic_add(dq_ptrs, dq, mask=(offs_m_curr[:, None] < seqlen_q) & (offs_d[None, :] < headdim)) | |
dq_ptrs += BLOCK_M * stride_dqm | |
q_ptrs += BLOCK_M * stride_qm | |
do_ptrs += BLOCK_M * stride_dom | |
if BIAS_TYPE == 'matrix': | |
b_ptrs += BLOCK_M * stride_bm | |
dv_ptrs = DV + (offs_n[:, None] * stride_dvn + offs_d[None, :]) | |
dk_ptrs = DK + (offs_n[:, None] * stride_dkn + offs_d[None, :]) | |
_bwd_store_dk_dv(dk_ptrs, dv_ptrs, dk, dv, offs_n, offs_d, seqlen_k, headdim, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM) | |
def init_to_zero(name): | |
return lambda nargs: nargs[name].zero_() | |
def _bwd_kernel(Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qb, stride_qh, stride_qm, stride_kb, stride_kh, stride_kn, stride_vb, stride_vh, stride_vn, stride_bb, stride_bh, stride_bm, stride_dob, stride_doh, stride_dom, stride_dqb, stride_dqh, stride_dqm, stride_dkb, stride_dkh, stride_dkn, stride_dvb, stride_dvh, stride_dvn, nheads, seqlen_q, seqlen_k, seqlen_q_rounded, headdim, CACHE_KEY_SEQLEN_Q, CACHE_KEY_SEQLEN_K, BIAS_TYPE: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_HEADDIM: tl.constexpr, SEQUENCE_PARALLEL: tl.constexpr, EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_HEADDIM: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): | |
off_hb = tl.program_id(1) | |
off_b = off_hb // nheads | |
off_h = off_hb % nheads | |
Q += off_b * stride_qb + off_h * stride_qh | |
K += off_b * stride_kb + off_h * stride_kh | |
V += off_b * stride_vb + off_h * stride_vh | |
DO += off_b * stride_dob + off_h * stride_doh | |
DQ += off_b * stride_dqb + off_h * stride_dqh | |
DK += off_b * stride_dkb + off_h * stride_dkh | |
DV += off_b * stride_dvb + off_h * stride_dvh | |
if BIAS_TYPE != 'none': | |
Bias += off_b * stride_bb + off_h * stride_bh | |
D += off_hb * seqlen_q_rounded | |
LSE += off_hb * seqlen_q_rounded | |
if not SEQUENCE_PARALLEL: | |
num_block_n = tl.cdiv(seqlen_k, BLOCK_N) | |
for start_n in range(0, num_block_n): | |
_bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=False, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) | |
else: | |
start_n = tl.program_id(0) | |
_bwd_kernel_one_col_block(start_n, Q, K, V, Bias, DO, DQ, DK, DV, LSE, D, softmax_scale, stride_qm, stride_kn, stride_vn, stride_bm, stride_dom, stride_dqm, stride_dkn, stride_dvn, seqlen_q, seqlen_k, headdim, ATOMIC_ADD=True, BIAS_TYPE=BIAS_TYPE, IS_CAUSAL=IS_CAUSAL, BLOCK_HEADDIM=BLOCK_HEADDIM, EVEN_M=EVEN_M, EVEN_N=EVEN_N, EVEN_HEADDIM=EVEN_HEADDIM, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N) | |
def _flash_attn_forward(q, k, v, bias=None, causal=False, softmax_scale=None): | |
(batch, seqlen_q, nheads, d) = q.shape | |
(_, seqlen_k, _, _) = k.shape | |
assert k.shape == (batch, seqlen_k, nheads, d) | |
assert v.shape == (batch, seqlen_k, nheads, d) | |
assert d <= 128, 'FlashAttention only support head dimensions up to 128' | |
assert q.dtype == k.dtype == v.dtype, 'All tensors must have the same type' | |
assert q.dtype in [torch.float16, torch.bfloat16], 'Only support fp16 and bf16' | |
assert q.is_cuda and k.is_cuda and v.is_cuda | |
softmax_scale = softmax_scale or 1.0 / math.sqrt(d) | |
has_bias = bias is not None | |
bias_type = 'none' | |
if has_bias: | |
assert bias.dtype in [q.dtype, torch.float] | |
assert bias.is_cuda | |
assert bias.dim() == 4 | |
if bias.stride(-1) != 1: | |
bias = bias.contiguous() | |
if bias.shape[2:] == (1, seqlen_k): | |
bias_type = 'vector' | |
elif bias.shape[2:] == (seqlen_q, seqlen_k): | |
bias_type = 'matrix' | |
else: | |
raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)') | |
bias = bias.expand(batch, nheads, seqlen_q, seqlen_k) | |
bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0) | |
seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128 | |
lse = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32) | |
tmp = torch.empty((batch, nheads, seqlen_q_rounded), device=q.device, dtype=torch.float32) | |
o = torch.empty_like(q) | |
BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16) | |
BLOCK = 128 | |
num_warps = 4 if d <= 64 else 8 | |
grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads) | |
_fwd_kernel[grid](q, k, v, bias, o, lse, tmp, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, o.stride(0), o.stride(2), o.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM, BLOCK_M=BLOCK, BLOCK_N=BLOCK, num_warps=num_warps, num_stages=1) | |
return (o, lse, softmax_scale) | |
def _flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=None, causal=False, softmax_scale=None): | |
if do.stride(-1) != 1: | |
do = do.contiguous() | |
(batch, seqlen_q, nheads, d) = q.shape | |
(_, seqlen_k, _, _) = k.shape | |
assert d <= 128 | |
seqlen_q_rounded = math.ceil(seqlen_q / 128) * 128 | |
assert lse.shape == (batch, nheads, seqlen_q_rounded) | |
assert q.stride(-1) == k.stride(-1) == v.stride(-1) == o.stride(-1) == 1 | |
assert dq.stride(-1) == dk.stride(-1) == dv.stride(-1) == 1 | |
softmax_scale = softmax_scale or 1.0 / math.sqrt(d) | |
dq_accum = torch.empty_like(q, dtype=torch.float32) | |
delta = torch.empty_like(lse) | |
BLOCK_HEADDIM = max(triton.next_power_of_2(d), 16) | |
grid = lambda META: (triton.cdiv(seqlen_q, META['BLOCK_M']), batch * nheads) | |
_bwd_preprocess_do_o_dot[grid](o, do, delta, o.stride(0), o.stride(2), o.stride(1), do.stride(0), do.stride(2), do.stride(1), nheads, seqlen_q, seqlen_q_rounded, d, BLOCK_M=128, BLOCK_HEADDIM=BLOCK_HEADDIM) | |
has_bias = bias is not None | |
bias_type = 'none' | |
if has_bias: | |
assert bias.dtype in [q.dtype, torch.float] | |
assert bias.is_cuda | |
assert bias.dim() == 4 | |
assert bias.stride(-1) == 1 | |
if bias.shape[2:] == (1, seqlen_k): | |
bias_type = 'vector' | |
elif bias.shape[2:] == (seqlen_q, seqlen_k): | |
bias_type = 'matrix' | |
else: | |
raise RuntimeError('Last 2 dimensions of bias must be (1, seqlen_k) or (seqlen_q, seqlen_k)') | |
bias = bias.expand(batch, nheads, seqlen_q, seqlen_k) | |
bias_strides = (bias.stride(0), bias.stride(1), bias.stride(2)) if has_bias else (0, 0, 0) | |
grid = lambda META: (triton.cdiv(seqlen_k, META['BLOCK_N']) if META['SEQUENCE_PARALLEL'] else 1, batch * nheads) | |
_bwd_kernel[grid](q, k, v, bias, do, dq_accum, dk, dv, lse, delta, softmax_scale, q.stride(0), q.stride(2), q.stride(1), k.stride(0), k.stride(2), k.stride(1), v.stride(0), v.stride(2), v.stride(1), *bias_strides, do.stride(0), do.stride(2), do.stride(1), dq_accum.stride(0), dq_accum.stride(2), dq_accum.stride(1), dk.stride(0), dk.stride(2), dk.stride(1), dv.stride(0), dv.stride(2), dv.stride(1), nheads, seqlen_q, seqlen_k, seqlen_q_rounded, d, seqlen_q // 32, seqlen_k // 32, bias_type, causal, BLOCK_HEADDIM) | |
dq.copy_(dq_accum) | |
class FlashAttnQKVPackedFunc(torch.autograd.Function): | |
def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None): | |
""" | |
qkv: (batch, seqlen, 3, nheads, headdim) | |
bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen). | |
For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen). | |
ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen) | |
""" | |
if qkv.stride(-1) != 1: | |
qkv = qkv.contiguous() | |
(o, lse, ctx.softmax_scale) = _flash_attn_forward(qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], bias=bias, causal=causal, softmax_scale=softmax_scale) | |
ctx.save_for_backward(qkv, o, lse, bias) | |
ctx.causal = causal | |
return o | |
def backward(ctx, do): | |
(qkv, o, lse, bias) = ctx.saved_tensors | |
assert not ctx.needs_input_grad[1], 'FlashAttention does not support bias gradient yet' | |
with torch.inference_mode(): | |
dqkv = torch.empty_like(qkv) | |
_flash_attn_backward(do, qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], o, lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale) | |
return (dqkv, None, None, None) | |
flash_attn_qkvpacked_func = FlashAttnQKVPackedFunc.apply | |
class FlashAttnKVPackedFunc(torch.autograd.Function): | |
def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None): | |
""" | |
q: (batch, seqlen_q, nheads, headdim) | |
kv: (batch, seqlen_k, 2, nheads, headdim) | |
bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k). | |
For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k). | |
ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k) | |
""" | |
(q, kv) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, kv]] | |
(o, lse, ctx.softmax_scale) = _flash_attn_forward(q, kv[:, :, 0], kv[:, :, 1], bias=bias, causal=causal, softmax_scale=softmax_scale) | |
ctx.save_for_backward(q, kv, o, lse, bias) | |
ctx.causal = causal | |
return o | |
def backward(ctx, do): | |
(q, kv, o, lse, bias) = ctx.saved_tensors | |
if len(ctx.needs_input_grad) >= 3: | |
assert not ctx.needs_input_grad[2], 'FlashAttention does not support bias gradient yet' | |
with torch.inference_mode(): | |
dq = torch.empty_like(q) | |
dkv = torch.empty_like(kv) | |
_flash_attn_backward(do, q, kv[:, :, 0], kv[:, :, 1], o, lse, dq, dkv[:, :, 0], dkv[:, :, 1], bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale) | |
return (dq, dkv, None, None, None) | |
flash_attn_kvpacked_func = FlashAttnKVPackedFunc.apply | |
class FlashAttnFunc(torch.autograd.Function): | |
def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None): | |
""" | |
q: (batch_size, seqlen_q, nheads, headdim) | |
k, v: (batch_size, seqlen_k, nheads, headdim) | |
bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k). | |
For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k). | |
ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k) | |
""" | |
(q, k, v) = [x if x.stride(-1) == 1 else x.contiguous() for x in [q, k, v]] | |
(o, lse, ctx.softmax_scale) = _flash_attn_forward(q, k, v, bias=bias, causal=causal, softmax_scale=softmax_scale) | |
ctx.save_for_backward(q, k, v, o, lse, bias) | |
ctx.causal = causal | |
return o | |
def backward(ctx, do): | |
(q, k, v, o, lse, bias) = ctx.saved_tensors | |
assert not ctx.needs_input_grad[3], 'FlashAttention does not support bias gradient yet' | |
with torch.inference_mode(): | |
dq = torch.empty_like(q) | |
dk = torch.empty_like(k) | |
dv = torch.empty_like(v) | |
_flash_attn_backward(do, q, k, v, o, lse, dq, dk, dv, bias=bias, causal=ctx.causal, softmax_scale=ctx.softmax_scale) | |
return (dq, dk, dv, None, None, None) | |
flash_attn_func = FlashAttnFunc.apply |