294 lines
14 KiB
Python
294 lines
14 KiB
Python
from typing import *
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import math
|
|
from .. import SparseTensor
|
|
from .. import config
|
|
|
|
|
|
__all__ = [
|
|
'sparse_windowed_scaled_dot_product_self_attention',
|
|
'sparse_windowed_scaled_dot_product_cross_attention',
|
|
]
|
|
|
|
|
|
def _sdpa_varlen_qkvpacked(qkv_feats: torch.Tensor, attn_func_args: dict) -> torch.Tensor:
|
|
"""sdpa for variable-length qkv-packed input (self-attention within windows)."""
|
|
q, k, v = qkv_feats.unbind(dim=1) # each [M, H, C]
|
|
cu_seqlens = attn_func_args['cu_seqlens']
|
|
seq_lens = attn_func_args['seq_lens']
|
|
N = len(seq_lens)
|
|
max_len = attn_func_args['max_seqlen'].item() if isinstance(attn_func_args['max_seqlen'], torch.Tensor) else attn_func_args['max_seqlen']
|
|
H, C = q.shape[-2], q.shape[-1]
|
|
|
|
# Pad into dense batch [N, max_len, H, C]
|
|
q_dense = q.new_zeros(N, max_len, H, C)
|
|
k_dense = k.new_zeros(N, max_len, H, C)
|
|
v_dense = v.new_zeros(N, max_len, H, C)
|
|
mask = torch.zeros(N, max_len, dtype=torch.bool, device=q.device)
|
|
for i in range(N):
|
|
sl = seq_lens[i].item() if isinstance(seq_lens[i], torch.Tensor) else seq_lens[i]
|
|
start = cu_seqlens[i].item()
|
|
q_dense[i, :sl] = q[start:start + sl]
|
|
k_dense[i, :sl] = k[start:start + sl]
|
|
v_dense[i, :sl] = v[start:start + sl]
|
|
mask[i, :sl] = True
|
|
|
|
# [N, H, L, C]
|
|
q_dense = q_dense.permute(0, 2, 1, 3)
|
|
k_dense = k_dense.permute(0, 2, 1, 3)
|
|
v_dense = v_dense.permute(0, 2, 1, 3)
|
|
|
|
# Build float mask for MPS compatibility
|
|
sdpa_mask = mask.unsqueeze(1).unsqueeze(2) # [N, 1, 1, L]
|
|
float_mask = torch.zeros(N, 1, max_len, max_len, dtype=q_dense.dtype, device=q.device)
|
|
float_mask.masked_fill_(~(mask.unsqueeze(1).unsqueeze(2) & mask.unsqueeze(1).unsqueeze(3)), float('-inf'))
|
|
|
|
out_dense = F.scaled_dot_product_attention(q_dense, k_dense, v_dense, attn_mask=float_mask)
|
|
out_dense = out_dense.permute(0, 2, 1, 3) # [N, L, H, C]
|
|
|
|
# Unpad
|
|
parts = []
|
|
for i in range(N):
|
|
sl = seq_lens[i].item() if isinstance(seq_lens[i], torch.Tensor) else seq_lens[i]
|
|
parts.append(out_dense[i, :sl])
|
|
return torch.cat(parts, dim=0)
|
|
|
|
|
|
def _sdpa_varlen(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
|
q_args: dict, kv_args: dict) -> torch.Tensor:
|
|
"""sdpa for variable-length cross-attention within windows."""
|
|
q_cu = q_args['cu_seqlens']
|
|
kv_cu = kv_args['cu_seqlens']
|
|
q_seq_lens = q_args['seq_lens']
|
|
kv_seq_lens = kv_args['seq_lens']
|
|
N = len(q_seq_lens)
|
|
max_q = q_args['max_seqlen'].item() if isinstance(q_args['max_seqlen'], torch.Tensor) else q_args['max_seqlen']
|
|
max_kv = kv_args['max_seqlen'].item() if isinstance(kv_args['max_seqlen'], torch.Tensor) else kv_args['max_seqlen']
|
|
H, C_q = q.shape[-2], q.shape[-1]
|
|
C_v = v.shape[-1]
|
|
|
|
q_dense = q.new_zeros(N, max_q, H, C_q)
|
|
k_dense = k.new_zeros(N, max_kv, H, C_q)
|
|
v_dense = v.new_zeros(N, max_kv, H, C_v)
|
|
q_mask = torch.zeros(N, max_q, dtype=torch.bool, device=q.device)
|
|
kv_mask = torch.zeros(N, max_kv, dtype=torch.bool, device=q.device)
|
|
for i in range(N):
|
|
ql = q_seq_lens[i].item() if isinstance(q_seq_lens[i], torch.Tensor) else q_seq_lens[i]
|
|
kvl = kv_seq_lens[i].item() if isinstance(kv_seq_lens[i], torch.Tensor) else kv_seq_lens[i]
|
|
qs = q_cu[i].item()
|
|
kvs = kv_cu[i].item()
|
|
q_dense[i, :ql] = q[qs:qs + ql]
|
|
k_dense[i, :kvl] = k[kvs:kvs + kvl]
|
|
v_dense[i, :kvl] = v[kvs:kvs + kvl]
|
|
q_mask[i, :ql] = True
|
|
kv_mask[i, :kvl] = True
|
|
|
|
q_dense = q_dense.permute(0, 2, 1, 3)
|
|
k_dense = k_dense.permute(0, 2, 1, 3)
|
|
v_dense = v_dense.permute(0, 2, 1, 3)
|
|
|
|
# q_mask: (N, max_q), kv_mask: (N, max_kv) -> cross mask: (N, 1, max_q, max_kv)
|
|
cross_mask = q_mask.unsqueeze(2) & kv_mask.unsqueeze(1) # (N, max_q, max_kv)
|
|
float_mask = torch.zeros(N, 1, max_q, max_kv, dtype=q_dense.dtype, device=q.device)
|
|
float_mask.masked_fill_(~cross_mask.unsqueeze(1), float('-inf'))
|
|
|
|
out_dense = F.scaled_dot_product_attention(q_dense, k_dense, v_dense, attn_mask=float_mask)
|
|
out_dense = out_dense.permute(0, 2, 1, 3)
|
|
|
|
parts = []
|
|
for i in range(N):
|
|
ql = q_seq_lens[i].item() if isinstance(q_seq_lens[i], torch.Tensor) else q_seq_lens[i]
|
|
parts.append(out_dense[i, :ql])
|
|
return torch.cat(parts, dim=0)
|
|
|
|
|
|
def calc_window_partition(
|
|
tensor: SparseTensor,
|
|
window_size: Union[int, Tuple[int, ...]],
|
|
shift_window: Union[int, Tuple[int, ...]] = 0,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
|
|
"""
|
|
Calculate serialization and partitioning for a set of coordinates.
|
|
|
|
Args:
|
|
tensor (SparseTensor): The input tensor.
|
|
window_size (int): The window size to use.
|
|
shift_window (Tuple[int, ...]): The shift of serialized coordinates.
|
|
|
|
Returns:
|
|
(torch.Tensor): Forwards indices.
|
|
(torch.Tensor): Backwards indices.
|
|
(torch.Tensor): Sequence lengths.
|
|
(dict): Attn func args.
|
|
"""
|
|
DIM = tensor.coords.shape[1] - 1
|
|
shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
|
|
window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
|
|
shifted_coords = tensor.coords.clone().detach()
|
|
shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
|
|
|
|
MAX_COORDS = [i + j for i, j in zip(tensor.spatial_shape, shift_window)]
|
|
NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
|
|
OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
|
|
|
|
shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
|
|
shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
|
|
fwd_indices = torch.argsort(shifted_indices)
|
|
bwd_indices = torch.empty_like(fwd_indices)
|
|
bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
|
|
seq_lens = torch.bincount(shifted_indices)
|
|
mask = seq_lens != 0
|
|
seq_lens = seq_lens[mask]
|
|
|
|
if config.ATTN == 'xformers':
|
|
if 'xops' not in globals():
|
|
import xformers.ops as xops
|
|
attn_func_args = {
|
|
'attn_bias': xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
|
|
}
|
|
elif config.ATTN == 'flash_attn':
|
|
attn_func_args = {
|
|
'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
|
|
'max_seqlen': torch.max(seq_lens)
|
|
}
|
|
elif config.ATTN == 'sdpa':
|
|
attn_func_args = {
|
|
'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
|
|
'max_seqlen': torch.max(seq_lens),
|
|
'seq_lens': seq_lens,
|
|
}
|
|
|
|
return fwd_indices, bwd_indices, seq_lens, attn_func_args
|
|
|
|
|
|
def sparse_windowed_scaled_dot_product_self_attention(
|
|
qkv: SparseTensor,
|
|
window_size: int,
|
|
shift_window: Tuple[int, int, int] = (0, 0, 0)
|
|
) -> SparseTensor:
|
|
"""
|
|
Apply windowed scaled dot product self attention to a sparse tensor.
|
|
|
|
Args:
|
|
qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
|
|
window_size (int): The window size to use.
|
|
shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
|
|
|
|
Returns:
|
|
(SparseTensor): [N, *, H, C] sparse tensor containing the output features.
|
|
"""
|
|
assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
|
|
|
|
serialization_spatial_cache_name = f'windowed_attention_{window_size}_{shift_window}'
|
|
serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
|
|
if serialization_spatial_cache is None:
|
|
fwd_indices, bwd_indices, seq_lens, attn_func_args = calc_window_partition(qkv, window_size, shift_window)
|
|
qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, attn_func_args))
|
|
else:
|
|
fwd_indices, bwd_indices, seq_lens, attn_func_args = serialization_spatial_cache
|
|
|
|
qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
|
|
|
|
if config.DEBUG:
|
|
start = 0
|
|
qkv_coords = qkv.coords[fwd_indices]
|
|
for i in range(len(seq_lens)):
|
|
seq_coords = qkv_coords[start:start+seq_lens[i]]
|
|
assert (seq_coords[:, 1:].max(dim=0).values - seq_coords[:, 1:].min(dim=0).values < window_size).all(), \
|
|
f"SparseWindowedScaledDotProductSelfAttention: window size exceeded"
|
|
start += seq_lens[i]
|
|
|
|
if config.ATTN == 'xformers':
|
|
if 'xops' not in globals():
|
|
import xformers.ops as xops
|
|
q, k, v = qkv_feats.unbind(dim=1) # [M, H, C]
|
|
q = q.unsqueeze(0) # [1, M, H, C]
|
|
k = k.unsqueeze(0) # [1, M, H, C]
|
|
v = v.unsqueeze(0) # [1, M, H, C]
|
|
out = xops.memory_efficient_attention(q, k, v, **attn_func_args)[0] # [M, H, C]
|
|
elif config.ATTN == 'flash_attn':
|
|
if 'flash_attn' not in globals():
|
|
import flash_attn
|
|
out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args) # [M, H, C]
|
|
elif config.ATTN == 'sdpa':
|
|
out = _sdpa_varlen_qkvpacked(qkv_feats, attn_func_args) # [M, H, C]
|
|
|
|
out = out[bwd_indices] # [T, H, C]
|
|
|
|
if config.DEBUG:
|
|
qkv_coords = qkv_coords[bwd_indices]
|
|
assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
|
|
|
|
return qkv.replace(out)
|
|
|
|
|
|
def sparse_windowed_scaled_dot_product_cross_attention(
|
|
q: SparseTensor,
|
|
kv: SparseTensor,
|
|
q_window_size: int,
|
|
kv_window_size: int,
|
|
q_shift_window: Tuple[int, int, int] = (0, 0, 0),
|
|
kv_shift_window: Tuple[int, int, int] = (0, 0, 0),
|
|
) -> SparseTensor:
|
|
"""
|
|
Apply windowed scaled dot product cross attention to two sparse tensors.
|
|
|
|
Args:
|
|
q (SparseTensor): [N, *, H, C] sparse tensor containing Qs.
|
|
kv (SparseTensor): [N, *, 2, H, C] sparse tensor containing Ks and Vs.
|
|
q_window_size (int): The window size to use for Qs.
|
|
kv_window_size (int): The window size to use for Ks and Vs.
|
|
q_shift_window (Tuple[int, int, int]): The shift of serialized coordinates for Qs.
|
|
kv_shift_window (Tuple[int, int, int]): The shift of serialized coordinates for Ks and Vs.
|
|
|
|
Returns:
|
|
(SparseTensor): [N, *, H, C] sparse tensor containing the output features.
|
|
"""
|
|
assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
|
|
assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
|
|
|
|
q_serialization_spatial_cache_name = f'windowed_attention_{q_window_size}_{q_shift_window}'
|
|
q_serialization_spatial_cache = q.get_spatial_cache(q_serialization_spatial_cache_name)
|
|
if q_serialization_spatial_cache is None:
|
|
q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args = calc_window_partition(q, q_window_size, q_shift_window)
|
|
q.register_spatial_cache(q_serialization_spatial_cache_name, (q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args))
|
|
else:
|
|
q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args = q_serialization_spatial_cache
|
|
kv_serialization_spatial_cache_name = f'windowed_attention_{kv_window_size}_{kv_shift_window}'
|
|
kv_serialization_spatial_cache = kv.get_spatial_cache(kv_serialization_spatial_cache_name)
|
|
if kv_serialization_spatial_cache is None:
|
|
kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args = calc_window_partition(kv, kv_window_size, kv_shift_window)
|
|
kv.register_spatial_cache(kv_serialization_spatial_cache_name, (kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args))
|
|
else:
|
|
kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args = kv_serialization_spatial_cache
|
|
|
|
assert len(q_seq_lens) == len(kv_seq_lens), "Number of sequences in q and kv must match"
|
|
|
|
q_feats = q.feats[q_fwd_indices] # [M, H, C]
|
|
kv_feats = kv.feats[kv_fwd_indices] # [M, 2, H, C]
|
|
|
|
if config.ATTN == 'xformers':
|
|
if 'xops' not in globals():
|
|
import xformers.ops as xops
|
|
k, v = kv_feats.unbind(dim=1) # [M, H, C]
|
|
q_feats_u = q_feats.unsqueeze(0) # [1, M, H, C]
|
|
k = k.unsqueeze(0) # [1, M, H, C]
|
|
v = v.unsqueeze(0) # [1, M, H, C]
|
|
mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seq_lens, kv_seq_lens)
|
|
out = xops.memory_efficient_attention(q_feats_u, k, v, attn_bias=mask)[0] # [M, H, C]
|
|
elif config.ATTN == 'flash_attn':
|
|
if 'flash_attn' not in globals():
|
|
import flash_attn
|
|
out = flash_attn.flash_attn_varlen_kvpacked_func(q_feats, kv_feats,
|
|
cu_seqlens_q=q_attn_func_args['cu_seqlens'], cu_seqlens_k=kv_attn_func_args['cu_seqlens'],
|
|
max_seqlen_q=q_attn_func_args['max_seqlen'], max_seqlen_k=kv_attn_func_args['max_seqlen'],
|
|
) # [M, H, C]
|
|
elif config.ATTN == 'sdpa':
|
|
k, v = kv_feats.unbind(dim=1)
|
|
out = _sdpa_varlen(q_feats, k, v, q_attn_func_args, kv_attn_func_args) # [M, H, C]
|
|
|
|
out = out[q_bwd_indices] # [T, H, C]
|
|
|
|
return q.replace(out)
|