flash_attn_triton.py

'\nCopied from https://github.com/HazyResearch/flash-attention/blob/eff9fe6b8076df59d64d7a3f464696738a3c7c24/flash_attn/flash_attn_triton.py\nupdate imports to use \'triton_pre_mlir\'\n\n*Experimental* implementation of FlashAttention in Triton.\nTested with triton==2.0.0.dev20221202.\nTriton 2.0 has a new backend (MLIR) but seems like it doesn\'t yet work for head dimensions\nother than 64:\nhttps://github.com/openai/triton/blob/d376020f90002757eea3ea9475d4f7cfc2ec5ead/python/triton/ops/flash_attention.py#L207\nWe\'ll update this implementation with the new Triton backend once this is fixed.\n\nWe use the FlashAttention implementation from Phil Tillet a starting point.\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nChanges:\n- Implement both causal and non-causal attention.\n- Implement both self-attention and cross-attention.\n- Support arbitrary seqlens (not just multiples of 128), for both forward and backward.\n- Support all head dimensions up to 128 (not just 16, 32, 64, 128), for both forward and backward.\n- Support attention bias.\n- Speed up the forward pass a bit, and only store the LSE instead of m and l.\n- Make the backward for d=128 much faster by reducing register spilling.\n- Optionally parallelize the backward pass across seqlen_k, to deal with the case of\nsmall batch size * nheads.\n\nCaution:\n- This is an *experimental* implementation. The forward pass should be quite robust but\nI\'m not 100% sure that the backward pass doesn\'t have race conditions (due to the Triton compiler).\n- This implementation has only been tested on A100.\n- If you plan to use headdim other than 64 and 128, you should test for race conditions\n(due to the Triton compiler), as done in tests/test_flash_attn.py\n"test_flash_attn_triton_race_condition". I\'ve tested and fixed many race conditions\nfor different head dimensions (40, 48, 64, 128, 80, 88, 96), but I\'m still not 100% confident\nthat there are none left for other head dimensions.\n\nDifferences between this Triton version and the CUDA version:\n- Triton version doesn\'t support dropout.\n- Triton forward is generally faster than CUDA forward, while Triton backward is\ngenerally slower than CUDA backward. Overall Triton forward + backward is slightly slower\nthan CUDA forward + backward.\n- Triton version doesn\'t support different sequence lengths in a batch (i.e., RaggedTensor/NestedTensor).\n- Triton version supports attention bias, while CUDA version doesn\'t.\n'
import math
import torch
import triton_pre_mlir as triton
import triton_pre_mlir.language as tl

@triton.heuristics({'EVEN_M': (lambda args: ((args['seqlen_q'] % args['BLOCK_M']) == 0)), 'EVEN_N': (lambda args: ((args['seqlen_k'] % args['BLOCK_N']) == 0)), 'EVEN_HEADDIM': (lambda args: (args['headdim'] == args['BLOCK_HEADDIM']))})
@triton.jit
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)))

@triton.jit
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)

@triton.jit
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)))

@triton.jit
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_())

@triton.autotune(configs=[triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': False}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ')), triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'SEQUENCE_PARALLEL': True}, num_warps=8, num_stages=1, pre_hook=init_to_zero('DQ'))], key=['CACHE_KEY_SEQLEN_Q', 'CACHE_KEY_SEQLEN_K', 'BIAS_TYPE', 'IS_CAUSAL', 'BLOCK_HEADDIM'])
@triton.heuristics({'EVEN_M': (lambda args: ((args['seqlen_q'] % args['BLOCK_M']) == 0)), 'EVEN_N': (lambda args: ((args['seqlen_k'] % args['BLOCK_N']) == 0)), 'EVEN_HEADDIM': (lambda args: (args['headdim'] == args['BLOCK_HEADDIM']))})
@triton.jit
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):

    @staticmethod
    def forward(ctx, qkv, bias=None, causal=False, softmax_scale=None):
        '\n            qkv: (batch, seqlen, 3, nheads, headdim)\n            bias: optional, shape broadcastible to (batch, nheads, seqlen, seqlen).\n                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen).\n                ALiBi mask for non-causal would have shape (1, nheads, seqlen, seqlen)\n        '
        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

    @staticmethod
    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):

    @staticmethod
    def forward(ctx, q, kv, bias=None, causal=False, softmax_scale=None):
        '\n            q: (batch, seqlen_q, nheads, headdim)\n            kv: (batch, seqlen_k, 2, nheads, headdim)\n            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n        '
        (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

    @staticmethod
    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):

    @staticmethod
    def forward(ctx, q, k, v, bias=None, causal=False, softmax_scale=None):
        '\n            q: (batch_size, seqlen_q, nheads, headdim)\n            k, v: (batch_size, seqlen_k, nheads, headdim)\n            bias: optional, shape broadcastible to (batch, nheads, seqlen_q, seqlen_k).\n                For example, ALiBi mask for causal would have shape (1, nheads, 1, seqlen_k).\n                ALiBi mask for non-causal would have shape (1, nheads, seqlen_q, seqlen_k)\n        '
        (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

    @staticmethod
    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