1
This commit is contained in:
193
trellis/modules/sparse/attention/serialized_attn.py
Executable file
193
trellis/modules/sparse/attention/serialized_attn.py
Executable file
@@ -0,0 +1,193 @@
|
||||
from typing import *
|
||||
from enum import Enum
|
||||
import torch
|
||||
import math
|
||||
from .. import SparseTensor
|
||||
from .. import DEBUG, ATTN
|
||||
|
||||
if ATTN == 'xformers':
|
||||
import xformers.ops as xops
|
||||
elif ATTN == 'flash_attn':
|
||||
import flash_attn
|
||||
else:
|
||||
raise ValueError(f"Unknown attention module: {ATTN}")
|
||||
|
||||
|
||||
__all__ = [
|
||||
'sparse_serialized_scaled_dot_product_self_attention',
|
||||
]
|
||||
|
||||
|
||||
class SerializeMode(Enum):
|
||||
Z_ORDER = 0
|
||||
Z_ORDER_TRANSPOSED = 1
|
||||
HILBERT = 2
|
||||
HILBERT_TRANSPOSED = 3
|
||||
|
||||
|
||||
SerializeModes = [
|
||||
SerializeMode.Z_ORDER,
|
||||
SerializeMode.Z_ORDER_TRANSPOSED,
|
||||
SerializeMode.HILBERT,
|
||||
SerializeMode.HILBERT_TRANSPOSED
|
||||
]
|
||||
|
||||
|
||||
def calc_serialization(
|
||||
tensor: SparseTensor,
|
||||
window_size: int,
|
||||
serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
|
||||
shift_sequence: int = 0,
|
||||
shift_window: Tuple[int, int, int] = (0, 0, 0)
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, 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.
|
||||
serialize_mode (SerializeMode): The serialization mode to use.
|
||||
shift_sequence (int): The shift of serialized sequence.
|
||||
shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
|
||||
|
||||
Returns:
|
||||
(torch.Tensor, torch.Tensor): Forwards and backwards indices.
|
||||
"""
|
||||
fwd_indices = []
|
||||
bwd_indices = []
|
||||
seq_lens = []
|
||||
seq_batch_indices = []
|
||||
offsets = [0]
|
||||
|
||||
if 'vox2seq' not in globals():
|
||||
import vox2seq
|
||||
|
||||
# Serialize the input
|
||||
serialize_coords = tensor.coords[:, 1:].clone()
|
||||
serialize_coords += torch.tensor(shift_window, dtype=torch.int32, device=tensor.device).reshape(1, 3)
|
||||
if serialize_mode == SerializeMode.Z_ORDER:
|
||||
code = vox2seq.encode(serialize_coords, mode='z_order', permute=[0, 1, 2])
|
||||
elif serialize_mode == SerializeMode.Z_ORDER_TRANSPOSED:
|
||||
code = vox2seq.encode(serialize_coords, mode='z_order', permute=[1, 0, 2])
|
||||
elif serialize_mode == SerializeMode.HILBERT:
|
||||
code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[0, 1, 2])
|
||||
elif serialize_mode == SerializeMode.HILBERT_TRANSPOSED:
|
||||
code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[1, 0, 2])
|
||||
else:
|
||||
raise ValueError(f"Unknown serialize mode: {serialize_mode}")
|
||||
|
||||
for bi, s in enumerate(tensor.layout):
|
||||
num_points = s.stop - s.start
|
||||
num_windows = (num_points + window_size - 1) // window_size
|
||||
valid_window_size = num_points / num_windows
|
||||
to_ordered = torch.argsort(code[s.start:s.stop])
|
||||
if num_windows == 1:
|
||||
fwd_indices.append(to_ordered)
|
||||
bwd_indices.append(torch.zeros_like(to_ordered).scatter_(0, to_ordered, torch.arange(num_points, device=tensor.device)))
|
||||
fwd_indices[-1] += s.start
|
||||
bwd_indices[-1] += offsets[-1]
|
||||
seq_lens.append(num_points)
|
||||
seq_batch_indices.append(bi)
|
||||
offsets.append(offsets[-1] + seq_lens[-1])
|
||||
else:
|
||||
# Partition the input
|
||||
offset = 0
|
||||
mids = [(i + 0.5) * valid_window_size + shift_sequence for i in range(num_windows)]
|
||||
split = [math.floor(i * valid_window_size + shift_sequence) for i in range(num_windows + 1)]
|
||||
bwd_index = torch.zeros((num_points,), dtype=torch.int64, device=tensor.device)
|
||||
for i in range(num_windows):
|
||||
mid = mids[i]
|
||||
valid_start = split[i]
|
||||
valid_end = split[i + 1]
|
||||
padded_start = math.floor(mid - 0.5 * window_size)
|
||||
padded_end = padded_start + window_size
|
||||
fwd_indices.append(to_ordered[torch.arange(padded_start, padded_end, device=tensor.device) % num_points])
|
||||
offset += valid_start - padded_start
|
||||
bwd_index.scatter_(0, fwd_indices[-1][valid_start-padded_start:valid_end-padded_start], torch.arange(offset, offset + valid_end - valid_start, device=tensor.device))
|
||||
offset += padded_end - valid_start
|
||||
fwd_indices[-1] += s.start
|
||||
seq_lens.extend([window_size] * num_windows)
|
||||
seq_batch_indices.extend([bi] * num_windows)
|
||||
bwd_indices.append(bwd_index + offsets[-1])
|
||||
offsets.append(offsets[-1] + num_windows * window_size)
|
||||
|
||||
fwd_indices = torch.cat(fwd_indices)
|
||||
bwd_indices = torch.cat(bwd_indices)
|
||||
|
||||
return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
|
||||
|
||||
|
||||
def sparse_serialized_scaled_dot_product_self_attention(
|
||||
qkv: SparseTensor,
|
||||
window_size: int,
|
||||
serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
|
||||
shift_sequence: int = 0,
|
||||
shift_window: Tuple[int, int, int] = (0, 0, 0)
|
||||
) -> SparseTensor:
|
||||
"""
|
||||
Apply serialized 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.
|
||||
serialize_mode (SerializeMode): The serialization mode to use.
|
||||
shift_sequence (int): The shift of serialized sequence.
|
||||
shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
|
||||
shift (int): The shift to use.
|
||||
"""
|
||||
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'serialization_{serialize_mode}_{window_size}_{shift_sequence}_{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, seq_batch_indices = calc_serialization(qkv, window_size, serialize_mode, shift_sequence, shift_window)
|
||||
qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
|
||||
else:
|
||||
fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache
|
||||
|
||||
M = fwd_indices.shape[0]
|
||||
T = qkv.feats.shape[0]
|
||||
H = qkv.feats.shape[2]
|
||||
C = qkv.feats.shape[3]
|
||||
|
||||
qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
|
||||
|
||||
if DEBUG:
|
||||
start = 0
|
||||
qkv_coords = qkv.coords[fwd_indices]
|
||||
for i in range(len(seq_lens)):
|
||||
assert (qkv_coords[start:start+seq_lens[i], 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch"
|
||||
start += seq_lens[i]
|
||||
|
||||
if all([seq_len == window_size for seq_len in seq_lens]):
|
||||
B = len(seq_lens)
|
||||
N = window_size
|
||||
qkv_feats = qkv_feats.reshape(B, N, 3, H, C)
|
||||
if ATTN == 'xformers':
|
||||
q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C]
|
||||
out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C]
|
||||
elif ATTN == 'flash_attn':
|
||||
out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C]
|
||||
else:
|
||||
raise ValueError(f"Unknown attention module: {ATTN}")
|
||||
out = out.reshape(B * N, H, C) # [M, H, C]
|
||||
else:
|
||||
if ATTN == 'xformers':
|
||||
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]
|
||||
mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
|
||||
out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C]
|
||||
elif ATTN == 'flash_attn':
|
||||
cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
|
||||
.to(qkv.device).int()
|
||||
out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
|
||||
|
||||
out = out[bwd_indices] # [T, H, C]
|
||||
|
||||
if DEBUG:
|
||||
qkv_coords = qkv_coords[bwd_indices]
|
||||
assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
|
||||
|
||||
return qkv.replace(out)
|
||||
Reference in New Issue
Block a user