123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508 |
- from typing import Optional, Union
- import torch
- import torch.nn as nn
- import flashattn_hopper_cuda
- def _flash_attn_forward(q, k, v, softmax_scale, causal,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0):
- maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
- q, k = [maybe_contiguous(x) for x in (q, k)]
- v = v.contiguous() if v.stride(-1) != 1 and v.stride(-3) != 1 else v
- out, q, k, v, out_padded, softmax_lse = flashattn_hopper_cuda.fwd(
- q,
- k,
- v,
- None,
- softmax_scale,
- causal,
- q_scale, k_scale, v_scale,
- window_size[0], window_size[1],
- softcap
- )
- return out, q, k, v, out_padded, softmax_lse
- def _flash_attn_varlen_forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, softmax_scale, causal,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1), softcap=0.0):
- maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
- q, k, v = [maybe_contiguous(x) for x in (q, k, v)]
- out, q, k, v, out_padded, softmax_lse = flashattn_hopper_cuda.fwd_varlen(
- q,
- k,
- v,
- None,
- cu_seqlens_q, cu_seqlens_k, None, None, max_seqlen_q, max_seqlen_k,
- softmax_scale,
- causal,
- q_scale, k_scale, v_scale,
- window_size[0], window_size[1],
- softcap,
- )
- return out, q, k, v, out_padded, softmax_lse
- def _flash_attn_backward(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- dq,
- dk,
- dv,
- softmax_scale,
- causal,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False
- ):
- maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
-
- dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
- dq, dk, dv, softmax_d, *rest = flashattn_hopper_cuda.bwd(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- dq,
- dk,
- dv,
- softmax_scale,
- causal,
- window_size[0],
- window_size[1],
- softcap,
- deterministic,
- )
- return dq, dk, dv, softmax_d
- def _flash_attn_varlen_backward(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- dq,
- dk,
- dv,
- softmax_scale,
- causal,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False
- ):
- maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
-
- dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
- dq, dk, dv, softmax_d, *rest = flashattn_hopper_cuda.bwd_varlen(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- dq,
- dk,
- dv,
- cu_seqlens_q,
- cu_seqlens_k,
- None, None,
- max_seqlen_q,
- max_seqlen_k,
- softmax_scale,
- causal,
- window_size[0],
- window_size[1],
- softcap,
- deterministic,
- )
- return dq, dk, dv, softmax_d
- class FlashAttnQKVPackedFunc(torch.autograd.Function):
- @staticmethod
- def forward(
- ctx,
- qkv,
- softmax_scale,
- causal,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False,
- num_heads_q=None,
- ):
- if softmax_scale is None:
- softmax_scale = qkv.shape[-1] ** (-0.5)
- if qkv.dim() == 5:
- assert qkv.shape[-3] == 3
- q, k, v = qkv.unbind(dim=-3)
- else:
- assert qkv.dim() == 4
- assert num_heads_q is not None
- num_heads_k = (qkv.shape[2] - num_heads_q) // 2
- assert num_heads_k * 2 + num_heads_q == qkv.shape[2]
- q, k, v = qkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2)
- out, q, k, v, out_padded, softmax_lse = _flash_attn_forward(
- q,
- k,
- v,
- softmax_scale,
- causal=causal,
- q_scale=q_scale, k_scale=k_scale, v_scale=v_scale,
- window_size=window_size,
- softcap=softcap,
- )
- ctx.save_for_backward(q, k, v, out_padded, softmax_lse)
- ctx.softmax_scale = softmax_scale
- ctx.causal = causal
- ctx.window_size = window_size
- ctx.softcap = softcap
- ctx.deterministic = deterministic
- ctx.ndim = qkv.dim()
-
- return out
- @staticmethod
- def backward(ctx, dout, *args):
- q, k, v, out, softmax_lse = ctx.saved_tensors
- if ctx.ndim == 5:
- qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
- dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
- dq, dk, dv = dqkv.unbind(dim=-3)
- else:
- num_heads_q = q.shape[2]
- num_heads_k = k.shape[2]
- qkv_shape = q.shape[:-2] + (num_heads_q + num_heads_k * 2, *q.shape[-1:])
- dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
- dq, dk, dv = dqkv.split([num_heads_q, num_heads_k, num_heads_k], dim=-2)
- _flash_attn_backward(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- dq,
- dk,
- dv,
- ctx.softmax_scale,
- ctx.causal,
- ctx.window_size,
- ctx.softcap,
- ctx.deterministic,
- )
- dqkv = dqkv[..., : dout.shape[-1]]
- return dqkv, None, None, None, None, None, None, None, None, None
- class FlashAttnFunc(torch.autograd.Function):
- @staticmethod
- def forward(
- ctx,
- q,
- k,
- v,
- softmax_scale,
- causal,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False,
- ):
- if softmax_scale is None:
- softmax_scale = q.shape[-1] ** (-0.5)
- out, q, k, v, out_padded, softmax_lse = _flash_attn_forward(
- q,
- k,
- v,
- softmax_scale,
- causal=causal,
- q_scale=q_scale, k_scale=k_scale, v_scale=v_scale,
- window_size=window_size,
- softcap=softcap,
- )
- ctx.save_for_backward(q, k, v, out_padded, softmax_lse)
- ctx.softmax_scale = softmax_scale
- ctx.causal = causal
- ctx.window_size = window_size
- ctx.softcap = softcap
- ctx.deterministic = deterministic
- return out, softmax_lse
- @staticmethod
- def backward(ctx, dout, *args):
- q, k, v, out, softmax_lse = ctx.saved_tensors
- dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
- _flash_attn_backward(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- dq,
- dk,
- dv,
- ctx.softmax_scale,
- ctx.causal,
- ctx.window_size,
- ctx.softcap,
- ctx.deterministic,
- )
- dq = dq[..., : dout.shape[-1]]
- dk = dk[..., : dout.shape[-1]]
- dv = dv[..., : dout.shape[-1]]
- return dq, dk, dv, None, None, None, None, None, None, None, None
- class FlashAttnVarlenFunc(torch.autograd.Function):
- @staticmethod
- def forward(
- ctx,
- q,
- k,
- v,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- softmax_scale,
- causal,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False,
- ):
- if softmax_scale is None:
- softmax_scale = q.shape[-1] ** (-0.5)
- out, q, k, v, out_padded, softmax_lse = _flash_attn_varlen_forward(
- q,
- k,
- v,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- softmax_scale,
- causal=causal,
- q_scale=q_scale, k_scale=k_scale, v_scale=v_scale,
- window_size=window_size,
- softcap=softcap,
- )
- ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k)
- ctx.max_seqlen_q = max_seqlen_q
- ctx.max_seqlen_k = max_seqlen_k
- ctx.softmax_scale = softmax_scale
- ctx.causal = causal
- ctx.window_size = window_size
- ctx.softcap = softcap
- ctx.deterministic = deterministic
- return out, softmax_lse
- @staticmethod
- def backward(ctx, dout, *args):
- q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
- dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
- _flash_attn_varlen_backward(
- dout,
- q,
- k,
- v,
- out,
- softmax_lse,
- cu_seqlens_q,
- cu_seqlens_k,
- ctx.max_seqlen_q,
- ctx.max_seqlen_k,
- dq,
- dk,
- dv,
- ctx.softmax_scale,
- ctx.causal,
- ctx.window_size,
- ctx.softcap,
- ctx.deterministic,
- )
- dq = dq[..., : dout.shape[-1]]
- dk = dk[..., : dout.shape[-1]]
- dv = dv[..., : dout.shape[-1]]
- return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None
- def flash_attn_qkvpacked_func(
- qkv,
- softmax_scale=None,
- causal=False,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False,
- num_heads_q=None,
- ):
- """dropout_p should be set to 0.0 during evaluation
- If Q, K, V are already stacked into 1 tensor, this function will be faster than
- calling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation
- of the gradients of Q, K, V.
- For multi-query and grouped-query attention (MQA/GQA), please see
- flash_attn_kvpacked_func and flash_attn_func.
- If window_size != (-1, -1), implements sliding window local attention. Query at position i
- will only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.
- Arguments:
- qkv: (batch_size, seqlen, 3, nheads, headdim)
- dropout_p: float. Dropout probability.
- softmax_scale: float. The scaling of QK^T before applying softmax.
- Default to 1 / sqrt(headdim).
- causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
- window_size: (left, right). If not (-1, -1), implements sliding window local attention.
- softcap: float. Anything > 0 activates softcapping attention.
- alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to
- the attention score of query i and key j.
- deterministic: bool. Whether to use the deterministic implementation of the backward pass,
- which is slightly slower and uses more memory. The forward pass is always deterministic.
- return_attn_probs: bool. Whether to return the attention probabilities. This option is for
- testing only. The returned probabilities are not guaranteed to be correct
- (they might not have the right scaling).
- Return:
- out: (batch_size, seqlen, nheads, headdim).
- softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
- logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
- normalization factor).
- S_dmask [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen, seqlen).
- The output of softmax (possibly with different scaling). It also encodes the dropout
- pattern (negative means that location was dropped, nonnegative means it was kept).
- """
- return FlashAttnQKVPackedFunc.apply(
- qkv,
- softmax_scale,
- causal,
- q_scale, k_scale, v_scale,
- window_size,
- softcap,
- deterministic,
- num_heads_q,
- )
- def flash_attn_func(
- q,
- k,
- v,
- softmax_scale=None,
- causal=False,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False
- ):
- """dropout_p should be set to 0.0 during evaluation
- Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads
- than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.
- For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head
- 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.
- If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.
- For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:
- 1 1 1 1 0
- 1 1 1 1 1
- If seqlen_q = 5 and seqlen_k = 2, the causal mask is:
- 0 0
- 0 0
- 0 0
- 1 0
- 1 1
- If the row of the mask is all zero, the output will be zero.
- If window_size != (-1, -1), implements sliding window local attention. Query at position i
- will only attend to keys between
- [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.
- Arguments:
- q: (batch_size, seqlen, nheads, headdim)
- k: (batch_size, seqlen, nheads_k, headdim)
- v: (batch_size, seqlen, nheads_k, headdim)
- dropout_p: float. Dropout probability.
- softmax_scale: float. The scaling of QK^T before applying softmax.
- Default to 1 / sqrt(headdim).
- causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).
- window_size: (left, right). If not (-1, -1), implements sliding window local attention.
- alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of
- (-alibi_slope * |i + seqlen_k - seqlen_q - j|)
- is added to the attention score of query i and key j.
- deterministic: bool. Whether to use the deterministic implementation of the backward pass,
- which is slightly slower and uses more memory. The forward pass is always deterministic.
- return_attn_probs: bool. Whether to return the attention probabilities. This option is for
- testing only. The returned probabilities are not guaranteed to be correct
- (they might not have the right scaling).
- Return:
- out: (batch_size, seqlen, nheads, headdim).
- softmax_lse [optional, if return_attn_probs=True]: (batch_size, nheads, seqlen). The
- logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
- normalization factor).
- """
- return FlashAttnFunc.apply(
- q,
- k,
- v,
- softmax_scale,
- causal,
- q_scale, k_scale, v_scale,
- window_size,
- softcap,
- deterministic,
- )
- def flash_attn_varlen_func(
- q,
- k,
- v,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- softmax_scale=None,
- causal=False,
- q_scale=None, k_scale=None, v_scale=None,
- window_size=(-1, -1),
- softcap=0.0,
- deterministic=False
- ):
- return FlashAttnVarlenFunc.apply(
- q,
- k,
- v,
- cu_seqlens_q,
- cu_seqlens_k,
- max_seqlen_q,
- max_seqlen_k,
- softmax_scale,
- causal,
- q_scale, k_scale, v_scale,
- window_size,
- softcap,
- deterministic,
- )
|