967 lines
56 KiB
Python
967 lines
56 KiB
Python
|
import os
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
|
||
|
|
import flash_attn_2_cuda as flash_attn_cuda
|
||
|
|
from einops import rearrange
|
||
|
|
import warnings
|
||
|
|
|
||
|
|
global enable_ixdnn
|
||
|
|
enable_ixdnn = os.getenv('ENABLE_FLASH_ATTENTION_WITH_IXDNN', '1') != '0'
|
||
|
|
|
||
|
|
def _get_block_size(device, head_dim, is_dropout, is_causal):
|
||
|
|
# This should match the block sizes in the CUDA kernel
|
||
|
|
assert head_dim <= 256
|
||
|
|
major, minor = torch.cuda.get_device_capability(device)
|
||
|
|
is_sm8x = major == 8 and minor > 0 # Only include sm86 and sm89, exclude sm80 (A100)
|
||
|
|
is_sm80 = major == 8 and minor == 0
|
||
|
|
is_sm90 = major == 9 and minor == 0
|
||
|
|
if head_dim <= 32:
|
||
|
|
return 128, 128
|
||
|
|
if head_dim <= 64:
|
||
|
|
return (128, 128) if not is_dropout else (128, 64)
|
||
|
|
elif head_dim <= 96:
|
||
|
|
return (64, 64) if (is_sm8x and is_causal) else (128, 64)
|
||
|
|
elif head_dim <= 128:
|
||
|
|
if is_sm8x:
|
||
|
|
return (64, 64) if (not is_dropout and is_causal) else (128, 32)
|
||
|
|
else:
|
||
|
|
return 128, (64 if not is_dropout else 32)
|
||
|
|
elif head_dim <= 160:
|
||
|
|
if is_sm8x:
|
||
|
|
return (128, 64) if not is_causal else (64, 64)
|
||
|
|
else:
|
||
|
|
return 128, 32
|
||
|
|
elif head_dim <= 192:
|
||
|
|
return (128, 64) if not is_dropout else (64, 64)
|
||
|
|
elif head_dim <= 224:
|
||
|
|
return (128, 64) if (is_sm80 or is_sm90) else (64, 64)
|
||
|
|
elif head_dim <= 256:
|
||
|
|
return (128, 64) if is_sm80 else (64, 64)
|
||
|
|
|
||
|
|
|
||
|
|
def _flash_attn_forward(q, k, v, dropout_p, softmax_scale, causal, return_softmax, use_alibi, alibi_mode):
|
||
|
|
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, S_dmask = flash_attn_cuda.fwd(
|
||
|
|
q, k, v, None, dropout_p, softmax_scale, causal, return_softmax, use_alibi, alibi_mode, None
|
||
|
|
)
|
||
|
|
return out, q, k, v, out_padded, softmax_lse, S_dmask
|
||
|
|
|
||
|
|
|
||
|
|
def _flash_attn_varlen_forward(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal, return_softmax, use_alibi, alibi_mode):
|
||
|
|
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, S_dmask = flash_attn_cuda.varlen_fwd(
|
||
|
|
q, k, v, None, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p,
|
||
|
|
softmax_scale, False, causal, return_softmax, use_alibi, alibi_mode, None
|
||
|
|
)
|
||
|
|
# if out.isnan().any() or softmax_lse.isnan().any():
|
||
|
|
# breakpoint()
|
||
|
|
return out, q, k, v, out_padded, softmax_lse, S_dmask
|
||
|
|
|
||
|
|
|
||
|
|
def _flash_attn_backward(dout, q, k, v, out, softmax_lse, dq, dk, dv,
|
||
|
|
dropout_p, softmax_scale, causal, use_alibi, alibi_mode):
|
||
|
|
maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
|
||
|
|
# dq, dk, dv are allocated by us so they should already be contiguous
|
||
|
|
dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
|
||
|
|
dq, dk, dv, softmax_d, = flash_attn_cuda.bwd(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, dropout_p, softmax_scale, causal, use_alibi, alibi_mode, None
|
||
|
|
)
|
||
|
|
return dq, dk, dv, softmax_d
|
||
|
|
|
||
|
|
|
||
|
|
def _flash_attn_varlen_backward(dout, q, k, v, out, softmax_lse, dq, dk, dv,
|
||
|
|
cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal, use_alibi, alibi_mode):
|
||
|
|
maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
|
||
|
|
# dq, dk, dv are allocated by us so they should already be contiguous
|
||
|
|
dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
|
||
|
|
dq, dk, dv, softmax_d, = flash_attn_cuda.varlen_bwd(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
|
||
|
|
max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale, False, causal, use_alibi, alibi_mode, None
|
||
|
|
)
|
||
|
|
# if dk.isnan().any() or dk.isnan().any() or dv.isnan().any() or softmax_d.isnan().any():
|
||
|
|
# breakpoint()
|
||
|
|
return dq, dk, dv, softmax_d
|
||
|
|
|
||
|
|
def _flash_attn_forward_ixdnn(q, k, v, cu_seqlens_q, cu_seqlens_k, dropout_p,
|
||
|
|
causal, return_softmax, use_alibi, alibi_mode, imp_mode, window_size):
|
||
|
|
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, S_dmask = flash_attn_cuda.fwd_ixdnn(
|
||
|
|
q, k, v, None, cu_seqlens_q, cu_seqlens_k, dropout_p, causal, return_softmax,
|
||
|
|
use_alibi, alibi_mode, imp_mode, window_size[0], window_size[1], None
|
||
|
|
)
|
||
|
|
return out, q, k, v, out_padded, softmax_lse, S_dmask
|
||
|
|
|
||
|
|
def _flash_attn_backward_ixdnn(dout, q, k, v, out, softmax_lse, dq, dk, dv,
|
||
|
|
cu_seqlens_q, cu_seqlens_k, dropout_p, causal, use_alibi, alibi_mode, imp_mode, window_size):
|
||
|
|
maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
|
||
|
|
# dq, dk, dv are allocated by us so they should already be contiguous
|
||
|
|
dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
|
||
|
|
dq, dk, dv, softmax_d, = flash_attn_cuda.bwd_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k, dropout_p,
|
||
|
|
causal, use_alibi, alibi_mode, imp_mode, window_size[0], window_size[1], None
|
||
|
|
)
|
||
|
|
return dq, dk, dv, softmax_d
|
||
|
|
|
||
|
|
def _flash_attn_varlen_forward_ixdnn(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, causal, use_alibi, alibi_mode, imp_mode, window_size):
|
||
|
|
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, S_dmask = flash_attn_cuda.varlen_fwd_ixdnn(
|
||
|
|
q, k, v, None, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p,
|
||
|
|
causal, use_alibi, alibi_mode, imp_mode, window_size[0], window_size[1], None
|
||
|
|
)
|
||
|
|
return out, q, k, v, out_padded, softmax_lse, S_dmask
|
||
|
|
|
||
|
|
def _flash_attn_varlen_backward_ixdnn(dout, q, k, v, out, softmax_lse, dq, dk, dv,
|
||
|
|
cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, causal, use_alibi, alibi_mode, imp_mode, window_size):
|
||
|
|
maybe_contiguous = lambda x: x.contiguous() if x.stride(-1) != 1 else x
|
||
|
|
# dq, dk, dv are allocated by us so they should already be contiguous
|
||
|
|
dout, q, k, v, out = [maybe_contiguous(x) for x in (dout, q, k, v, out)]
|
||
|
|
dq, dk, dv, softmax_d = flash_attn_cuda.varlen_bwd_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
|
||
|
|
max_seqlen_q, max_seqlen_k, dropout_p, causal, use_alibi, alibi_mode, imp_mode,
|
||
|
|
window_size[0], window_size[1], None
|
||
|
|
)
|
||
|
|
return dq, dk, dv, softmax_d
|
||
|
|
|
||
|
|
class FlashAttnQKVPackedFunc(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, qkv, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax, use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens = torch.full((qkv.shape[0],), qkv.shape[1], dtype=torch.int32, device=qkv.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward_ixdnn(
|
||
|
|
qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], cu_seqlens, cu_seqlens, dropout_p,
|
||
|
|
causal=causal, return_softmax=return_softmax and dropout_p > 0,
|
||
|
|
use_alibi=use_alibi, alibi_mode=alibi_mode, imp_mode=imp_mode,
|
||
|
|
window_size=window_size
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = qkv.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward(
|
||
|
|
qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2], dropout_p, softmax_scale,
|
||
|
|
causal=causal, return_softmax=return_softmax and dropout_p > 0,
|
||
|
|
use_alibi=use_alibi, alibi_mode=alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
if enable_ixdnn:
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
|
||
|
|
else:
|
||
|
|
q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
|
||
|
|
dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2],
|
||
|
|
cu_seqlens, cu_seqlens, ctx.dropout_p, ctx.causal, ctx.use_alibi, ctx.alibi_mode, ctx.imp_mode,
|
||
|
|
ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse, dqkv[:, :, 0], dqkv[:, :, 1], dqkv[:, :, 2],
|
||
|
|
ctx.dropout_p, ctx.softmax_scale, ctx.causal, ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
dqkv = dqkv[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dqkv, None, None, None, None, None, None, None, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
class FlashAttnVarlenQKVPackedFunc(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, qkv, cu_seqlens, max_seqlen, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax, use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens = torch.full((cu_seqlens.numel(),), max_seqlen, dtype=torch.int32, device=cu_seqlens.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward_ixdnn(
|
||
|
|
qkv[:, 0], qkv[:, 1], qkv[:, 2], cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
|
||
|
|
dropout_p, causal=causal, use_alibi=use_alibi, alibi_mode=alibi_mode, imp_mode=imp_mode,
|
||
|
|
window_size=window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = qkv.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward(
|
||
|
|
qkv[:, 0], qkv[:, 1], qkv[:, 2], cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
|
||
|
|
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax and dropout_p > 0,
|
||
|
|
use_alibi=use_alibi, alibi_mode=alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.max_seqlen = max_seqlen
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
qkv_shape = q.shape[:-2] + (3, *q.shape[-2:])
|
||
|
|
dqkv = torch.empty(qkv_shape, dtype=q.dtype, device=q.device)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_varlen_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dqkv[:, 0], dqkv[:, 1], dqkv[:, 2],
|
||
|
|
cu_seqlens, cu_seqlens, ctx.max_seqlen, ctx.max_seqlen,
|
||
|
|
ctx.dropout_p, ctx.causal, ctx.use_alibi, ctx.alibi_mode, ctx.imp_mode,
|
||
|
|
ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_varlen_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse, dqkv[:, 0], dqkv[:, 1], dqkv[:, 2],
|
||
|
|
cu_seqlens, cu_seqlens, ctx.max_seqlen, ctx.max_seqlen,
|
||
|
|
ctx.dropout_p, ctx.softmax_scale, ctx.causal, ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
dqkv = dqkv[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dqkv, None, None, None, None, None, None, None, None, None, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
class FlashAttnKVPackedFunc(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, q, kv, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax, use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens_q = torch.full((q.shape[0],), q.shape[1], dtype=torch.int32, device=q.device)
|
||
|
|
cu_seqlens_k = torch.full((kv.shape[0],), kv.shape[1], dtype=torch.int32, device=kv.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward_ixdnn(
|
||
|
|
q, kv[:, :, 0], kv[:, :, 1], cu_seqlens_q, cu_seqlens_k, dropout_p, causal=causal,
|
||
|
|
return_softmax=return_softmax and dropout_p > 0, use_alibi=use_alibi,
|
||
|
|
alibi_mode=alibi_mode, imp_mode=imp_mode, window_size=window_size
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = q.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward(
|
||
|
|
q, kv[:, :, 0], kv[:, :, 1], dropout_p, softmax_scale, causal=causal,
|
||
|
|
return_softmax=return_softmax and dropout_p > 0, use_alibi=use_alibi,
|
||
|
|
alibi_mode=alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
if enable_ixdnn:
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors
|
||
|
|
else:
|
||
|
|
q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
dq = torch.empty_like(q)
|
||
|
|
kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
|
||
|
|
dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dkv[:, :, 0], dkv[:, :, 1],
|
||
|
|
cu_seqlens_q, cu_seqlens_k, ctx.dropout_p, ctx.causal, ctx.use_alibi, ctx.alibi_mode, ctx.imp_mode, ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse,
|
||
|
|
dq, dkv[:, :, 0], dkv[:, :, 1], ctx.dropout_p, ctx.softmax_scale, ctx.causal, ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
dq = dq[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
dkv = dkv[..., :dout.shape[-1]]
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dq, dkv, None, None, None, None, None, None, None, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
class FlashAttnVarlenKVPackedFunc(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, q, kv, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p,
|
||
|
|
softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax,
|
||
|
|
use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens_q = torch.full((cu_seqlens_q.numel(),), max_seqlen_q, dtype=torch.int32, device=cu_seqlens_q.device)
|
||
|
|
cu_seqlens_k = torch.full((cu_seqlens_k.numel(),), max_seqlen_k, dtype=torch.int32, device=cu_seqlens_k.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward_ixdnn(
|
||
|
|
q, kv[:, 0], kv[:, 1], cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, causal=causal, use_alibi=use_alibi, alibi_mode=alibi_mode, imp_mode=imp_mode,
|
||
|
|
window_size=window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = q.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward(
|
||
|
|
q, kv[:, 0], kv[:, 1], cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax and dropout_p > 0,
|
||
|
|
use_alibi=use_alibi, alibi_mode=alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse,
|
||
|
|
cu_seqlens_q, cu_seqlens_k, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.max_seqlen_q = max_seqlen_q
|
||
|
|
ctx.max_seqlen_k = max_seqlen_k
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
dq = torch.empty_like(q)
|
||
|
|
kv_shape = k.shape[:-2] + (2, *k.shape[-2:])
|
||
|
|
dkv = torch.empty(kv_shape, dtype=k.dtype, device=k.device)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_varlen_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dkv[:, 0], dkv[:, 1],
|
||
|
|
cu_seqlens_q, cu_seqlens_k, ctx.max_seqlen_q, ctx.max_seqlen_k,
|
||
|
|
ctx.dropout_p, ctx.causal, ctx.use_alibi, ctx.alibi_mode, ctx.imp_mode,
|
||
|
|
ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_varlen_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dkv[:, 0], dkv[:, 1],
|
||
|
|
cu_seqlens_q, cu_seqlens_k, ctx.max_seqlen_q, ctx.max_seqlen_k,
|
||
|
|
ctx.dropout_p, ctx.softmax_scale, ctx.causal, ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
dq = dq[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
dkv = dkv[..., :dout.shape[-1]]
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dq, dkv, None, None, None, None, None, None, None, None, None, None, None, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
class FlashAttnFunc(torch.autograd.Function):
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def forward(ctx, q, k, v, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax, use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens_q = torch.full((q.shape[0],), q.shape[1], dtype=torch.int32, device=q.device)
|
||
|
|
cu_seqlens_k = torch.full((k.shape[0],), k.shape[1], dtype=torch.int32, device=k.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward_ixdnn(
|
||
|
|
q, k, v, cu_seqlens_q, cu_seqlens_k, dropout_p, causal=causal,
|
||
|
|
return_softmax=return_softmax and dropout_p > 0, use_alibi=use_alibi,
|
||
|
|
alibi_mode = alibi_mode, imp_mode = imp_mode, window_size=window_size
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = q.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_forward(
|
||
|
|
q, k, v, dropout_p, softmax_scale, causal=causal,
|
||
|
|
return_softmax=return_softmax and dropout_p > 0, use_alibi=use_alibi,
|
||
|
|
alibi_mode = alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
if enable_ixdnn:
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors
|
||
|
|
else:
|
||
|
|
q, k, v, out, softmax_lse, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse,
|
||
|
|
dq, dk, dv, cu_seqlens_q, cu_seqlens_k, ctx.dropout_p,
|
||
|
|
ctx.causal, ctx.use_alibi, ctx.alibi_mode, ctx.imp_mode,
|
||
|
|
ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse,
|
||
|
|
dq, dk, dv, ctx.dropout_p, ctx.softmax_scale, ctx.causal, ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
dq = dq[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
dk = dk[..., :dout.shape[-1]]
|
||
|
|
dv = dv[..., :dout.shape[-1]]
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dq, dk, dv, None, None, None, None, None, None, 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, dropout_p,
|
||
|
|
softmax_scale, causal, window_size, alibi_slopes, deterministic, return_softmax,
|
||
|
|
use_alibi, alibi_mode, imp_mode):
|
||
|
|
# Save rng_state because the backward pass will regenerate the dropout mask
|
||
|
|
rng_state = torch.cuda.get_rng_state() if dropout_p > 0 else None
|
||
|
|
if alibi_slopes is not None and use_alibi == False:
|
||
|
|
warnings.warn("Parameter 'alibi_slopes' is not supported, automatically switching to the supported parameters 'use_alibi=True' and 'alibi_mode=1'", UserWarning)
|
||
|
|
use_alibi = True
|
||
|
|
imp_mode = 1
|
||
|
|
if enable_ixdnn:
|
||
|
|
cu_seqlens_q = torch.full((cu_seqlens_q.numel(),), max_seqlen_q, dtype=torch.int32, device=cu_seqlens_q.device)
|
||
|
|
cu_seqlens_k = torch.full((cu_seqlens_k.numel(),), max_seqlen_k, dtype=torch.int32, device=cu_seqlens_k.device)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward_ixdnn(
|
||
|
|
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, causal=causal, use_alibi=use_alibi, alibi_mode = alibi_mode, imp_mode=imp_mode,
|
||
|
|
window_size=window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
if softmax_scale is None:
|
||
|
|
softmax_scale = q.shape[-1] ** (-0.5)
|
||
|
|
out, q, k, v, out_padded, softmax_lse, S_dmask = _flash_attn_varlen_forward(
|
||
|
|
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal=causal, return_softmax=return_softmax and dropout_p > 0,
|
||
|
|
use_alibi = use_alibi, alibi_mode = alibi_mode
|
||
|
|
)
|
||
|
|
ctx.save_for_backward(q, k, v, out_padded, softmax_lse,
|
||
|
|
cu_seqlens_q, cu_seqlens_k, rng_state)
|
||
|
|
ctx.dropout_p = dropout_p
|
||
|
|
ctx.max_seqlen_q = max_seqlen_q
|
||
|
|
ctx.max_seqlen_k = max_seqlen_k
|
||
|
|
ctx.softmax_scale = softmax_scale
|
||
|
|
ctx.causal = causal
|
||
|
|
ctx.use_alibi = use_alibi
|
||
|
|
ctx.alibi_mode = alibi_mode
|
||
|
|
ctx.imp_mode = imp_mode
|
||
|
|
ctx.window_size = window_size
|
||
|
|
return out if not return_softmax else (out, softmax_lse, S_dmask)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def backward(ctx, dout, *args):
|
||
|
|
q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, rng_state = ctx.saved_tensors
|
||
|
|
if rng_state is not None:
|
||
|
|
cur_rng_state = torch.cuda.get_rng_state()
|
||
|
|
torch.cuda.set_rng_state(rng_state)
|
||
|
|
dq, dk, dv = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)
|
||
|
|
if enable_ixdnn:
|
||
|
|
_flash_attn_varlen_backward_ixdnn(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
|
||
|
|
ctx.max_seqlen_q, ctx.max_seqlen_k, ctx.dropout_p, ctx.causal, ctx.use_alibi,
|
||
|
|
ctx.alibi_mode, ctx.imp_mode, ctx.window_size
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
_flash_attn_varlen_backward(
|
||
|
|
dout, q, k, v, out, softmax_lse, dq, dk, dv, cu_seqlens_q, cu_seqlens_k,
|
||
|
|
ctx.max_seqlen_q, ctx.max_seqlen_k, ctx.dropout_p, ctx.softmax_scale, ctx.causal,
|
||
|
|
ctx.use_alibi, ctx.alibi_mode
|
||
|
|
)
|
||
|
|
|
||
|
|
dq = dq[..., :dout.shape[-1]] # We could have padded the head dimension
|
||
|
|
dk = dk[..., :dout.shape[-1]]
|
||
|
|
dv = dv[..., :dout.shape[-1]]
|
||
|
|
if rng_state is not None:
|
||
|
|
torch.cuda.set_rng_state(cur_rng_state)
|
||
|
|
return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None, None, None
|
||
|
|
|
||
|
|
|
||
|
|
def flash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1), alibi_slopes=None,
|
||
|
|
deterministic=False, return_attn_probs=False, use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""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.
|
||
|
|
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 KV must be divisible by the number of heads in Q.
|
||
|
|
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 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.
|
||
|
|
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.
|
||
|
|
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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
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, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_attn_probs, use_alibi, alibi_mode, imp_mode)
|
||
|
|
|
||
|
|
|
||
|
|
def flash_attn_kvpacked_func(q, kv, dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1), alibi_slopes=None,
|
||
|
|
deterministic=False, return_attn_probs=False, use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""dropout_p should be set to 0.0 during evaluation
|
||
|
|
If 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 K, V.
|
||
|
|
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 KV must be divisible by the number of heads in Q.
|
||
|
|
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)
|
||
|
|
kv: (batch_size, seqlen, 2, 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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
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 FlashAttnKVPackedFunc.apply(q, kv, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_attn_probs, use_alibi, alibi_mode, imp_mode)
|
||
|
|
|
||
|
|
def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1), alibi_slopes=None,
|
||
|
|
deterministic=False, return_attn_probs=False, use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""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 KV must be divisible by the number of heads in Q.
|
||
|
|
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.
|
||
|
|
|
||
|
|
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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
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 FlashAttnFunc.apply(q, k, v, dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic, return_attn_probs, use_alibi, alibi_mode, imp_mode)
|
||
|
|
|
||
|
|
|
||
|
|
def flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens, max_seqlen, dropout_p=0.0, softmax_scale=None,
|
||
|
|
causal=False, window_size=(-1, -1), alibi_slopes=None, deterministic=False,
|
||
|
|
return_attn_probs=False, use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""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_varlen_func on Q, K, V since the backward pass avoids explicit concatenation
|
||
|
|
of the gradients of Q, K, V.
|
||
|
|
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.
|
||
|
|
|
||
|
|
Arguments:
|
||
|
|
qkv: (total, 3, nheads, headdim), where total = total number of tokens in the batch.
|
||
|
|
cu_seqlens: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
|
||
|
|
of the sequences in the batch, used to index into qkv.
|
||
|
|
max_seqlen: int. Maximum sequence length in the batch.
|
||
|
|
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 - 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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
Return:
|
||
|
|
out: (total, 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 FlashAttnVarlenQKVPackedFunc.apply(
|
||
|
|
qkv, cu_seqlens, max_seqlen, dropout_p, softmax_scale, causal, window_size,
|
||
|
|
alibi_slopes, deterministic, return_attn_probs, use_alibi, alibi_mode, imp_mode
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1),
|
||
|
|
alibi_slopes=None, deterministic=False, return_attn_probs=False,
|
||
|
|
use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""dropout_p should be set to 0.0 during evaluation
|
||
|
|
If 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 K, V.
|
||
|
|
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 KV must be divisible by the number of heads in Q.
|
||
|
|
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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
|
||
|
|
kv: (total_k, 2, nheads_k, headdim), where total_k = total number of key tokens in the batch.
|
||
|
|
cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
|
||
|
|
of the sequences in the batch, used to index into q.
|
||
|
|
cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
|
||
|
|
of the sequences in the batch, used to index into kv.
|
||
|
|
max_seqlen_q: int. Maximum query sequence length in the batch.
|
||
|
|
max_seqlen_k: int. Maximum key sequence length in the batch.
|
||
|
|
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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
Return:
|
||
|
|
out: (total, 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 FlashAttnVarlenKVPackedFunc.apply(
|
||
|
|
q, kv, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal, window_size, alibi_slopes,
|
||
|
|
deterministic, return_attn_probs, use_alibi, alibi_mode, imp_mode
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p=0.0, softmax_scale=None, causal=False, window_size=(-1, -1),
|
||
|
|
alibi_slopes=None, deterministic=False, return_attn_probs=False,
|
||
|
|
use_alibi=False, alibi_mode=1, imp_mode=0):
|
||
|
|
"""dropout_p should be set to 0.0 during evaluation
|
||
|
|
Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
|
||
|
|
than Q. Note that the number of heads in K, V must be divisible by the number of heads in Q.
|
||
|
|
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: (total_q, nheads, headdim), where total_q = total number of query tokens in the batch.
|
||
|
|
k: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
|
||
|
|
v: (total_k, nheads_k, headdim), where total_k = total number of key tokens in the batch.
|
||
|
|
cu_seqlens_q: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
|
||
|
|
of the sequences in the batch, used to index into q.
|
||
|
|
cu_seqlens_k: (batch_size + 1,), dtype torch.int32. The cumulative sequence lengths
|
||
|
|
of the sequences in the batch, used to index into kv.
|
||
|
|
max_seqlen_q: int. Maximum query sequence length in the batch.
|
||
|
|
max_seqlen_k: int. Maximum key sequence length in the batch.
|
||
|
|
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).
|
||
|
|
use_alibi: bool. Whether to apply attention with linear bias(ALiBi). if True, causal must
|
||
|
|
be True. The calculation formula(QK^T + bias * slops_m) is shown as follows,
|
||
|
|
q1*k1 0
|
||
|
|
q2*k1 q2*k2 -1 0
|
||
|
|
q3*k1 q3*k2 q3*k3 + -2 -1 0 * slops_m
|
||
|
|
q4*k1 q4*k2 q4*k3 q4*k4 -3 -2 -1 0
|
||
|
|
q5*k1 q5*k2 q5*k3 q5*k4 q5*k5 -4 -3 -2 -1 0
|
||
|
|
|
||
|
|
alibi_mode: int. The bias mode of ALiBi, default to 1.
|
||
|
|
alibi_mode=0: alibi_mode=1:
|
||
|
|
0 -√0
|
||
|
|
-1 0 -√1 -√0
|
||
|
|
-2 -1 0 -√2 -√1 -√0
|
||
|
|
-3 -2 -1 0 -√3 -√2 -√1 -√0
|
||
|
|
-4 -3 -2 -1 0 -√4 -√3 -√2 -√1 -√0
|
||
|
|
-5 -4 -3 -2 -1 0 -√5 -√4 -√3 -√2 -√1 -√0
|
||
|
|
imp_mode: int. Support two modes of backward implementation, default to 0.
|
||
|
|
imp_mode=0(CUDNN_FATTN_BALANCE_MODE): Implement dQ reduction calculation by launching a new kernel,
|
||
|
|
which will occupy additional memory. This mode usually has better performance.
|
||
|
|
imp_mode=1(CUDNN_FATTN_LEAST_MEM_MODE): Implement dQ reduction calculation within the block,
|
||
|
|
which will not occupy additional memory. There will be a slight performance loss in this mode.
|
||
|
|
This mode can be considered in memory-limited scenarios.
|
||
|
|
Return:
|
||
|
|
out: (total, 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 FlashAttnVarlenFunc.apply(
|
||
|
|
q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
|
||
|
|
dropout_p, softmax_scale, causal, window_size, alibi_slopes, deterministic,
|
||
|
|
return_attn_probs, use_alibi, alibi_mode, imp_mode
|
||
|
|
)
|