1050 lines
45 KiB
Python
1050 lines
45 KiB
Python
import functools
|
|
import math
|
|
import os
|
|
import warnings
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
try:
|
|
import flash_attn_interface
|
|
FLASH_ATTN_3_AVAILABLE = True
|
|
except ModuleNotFoundError:
|
|
FLASH_ATTN_3_AVAILABLE = False
|
|
|
|
try:
|
|
import flash_attn
|
|
FLASH_ATTN_2_AVAILABLE = True
|
|
except ModuleNotFoundError:
|
|
FLASH_ATTN_2_AVAILABLE = False
|
|
|
|
try:
|
|
major, minor = torch.cuda.get_device_capability(0)
|
|
if f"{major}.{minor}" == "8.0":
|
|
from sageattention_sm80 import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
elif f"{major}.{minor}" == "8.6":
|
|
from sageattention_sm86 import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
elif f"{major}.{minor}" == "8.9":
|
|
from sageattention_sm89 import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
elif f"{major}.{minor}" == "9.0":
|
|
from sageattention_sm90 import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
elif major > 9:
|
|
from sageattention_sm120 import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
except Exception:
|
|
try:
|
|
from sageattention import sageattn
|
|
SAGE_ATTENTION_AVAILABLE = True
|
|
except Exception:
|
|
sageattn = None
|
|
SAGE_ATTENTION_AVAILABLE = False
|
|
|
|
|
|
from .attention_kernel import (_sparse_linear_attention, get_block_map,
|
|
mean_pool, sparse_linear_attention_varlen)
|
|
|
|
|
|
def convert_qkv_dtype(q, k, v):
|
|
try:
|
|
"""Unify the dtype of q, k, v tensors"""
|
|
dtypes = {q.dtype, k.dtype, v.dtype}
|
|
|
|
# If any tensor is float16/bfloat16
|
|
if torch.float16 in dtypes or torch.bfloat16 in dtypes:
|
|
target_dtype = torch.bfloat16 if torch.bfloat16 in dtypes else torch.float16
|
|
# If all tensors are float32
|
|
elif dtypes == {torch.float32}:
|
|
target_dtype = torch.bfloat16 if (torch.cuda.is_available() and
|
|
torch.cuda.get_device_capability()[0] >= 8) else torch.float16
|
|
else:
|
|
return q, k, v # No conversion for other cases
|
|
|
|
return q.to(target_dtype), k.to(target_dtype), v.to(target_dtype)
|
|
except Exception:
|
|
return q, k, v
|
|
|
|
|
|
def _convert_attn_mask_to_lens(attn_mask):
|
|
"""
|
|
Convert attention mask to sequence lengths for Flash Attention.
|
|
|
|
Args:
|
|
attn_mask: Attention mask, can be:
|
|
- [B, L] with 1=valid, 0=padding
|
|
- [B, 1, L] or [B, 1, 1, L] attention bias with 0=valid, -inf/-10000=padding
|
|
- [B, H, Lq, Lk] full attention mask
|
|
|
|
Returns:
|
|
k_lens: [B] tensor of valid sequence lengths, or None if not a simple padding mask
|
|
"""
|
|
if attn_mask is None:
|
|
return None
|
|
|
|
# Squeeze to simplest form
|
|
while attn_mask.ndim > 2 and attn_mask.shape[1] == 1:
|
|
attn_mask = attn_mask.squeeze(1)
|
|
|
|
# Only handle [B, L] case (simple padding mask)
|
|
if attn_mask.ndim != 2:
|
|
return None
|
|
|
|
# Check if it's attention bias format (0 and -inf/-10000) or binary mask (0/1)
|
|
unique_vals = torch.unique(attn_mask)
|
|
if len(unique_vals) > 2:
|
|
return None # Complex mask, can't convert
|
|
|
|
# Determine which value means "valid"
|
|
max_val = unique_vals.max().item()
|
|
|
|
if max_val <= 0: # Attention bias format: 0=valid, negative=padding
|
|
valid_mask = (attn_mask >= -1.0) # 0 is valid
|
|
else: # Binary format: 1=valid, 0=padding
|
|
valid_mask = (attn_mask > 0.5)
|
|
|
|
# Check if it's a simple left-padded or right-padded mask
|
|
# For right-padding: [1,1,1,0,0] -> valid tokens are contiguous from start
|
|
k_lens = valid_mask.sum(dim=-1).to(torch.int32)
|
|
|
|
# Verify it's actually a contiguous padding mask by reconstruction
|
|
B, L = valid_mask.shape
|
|
reconstructed = torch.arange(L, device=valid_mask.device).unsqueeze(0) < k_lens.unsqueeze(1)
|
|
if not torch.all(reconstructed == valid_mask):
|
|
return None # Not a simple contiguous padding mask
|
|
|
|
return k_lens
|
|
|
|
|
|
def flash_attention_naive(
|
|
q,
|
|
k,
|
|
v,
|
|
cu_seqlens_q=None,
|
|
cu_seqlens_k=None,
|
|
max_seqlen_q=None,
|
|
max_seqlen_k=None,
|
|
):
|
|
# apply attention
|
|
if FLASH_ATTN_3_AVAILABLE:
|
|
# Note: dropout_p, window_size are not supported in FA3 now.
|
|
x = flash_attn_interface.flash_attn_varlen_func(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_k=cu_seqlens_k,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_seqlen_k=max_seqlen_k,
|
|
)[0]
|
|
else:
|
|
assert FLASH_ATTN_2_AVAILABLE
|
|
x = flash_attn.flash_attn_varlen_func(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_k=cu_seqlens_k,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_seqlen_k=max_seqlen_k,
|
|
)
|
|
|
|
# output
|
|
return x.type(q.dtype)
|
|
|
|
|
|
def flash_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
q_lens=None,
|
|
k_lens=None,
|
|
dropout_p=0.,
|
|
softmax_scale=None,
|
|
q_scale=None,
|
|
causal=False,
|
|
window_size=(-1, -1),
|
|
deterministic=False,
|
|
dtype=torch.bfloat16,
|
|
version=None,
|
|
):
|
|
"""
|
|
q: [B, Lq, Nq, C1].
|
|
k: [B, Lk, Nk, C1].
|
|
v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
|
|
q_lens: [B].
|
|
k_lens: [B].
|
|
dropout_p: float. Dropout probability.
|
|
softmax_scale: float. The scaling of QK^T before applying softmax.
|
|
causal: bool. Whether to apply causal attention mask.
|
|
window_size: (left right). If not (-1, -1), apply sliding window local attention.
|
|
deterministic: bool. If True, slightly slower and uses more memory.
|
|
dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
|
|
"""
|
|
half_dtypes = (torch.float16, torch.bfloat16)
|
|
assert dtype in half_dtypes
|
|
assert q.device.type == 'cuda' and q.size(-1) <= 256
|
|
|
|
# params
|
|
b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
|
|
|
|
def half(x):
|
|
return x if x.dtype in half_dtypes else x.to(dtype)
|
|
|
|
# preprocess query
|
|
if q_lens is None:
|
|
q = half(q.flatten(0, 1))
|
|
q_lens = torch.tensor(
|
|
[lq] * b, dtype=torch.int32).to(
|
|
device=q.device, non_blocking=True)
|
|
else:
|
|
q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
|
|
|
|
# preprocess key, value
|
|
if k_lens is None:
|
|
k = half(k.flatten(0, 1))
|
|
v = half(v.flatten(0, 1))
|
|
k_lens = torch.tensor(
|
|
[lk] * b, dtype=torch.int32).to(
|
|
device=k.device, non_blocking=True)
|
|
else:
|
|
k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
|
|
v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
|
|
|
|
q = q.to(v.dtype)
|
|
k = k.to(v.dtype)
|
|
|
|
if q_scale is not None:
|
|
q = q * q_scale
|
|
|
|
if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
|
|
warnings.warn(
|
|
'Flash attention 3 is not available, use flash attention 2 instead.'
|
|
)
|
|
|
|
# apply attention
|
|
if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
|
|
# Note: dropout_p, window_size are not supported in FA3 now.
|
|
x = flash_attn_interface.flash_attn_varlen_func(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
|
|
0, dtype=torch.int32).to(q.device, non_blocking=True),
|
|
cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
|
|
0, dtype=torch.int32).to(q.device, non_blocking=True),
|
|
seqused_q=None,
|
|
seqused_k=None,
|
|
max_seqlen_q=lq,
|
|
max_seqlen_k=lk,
|
|
softmax_scale=softmax_scale,
|
|
causal=causal,
|
|
deterministic=deterministic)
|
|
if isinstance(x, (tuple, list)):
|
|
x = x[0].unflatten(0, (b, lq))
|
|
else:
|
|
x = x.unflatten(0, (b, lq))
|
|
else:
|
|
assert FLASH_ATTN_2_AVAILABLE
|
|
x = flash_attn.flash_attn_varlen_func(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
|
|
0, dtype=torch.int32).to(q.device, non_blocking=True),
|
|
cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
|
|
0, dtype=torch.int32).to(q.device, non_blocking=True),
|
|
max_seqlen_q=lq,
|
|
max_seqlen_k=lk,
|
|
dropout_p=dropout_p,
|
|
softmax_scale=softmax_scale,
|
|
causal=causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic).unflatten(0, (b, lq))
|
|
|
|
# output
|
|
return x.type(out_dtype)
|
|
|
|
|
|
def attention(
|
|
q,
|
|
k,
|
|
v,
|
|
q_lens=None,
|
|
k_lens=None,
|
|
dropout_p=0.,
|
|
softmax_scale=None,
|
|
q_scale=None,
|
|
causal=False,
|
|
window_size=(-1, -1),
|
|
deterministic=False,
|
|
dtype=torch.bfloat16,
|
|
fa_version=None,
|
|
attention_type=None,
|
|
attn_mask=None,
|
|
):
|
|
attention_type = os.environ.get("VIDEOX_ATTENTION_TYPE", "FLASH_ATTENTION") if attention_type is None else attention_type
|
|
if fa_version is None and os.environ.get("VIDEOX_FA_VERSION") is not None:
|
|
fa_version = int(os.environ["VIDEOX_FA_VERSION"])
|
|
if torch.is_grad_enabled() and attention_type == "SAGE_ATTENTION":
|
|
attention_type = "FLASH_ATTENTION"
|
|
|
|
# Convert attn_mask to k_lens for Flash Attention if possible
|
|
# Note: flash_attention doesn't support variable-length query, only set k_lens
|
|
if attn_mask is not None and k_lens is None and attention_type == "FLASH_ATTENTION":
|
|
converted_lens = _convert_attn_mask_to_lens(attn_mask)
|
|
if converted_lens is not None:
|
|
k_lens = converted_lens
|
|
attn_mask = None # Successfully converted, clear the mask
|
|
else:
|
|
# Conversion failed, fallback to SDPA which supports attn_mask
|
|
attention_type = "SDPA"
|
|
|
|
if attention_type == "SAGE_ATTENTION" and SAGE_ATTENTION_AVAILABLE:
|
|
if q_lens is not None or k_lens is not None:
|
|
warnings.warn(
|
|
'Padding mask is disabled when using SAGE_ATTENTION. It can have a significant impact on performance.'
|
|
)
|
|
|
|
q, k, v = convert_qkv_dtype(q, k, v)
|
|
out = sageattn(
|
|
q, k, v, attn_mask=attn_mask, tensor_layout="NHD", is_causal=causal, dropout_p=dropout_p)
|
|
|
|
elif attention_type == "FLASH_ATTENTION" and (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE):
|
|
return flash_attention(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
q_lens=q_lens,
|
|
k_lens=k_lens,
|
|
dropout_p=dropout_p,
|
|
softmax_scale=softmax_scale,
|
|
q_scale=q_scale,
|
|
causal=causal,
|
|
window_size=window_size,
|
|
deterministic=deterministic,
|
|
dtype=dtype,
|
|
version=fa_version,
|
|
)
|
|
else:
|
|
if q_lens is not None or k_lens is not None:
|
|
warnings.warn(
|
|
'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
|
|
)
|
|
q = q.transpose(1, 2)
|
|
k = k.transpose(1, 2)
|
|
v = v.transpose(1, 2)
|
|
|
|
out = torch.nn.functional.scaled_dot_product_attention(
|
|
q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
|
|
|
|
out = out.transpose(1, 2).contiguous()
|
|
return out
|
|
|
|
|
|
class SparseLinearAttention(nn.Module):
|
|
# Modified from https://github.com/thu-ml/TurboDiffusion/blob/main/turbodiffusion/SLA/kernel.py
|
|
def __init__(self, head_dim, topk, feature_map='softmax', BLKQ=64, BLKK=64, use_bf16=True, tie_feature_map_qk=True):
|
|
R'''
|
|
Args:
|
|
head_dim: dimension of each head.
|
|
topk: ratio of keys selected for sparse attention, shared across all queries.
|
|
feature_map: feature map for linear attention, one of ['hedgehog', 'elu', 'relu', 'softmax'].
|
|
BLKQ: block size for query.
|
|
BLKK: block size for key.
|
|
use_bf16: whether to use bfloat16 (default) or float16 for computation. The conversion to bf16/fp16 is done inside the module.
|
|
tie_feature_map_qk: whether to use the same feature map for query and key.
|
|
'''
|
|
super().__init__()
|
|
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
|
self.topk = topk
|
|
self.BLKQ = BLKQ
|
|
self.BLKK = BLKK
|
|
self.proj_l = nn.Linear(head_dim, head_dim, dtype=torch.float32)
|
|
|
|
if feature_map == 'elu':
|
|
def elu_feature_map(x):
|
|
return torch.nn.functional.elu(x) + 1
|
|
self.feature_map_q = elu_feature_map
|
|
self.feature_map_k = elu_feature_map
|
|
elif feature_map == 'relu':
|
|
self.feature_map_q = nn.ReLU()
|
|
self.feature_map_k = nn.ReLU()
|
|
elif feature_map == 'softmax':
|
|
def softmax_feature_map(x):
|
|
return torch.nn.functional.softmax(x, dim=-1)
|
|
self.feature_map_q = softmax_feature_map
|
|
self.feature_map_k = softmax_feature_map
|
|
else:
|
|
raise NotImplementedError(f'Not supported feature map {feature_map}.')
|
|
|
|
if tie_feature_map_qk:
|
|
self.feature_map_k = self.feature_map_q
|
|
|
|
self.init_weights_()
|
|
|
|
def init_weights_(self):
|
|
with torch.no_grad():
|
|
nn.init.zeros_(self.proj_l.weight)
|
|
nn.init.zeros_(self.proj_l.bias)
|
|
|
|
def forward(self, q, k, v, return_sparsity=False):
|
|
R'''
|
|
Args:
|
|
q: queries [B, L, H, D]
|
|
k: keys [B, L, H, D].
|
|
v: values [B, L, H, D].
|
|
return_sparsity: whether to return the actual sparsity.
|
|
'''
|
|
dtype = q.dtype
|
|
|
|
q = q.transpose(1, 2).contiguous()
|
|
k = k.transpose(1, 2).contiguous()
|
|
v = v.transpose(1, 2).contiguous()
|
|
|
|
sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
|
|
|
|
q = q.to(self.dtype)
|
|
k = k.to(self.dtype)
|
|
v = v.to(self.dtype)
|
|
o_s = _sparse_linear_attention.apply(q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK)
|
|
|
|
q = self.feature_map_q(q).contiguous().to(self.dtype) # c_q
|
|
k = self.feature_map_k(k).contiguous().to(self.dtype) # c_k
|
|
|
|
def calc_linear(q, k, v):
|
|
kvsum = k.transpose(-1, -2) @ v
|
|
ksum = torch.sum(k, dim=-2, keepdim=True)
|
|
return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True))
|
|
o_l = calc_linear(q, k, v)
|
|
|
|
with torch.amp.autocast('cuda', dtype=self.dtype):
|
|
o_l = self.proj_l(o_l)
|
|
o = (o_s + o_l).to(dtype).transpose(1, 2)
|
|
|
|
if return_sparsity:
|
|
return o, real_topk / sparse_map.shape[-1]
|
|
else:
|
|
return o
|
|
|
|
|
|
class RoLASparseAttention(nn.Module):
|
|
R'''
|
|
RoLa sparse + low-rank self-attention branch (the official SparkDiffusion default).
|
|
|
|
Mirrors ``sparkdiffusion/networks/wan_rola_attention.py::WanSelfAttentionRoLa``:
|
|
- sparse branch : block-sparse exact softmax via the in-repo SLA Triton kernel
|
|
(same ``get_block_map`` + ``_sparse_linear_attention`` used by
|
|
``SparseLinearAttention``; runs on post-RoPE q/k).
|
|
- low-rank branch : ``proj_q``/``proj_k`` (head_dim -> rank) -> SiLU -> rank-truncated
|
|
3D RoPE -> linear attention ``Q_lr (K_lr^T V)`` -> RMSNorm
|
|
(runs on PRE-RoPE, qk-normed q/k; RoPE is applied OUTSIDE the
|
|
non-linear feature map, on the low-rank bottleneck).
|
|
- token gate : ``sigmoid(gate_proj(x) + gate_bias)`` scales the low-rank output
|
|
magnitude per token/head; ``gate`` is computed from the raw hidden
|
|
states ``x`` (not q/k/v).
|
|
|
|
Fusion (per official)::
|
|
|
|
scale = sqrt(mean(o_s^2) + 1e-5) # per-token/head RMS of the sparse output
|
|
out = (o_s / scale + gate * o_l) * scale # == o_s + gate * o_l * scale
|
|
|
|
New learnable params per layer (all under ``attn_op.*``): ``proj_q``, ``proj_k``,
|
|
``gate_proj``, ``gate_bias``. They are absent from a dense checkpoint (loaded with
|
|
strict=False) and start at their init below. This replaces the earlier PureSLA path
|
|
(single zero-init ``proj_l`` + softmax feature-map linear attention).
|
|
|
|
All forward helpers take/return BHLD (``[B, H, L, *]``) tensors.
|
|
'''
|
|
def __init__(self, head_dim, num_heads, topk=0.03, rola_rank=64,
|
|
BLKQ=64, BLKK=64, use_bf16=True):
|
|
R'''
|
|
Args:
|
|
head_dim: dimension of each head.
|
|
num_heads: number of attention heads (sizes ``gate_bias``).
|
|
topk: ratio of key blocks kept for the sparse branch (0.1 -> 90% sparse).
|
|
rola_rank: low-rank bottleneck dimension of the linear branch.
|
|
BLKQ/BLKK: query/key block sizes for the sparse kernel.
|
|
use_bf16: compute the sparse kernel in bfloat16 (default) or float16.
|
|
'''
|
|
super().__init__()
|
|
self.head_dim = head_dim
|
|
self.num_heads = num_heads
|
|
self.topk = topk
|
|
self.rola_rank = rola_rank
|
|
self.BLKQ = BLKQ
|
|
self.BLKK = BLKK
|
|
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
|
|
|
r = rola_rank
|
|
self.proj_q = nn.Linear(head_dim, r, bias=False)
|
|
self.proj_k = nn.Linear(head_dim, r, bias=False)
|
|
self.gate_proj = nn.Linear(head_dim, 1, bias=False)
|
|
# -1.946 -> sigmoid ~= 0.125, so the sparse branch dominates at init.
|
|
self.gate_bias = nn.Parameter(torch.full((1, num_heads, 1, 1), -1.946))
|
|
|
|
self.init_weights_()
|
|
|
|
def init_weights_(self):
|
|
nn.init.kaiming_uniform_(self.proj_q.weight, a=math.sqrt(5))
|
|
nn.init.kaiming_uniform_(self.proj_k.weight, a=math.sqrt(5))
|
|
nn.init.zeros_(self.gate_proj.weight)
|
|
with torch.no_grad():
|
|
self.gate_bias.fill_(-1.946)
|
|
|
|
def sparse(self, q, k, v):
|
|
R'''Block-sparse exact softmax attention.
|
|
|
|
Args:
|
|
q, k: [B, H, L, D] POST-RoPE queries/keys.
|
|
v: [B, H, L, D] values.
|
|
Returns:
|
|
[B, H, L, D], or ``None`` when the sequence is shorter than one block
|
|
(the caller then falls back to ``gate * o_l``).
|
|
'''
|
|
L = q.shape[2]
|
|
if L < max(self.BLKQ, self.BLKK):
|
|
return None
|
|
orig_dtype = q.dtype
|
|
q = q.contiguous()
|
|
k = k.contiguous()
|
|
v = v.contiguous()
|
|
sparse_map, lut, real_topk = get_block_map(
|
|
q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
|
|
q = q.to(self.dtype)
|
|
k = k.to(self.dtype)
|
|
v = v.to(self.dtype)
|
|
o_s = _sparse_linear_attention.apply(
|
|
q, k, v, sparse_map, lut, real_topk, self.BLKQ, self.BLKK)
|
|
return o_s.to(orig_dtype)
|
|
|
|
def lowrank(self, q_normed, k_normed, v, grid_sizes, freqs):
|
|
R'''Low-rank linear attention branch with rank-truncated RoPE.
|
|
|
|
Args:
|
|
q_normed, k_normed: [B, H, L, D] PRE-RoPE, qk-normed queries/keys.
|
|
v: [B, H, L, D] values.
|
|
grid_sizes: [B, 3] (F, H, W) for the rank RoPE expansion.
|
|
freqs: complex rotary table [M, head_dim // 2].
|
|
Returns:
|
|
[B, H, L, D], RMS-normalized.
|
|
'''
|
|
from .spark_transformer3d import rope_apply_lowrank
|
|
|
|
q_lr = F.silu(self.proj_q(q_normed)) # [B, H, L, rank]
|
|
k_lr = F.silu(self.proj_k(k_normed))
|
|
# RoPE applied outside the non-linear feature map, on the rank bottleneck.
|
|
q_lr = rope_apply_lowrank(q_lr, grid_sizes, freqs, self.rola_rank)
|
|
k_lr = rope_apply_lowrank(k_lr, grid_sizes, freqs, self.rola_rank)
|
|
c = k_lr.transpose(-2, -1) @ v # [B, H, rank, D]
|
|
out = q_lr @ c # [B, H, L, D]
|
|
return F.rms_norm(out, [out.shape[-1]])
|
|
|
|
def gate(self, x_bhld, head_slice=None):
|
|
R'''Token-wise gate from the raw hidden states.
|
|
|
|
Args:
|
|
x_bhld: [B, H, L, D] hidden states (attention input), head layout.
|
|
head_slice: optional slice selecting this rank's head shard (multi-GPU
|
|
Ulysses head-parallel), so the per-head ``gate_bias`` lines up.
|
|
Returns:
|
|
[B, H, L, 1] gate in (0, 1).
|
|
'''
|
|
gb = self.gate_bias if head_slice is None else self.gate_bias[:, head_slice]
|
|
return torch.sigmoid(self.gate_proj(x_bhld) + gb)
|
|
|
|
|
|
# Logical tile geometries keyed by tile volume -- the official VSA_H3_TILE_SHAPES. 64 is what the
|
|
# official defaults actually select for the checkpoints this repo runs: the Wan backend's
|
|
# VSA_TILE_SIZE is (4,4,4), and both released FastH3 inference contracts pin "vsa_tile_size": 64.
|
|
# (256/(4,8,8) is the H3 backend's *code* default, aimed at the Blackwell CuTe fastpath.) 128 and
|
|
# 256 stay selectable and reach the same 64-block kernel through the route-A expansion below.
|
|
_H3_TILE_SHAPES = {64: (4, 4, 4), 128: (4, 4, 8), 256: (4, 8, 8)}
|
|
_H3_TILE_ELEMS = 64
|
|
# Physical block size the Triton kernels are instantiated at: a logical tile of N elems is N // 64
|
|
# CONTIGUOUS 64-blocks in the padded tile buffer.
|
|
_H3_KERNEL_BLOCK = 64
|
|
|
|
|
|
def _h3_compute_topk(sparsity, num_blocks):
|
|
"""Key tiles to KEEP for a sparsity level, clamped to [1, num_blocks]."""
|
|
return max(1, min(math.ceil((1.0 - sparsity) * num_blocks), num_blocks))
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _h3_tile_partition_indices(grid, tile_shape, device):
|
|
"""Region tile-order -> original flat region index (port of get_tile_partition_indices)."""
|
|
T, H, W = grid
|
|
ts, hs, ws = tile_shape
|
|
indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W)
|
|
parts = []
|
|
for t in range(math.ceil(T / ts)):
|
|
for h in range(math.ceil(H / hs)):
|
|
for w in range(math.ceil(W / ws)):
|
|
parts.append(indices[t * ts:min(t * ts + ts, T),
|
|
h * hs:min(h * hs + hs, H),
|
|
w * ws:min(w * ws + ws, W)].flatten())
|
|
return torch.cat(parts, dim=0)
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _h3_variable_block_sizes(grid, num_tiles, device, tile_shape):
|
|
"""Valid (non-pad) token count inside every 3-D tile (port of construct_variable_block_sizes)."""
|
|
t, h, w = grid
|
|
ts_t, ts_h, ts_w = tile_shape
|
|
n_t, n_h, n_w = num_tiles
|
|
|
|
def _sizes(dim_len, tile, n):
|
|
sizes = torch.full((n,), tile, dtype=torch.long, device=device)
|
|
rem = dim_len - (n - 1) * tile
|
|
sizes[-1] = rem if rem > 0 else tile
|
|
return sizes
|
|
|
|
t_s = _sizes(t, ts_t, n_t)
|
|
h_s = _sizes(h, ts_h, n_h)
|
|
w_s = _sizes(w, ts_w, n_w)
|
|
return (t_s[:, None, None] * h_s[None, :, None] * w_s[None, None, :]).reshape(-1)
|
|
|
|
|
|
def _h3_non_pad_index(variable_block_sizes, tile_elems):
|
|
"""Padded-buffer slot index of every real row, in tile order (port of get_non_pad_index)."""
|
|
n_win = variable_block_sizes.shape[0]
|
|
device = variable_block_sizes.device
|
|
starts = torch.arange(n_win, device=device) * tile_elems
|
|
index_pad = starts[:, None] + torch.arange(tile_elems, device=device)[None, :]
|
|
index_mask = torch.arange(tile_elems, device=device)[None, :] < variable_block_sizes[:, None]
|
|
return index_pad[index_mask]
|
|
|
|
|
|
def _h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems):
|
|
R'''Fail synchronously on out-of-bounds tile geometry (port of the official guard).
|
|
|
|
Invariants the block-sparse kernels trust without checking: every tile's valid size is in
|
|
(0, tile_elems]; the sizes sum to the packed sequence length; and ``untile_combined_index``
|
|
maps each packed row to exactly one non-pad slot of the padded tile buffer. A violation would
|
|
otherwise surface only as an async device fault at some later kernel or collective (e.g. an
|
|
FSDP all-gather), which is unattributable -- so raise here, once per cached geometry.
|
|
'''
|
|
total = sum(math.prod(seg) if isinstance(seg, tuple) else seg for seg in segments)
|
|
n_pad = variable_block_sizes.numel() * tile_elems
|
|
sizes_min = int(variable_block_sizes.min())
|
|
sizes_max = int(variable_block_sizes.max())
|
|
sizes_sum = int(variable_block_sizes.sum())
|
|
if sizes_min < 1 or sizes_max > tile_elems or sizes_sum != total:
|
|
raise ValueError(
|
|
f"VSA-H3 tile sizes out of bounds for segments={segments}, tile_elems={tile_elems}: "
|
|
f"min={sizes_min}, max={sizes_max}, sum={sizes_sum}, expected sum={total}.")
|
|
if untile_combined_index.numel() != total:
|
|
raise ValueError(
|
|
f"VSA-H3 untile index has {untile_combined_index.numel()} entries for a packed "
|
|
f"sequence of {total} rows (segments={segments}).")
|
|
idx_min = int(untile_combined_index.min())
|
|
idx_max = int(untile_combined_index.max())
|
|
if idx_min < 0 or idx_max >= n_pad:
|
|
# Range first: the pad-slot gather below would itself index out of bounds (the very async
|
|
# fault this guard exists to preempt).
|
|
raise ValueError(
|
|
f"VSA-H3 untile index is not an injective map into non-pad slots: range "
|
|
f"[{idx_min}, {idx_max}] vs padded length {n_pad} (segments={segments}).")
|
|
in_tile_offset = untile_combined_index % tile_elems
|
|
maps_into_pad = bool(
|
|
(in_tile_offset >= variable_block_sizes[untile_combined_index // tile_elems]).any())
|
|
if maps_into_pad or int(torch.unique(untile_combined_index).numel()) != total:
|
|
raise ValueError(
|
|
f"VSA-H3 untile index is not an injective map into non-pad slots: "
|
|
f"pad-slot hit={maps_into_pad} (segments={segments}).")
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def build_h3_tile_geometry(segments, device, tile_shape=_H3_TILE_SHAPES[_H3_TILE_ELEMS]):
|
|
R'''Tile a packed sequence of interleaved dense segments and video regions.
|
|
|
|
Port of the official ``_h3_segment_tile_geometry``. ``segments`` lists, in packed
|
|
order, dense segment row counts (``int``) and video-region token grids (``(t,h,w)``
|
|
triples). Returns ``(tile_partition_indices, variable_block_sizes,
|
|
untile_combined_index, num_prefix_tiles, num_video_tiles, video_tile_spans,
|
|
tile_elems)`` where ``video_tile_spans`` holds one ``[start, end)`` tile range per
|
|
video region (prefix tiles included in the numbering).
|
|
'''
|
|
tile_elems = math.prod(tile_shape)
|
|
ts_t, ts_h, ts_w = tile_shape
|
|
prefix_sizes, prefix_rows, video_entries = [], [], []
|
|
cursor = 0
|
|
for seg in segments:
|
|
if isinstance(seg, tuple):
|
|
video_entries.append((cursor, seg))
|
|
cursor += math.prod(seg)
|
|
continue
|
|
if seg <= 0:
|
|
continue
|
|
full, rem = divmod(seg, tile_elems)
|
|
prefix_sizes.extend([tile_elems] * full)
|
|
if rem:
|
|
prefix_sizes.append(rem)
|
|
prefix_rows.append(torch.arange(cursor, cursor + seg, device=device, dtype=torch.long))
|
|
cursor += seg
|
|
num_prefix_tiles = len(prefix_sizes)
|
|
|
|
video_index_parts, video_size_parts, video_tile_spans = [], [], []
|
|
tile_cursor = num_prefix_tiles
|
|
for start, grid in video_entries:
|
|
t, h, w = grid
|
|
num_tiles = (math.ceil(t / ts_t), math.ceil(h / ts_h), math.ceil(w / ts_w))
|
|
region_sizes = _h3_variable_block_sizes(grid, num_tiles, device, tile_shape)
|
|
video_index_parts.append(_h3_tile_partition_indices(grid, tile_shape, device) + start)
|
|
video_size_parts.append(region_sizes)
|
|
video_tile_spans.append((tile_cursor, tile_cursor + int(region_sizes.numel())))
|
|
tile_cursor += int(region_sizes.numel())
|
|
num_video_tiles = tile_cursor - num_prefix_tiles
|
|
|
|
tile_partition_indices = torch.cat(prefix_rows + video_index_parts)
|
|
variable_block_sizes = torch.cat(
|
|
[torch.tensor(prefix_sizes, dtype=torch.long, device=device)] + video_size_parts)
|
|
non_pad_index = _h3_non_pad_index(variable_block_sizes, tile_elems)
|
|
untile_combined_index = non_pad_index[torch.argsort(tile_partition_indices)]
|
|
# One-time (lru-cached) synchronous bounds check; see _h3_validate_segment_geometry.
|
|
_h3_validate_segment_geometry(segments, variable_block_sizes, untile_combined_index, tile_elems)
|
|
return (tile_partition_indices, variable_block_sizes, untile_combined_index,
|
|
num_prefix_tiles, num_video_tiles, tuple(video_tile_spans), tile_elems)
|
|
|
|
|
|
# Per-process reusable padded tile scratch for forward_h3's no-grad (inference) path.
|
|
# Mirrors the official FastVideo builder-owned tile_buf_holder: one buffer per
|
|
# (role, shape, dtype, device) shared across all layers and steps, so the per-layer
|
|
# torch.zeros((B,H,L_buf,D)) allocation+zeroing churn is paid once instead of every call.
|
|
# Keyed by role because q/k/v/gate are all live simultaneously within one forward_h3.
|
|
_H3_TILE_BUFS = {}
|
|
|
|
|
|
def _h3_tile_into_buf(t, dtype, idx, shape, device, role, reuse):
|
|
R'''Scatter packed [B,L,H,D] ``t`` into a [B,H,L_buf,D] tile buffer (pad/dummy stay zero).
|
|
|
|
Under ``reuse`` (no-grad inference) the buffer is cached in ``_H3_TILE_BUFS`` and only
|
|
re-zeroed when the tile geometry (``idx`` identity) changes, exactly like the official
|
|
holder: pad slots are never written and non-pad slots are fully overwritten per call, so
|
|
a reused buffer stays valid. Callers consume the result before the next layer overwrites
|
|
it. Under grad the buffer is fresh -- an in-place write into a buffer that outlives the
|
|
step would anchor that step's autograd graph (the official tile-cache OOM).
|
|
'''
|
|
B, H, L_buf, D = shape
|
|
if not reuse:
|
|
buf = torch.zeros(shape, device=device, dtype=dtype)
|
|
buf[:, :, idx, :] = t.transpose(1, 2).to(dtype)
|
|
return buf
|
|
key = (role, B, H, L_buf, D, dtype, str(device))
|
|
ent = _H3_TILE_BUFS.get(key)
|
|
if ent is None or ent[0].shape != shape or ent[0].dtype != dtype:
|
|
ent = [torch.zeros(shape, device=device, dtype=dtype), None]
|
|
_H3_TILE_BUFS[key] = ent
|
|
if ent[1] is not idx: # geometry identity changed -> clear stale valid rows once
|
|
ent[0].zero_()
|
|
ent[1] = idx
|
|
ent[0][:, :, idx, :] = t.transpose(1, 2).to(dtype)
|
|
return ent[0]
|
|
|
|
|
|
def _h3_pool_tiles(x, variable_block_sizes, tile_elems, n_tiles):
|
|
R'''fp32 masked mean over each tile. ``x``: [B,H,L_buf,D] (pad slots zero) -> [B,H,n_tiles,D].
|
|
|
|
Pad slots are guaranteed zero (zeros-init buffer, only non-pad slots written), so a
|
|
plain fp32 sum divided by the true tile size is exactly the masked mean.
|
|
'''
|
|
B, H, _, D = x.shape
|
|
xt = x[:, :, :n_tiles * tile_elems].reshape(B, H, n_tiles, tile_elems, D)
|
|
pooled = xt.sum(dim=3, dtype=torch.float32)
|
|
return pooled / variable_block_sizes.view(1, 1, -1, 1).to(torch.float32)
|
|
|
|
|
|
def _h3_build_block_mask(scores, num_prefix_tiles, sparsity, exempt,
|
|
video_tile_spans, span_sparsities):
|
|
R'''scores [B,H,n_tiles,n_tiles] -> bool mask (port of official ``_build_block_mask``).
|
|
|
|
Prefix rows/cols: with ``exempt`` every query keeps every prefix column, and every
|
|
query keeps its own top-k of EACH video region under that region's sparsity.
|
|
``sparsity <= 0`` keeps every tile (dense).
|
|
'''
|
|
n_tiles = scores.shape[-1]
|
|
if sparsity <= 0.0:
|
|
span_sparsities = tuple(sparsity for _ in video_tile_spans)
|
|
span_topk = [_h3_compute_topk(s, end - start)
|
|
for (start, end), s in zip(video_tile_spans, span_sparsities)]
|
|
if all(kk == end - start for kk, (start, end) in zip(span_topk, video_tile_spans)):
|
|
return torch.ones_like(scores, dtype=torch.bool)
|
|
mask = torch.zeros_like(scores, dtype=torch.bool)
|
|
if exempt or num_prefix_tiles == 0:
|
|
for (start, end), k_span in zip(video_tile_spans, span_topk):
|
|
if k_span == end - start:
|
|
mask[..., start:end] = True
|
|
continue
|
|
idx = scores[..., start:end].topk(k_span, dim=-1).indices + start
|
|
mask.scatter_(-1, idx, True)
|
|
mask[..., :num_prefix_tiles] = True
|
|
else:
|
|
k_total = min(span_topk[0] + num_prefix_tiles, n_tiles)
|
|
idx = scores.topk(k_total, dim=-1).indices
|
|
mask.scatter_(-1, idx, True)
|
|
mask[:, :, :num_prefix_tiles, :] = True
|
|
return mask
|
|
|
|
|
|
def _h3_expand_mask_and_sizes_to_64(block_mask, variable_block_sizes, tile_elems):
|
|
R'''Logical tile selection -> the physical 64-block map the Triton kernels run at.
|
|
|
|
Port of the official ``_expand_mask_and_sizes_256_to_64`` / ``_128_to_64`` (their "route A"),
|
|
generalized over ``r = tile_elems // 64``. A logical tile occupies ``r`` CONTIGUOUS 64-blocks in
|
|
the padded buffer, so selecting it selects all r sub-blocks on both the query and the key axis
|
|
(``repeat_interleave``), and its valid-token count chops into 64-token strides. The trailing
|
|
sub-blocks of a partial logical tile legitimately get size 0, which every kernel treats as
|
|
"contributes nothing" (the forward masks those columns to -inf, dQ skips them, dK/dV never
|
|
stores them). ``tile_elems == 64`` is the identity.
|
|
'''
|
|
r = tile_elems // _H3_KERNEL_BLOCK
|
|
if r == 1:
|
|
return block_mask, variable_block_sizes
|
|
mask_64 = block_mask.repeat_interleave(r, dim=2).repeat_interleave(r, dim=3)
|
|
sizes = variable_block_sizes.to(torch.int32)
|
|
offsets = torch.arange(r, device=sizes.device, dtype=torch.int32) * _H3_KERNEL_BLOCK
|
|
sizes_64 = torch.clamp(sizes[:, None] - offsets[None, :],
|
|
min=0, max=_H3_KERNEL_BLOCK).reshape(-1)
|
|
return mask_64, sizes_64
|
|
|
|
|
|
class VSASparseAttention(nn.Module):
|
|
R'''
|
|
Video Sparse Attention (VSA) -- the FastVideo/hao-ai-lab coarse-to-fine block-sparse
|
|
attention (arXiv:2505.13389) used by the official FastWan2.1 sparse-distilled students.
|
|
|
|
Two branches, mirroring the released FastWan ``to_gate_compress`` design:
|
|
- sparse branch : mean-pool q/k into blocks -> coarse block scores -> keep the top-k
|
|
key blocks per query block -> block-sparse EXACT softmax over the
|
|
selected blocks.
|
|
- compressed branch: a global mean-pooled ("compressed") attention over ALL key blocks,
|
|
broadcast back to token resolution and scaled per token/head by the
|
|
learned gate ``to_gate_compress(x)`` (x = attention-input hidden
|
|
states). This is FastVideo's learned compensation for the context
|
|
lost to sparsification. ``to_gate_compress`` is zero-init, so a
|
|
checkpoint without those weights reduces to sparse-only::
|
|
|
|
out = o_sparse + to_gate_compress(x) * o_compressed
|
|
|
|
New learnable params per layer (all under ``attn_op.*``): ``to_gate_compress`` weight+bias.
|
|
Absent from a dense checkpoint (loaded with strict=False); the official FastWan release
|
|
provides them.
|
|
|
|
NOTE on fidelity: the block-sparse forward IS the official FastVideo Triton kernel (vendored),
|
|
the tile geometry / exempt mask / top-k are ports of the official H3 backend, and the tile
|
|
volume is selectable over the official {64, 128, 256} -- 128/256 reach the same 64-block kernel
|
|
through the official route-A expansion.
|
|
|
|
``forward`` accepts BSND (``[B, L, H, D]``) q/k/v like the sibling attention modules;
|
|
internal helpers use BHLD (``[B, H, L, *]``).
|
|
'''
|
|
def __init__(self, head_dim, num_heads, dim, topk=0.2,
|
|
BLKQ=64, BLKK=64, use_bf16=True, tile_elems=_H3_TILE_ELEMS):
|
|
R'''
|
|
Args:
|
|
head_dim: dimension of each head.
|
|
num_heads: number of attention heads.
|
|
dim: hidden size of the attention input (sizes ``to_gate_compress``).
|
|
topk: keep-ratio of key blocks for the sparse branch (0.2 -> 80% sparse,
|
|
matching the official ``--VSA-sparsity 0.8``).
|
|
BLKQ/BLKK: query/key block sizes for the flat (non-segment) sparse kernel.
|
|
use_bf16: compute the sparse kernel in bfloat16 (default) or float16.
|
|
tile_elems: logical tile volume for the segment-aware path -- one of the official
|
|
``VSA_H3_TILE_SHAPES`` keys (64 / 128 / 256). Must match the volume the
|
|
checkpoint was distilled at; the released FastWan and FastH3 both use 64.
|
|
'''
|
|
super().__init__()
|
|
self.head_dim = head_dim
|
|
self.num_heads = num_heads
|
|
self.dim = dim
|
|
self.topk = topk
|
|
self.BLKQ = BLKQ
|
|
self.BLKK = BLKK
|
|
self.dtype = torch.bfloat16 if use_bf16 else torch.float16
|
|
if tile_elems not in _H3_TILE_SHAPES:
|
|
raise ValueError(
|
|
f"Unsupported VSA tile volume {tile_elems}; supported volumes are "
|
|
f"{sorted(_H3_TILE_SHAPES)} (the official VSA_H3_TILE_SHAPES).")
|
|
self.tile_elems = tile_elems
|
|
self.tile_shape = _H3_TILE_SHAPES[tile_elems]
|
|
|
|
inner = num_heads * head_dim
|
|
# Learned gate for the compressed global branch (FastVideo ``to_gate_compress``).
|
|
self.to_gate_compress = nn.Linear(dim, inner, bias=True)
|
|
self.init_weights_()
|
|
|
|
def init_weights_(self):
|
|
with torch.no_grad():
|
|
nn.init.zeros_(self.to_gate_compress.weight)
|
|
nn.init.zeros_(self.to_gate_compress.bias)
|
|
|
|
def sparse(self, q, k, v):
|
|
R'''Block-sparse exact softmax (VSA fine stage).
|
|
|
|
Args:
|
|
q, k: [B, H, L, D] POST-RoPE queries/keys.
|
|
v: [B, H, L, D] values.
|
|
Returns:
|
|
[B, H, L, D], or ``None`` when the sequence is shorter than one block.
|
|
'''
|
|
L = q.shape[2]
|
|
if L < max(self.BLKQ, self.BLKK):
|
|
return None
|
|
orig_dtype = q.dtype
|
|
with torch.no_grad():
|
|
pq = self._pool_flat(q, self.BLKQ)
|
|
pk = self._pool_flat(k, self.BLKK)
|
|
scores = pq @ pk.transpose(-1, -2)
|
|
real_topk = max(1, min(scores.shape[-1], math.ceil(self.topk * scores.shape[-1])))
|
|
lut = scores.topk(real_topk, dim=-1, sorted=False).indices.contiguous()
|
|
sparse_map = torch.zeros_like(scores, dtype=torch.bool)
|
|
sparse_map.scatter_(-1, lut, True)
|
|
if self.BLKQ % 64 or self.BLKK % 64:
|
|
raise ValueError("VSA 的 flat block 大小必须是 64 的正整数倍。")
|
|
B, H, L, D = q.shape
|
|
n_tiles = math.ceil(L / 64)
|
|
padding = n_tiles * 64 - L
|
|
q, k, v = (F.pad(t.to(self.dtype), (0, 0, 0, padding)).contiguous() for t in (q, k, v))
|
|
sizes = (L - torch.arange(n_tiles, device=q.device) * 64).clamp(max=64).to(torch.int32)
|
|
mask = sparse_map.repeat_interleave(self.BLKQ // 64, dim=-2).repeat_interleave(
|
|
self.BLKK // 64, dim=-1)[..., :n_tiles, :n_tiles].contiguous()
|
|
o_s = sparse_linear_attention_varlen(q, k, v, mask, sizes)
|
|
return o_s[:, :, :L].to(orig_dtype)
|
|
|
|
@staticmethod
|
|
def _pool_flat(x, block_size):
|
|
B, H, L, D = x.shape
|
|
n_blocks = math.ceil(L / block_size)
|
|
padded = F.pad(x.float(), (0, 0, 0, n_blocks * block_size - L))
|
|
counts = (L - torch.arange(n_blocks, device=x.device) * block_size).clamp(max=block_size)
|
|
return padded.reshape(B, H, n_blocks, block_size, D).sum(3) / counts.view(1, 1, -1, 1)
|
|
|
|
def compressed(self, q, k, v):
|
|
R'''Mean-pooled global ("compressed") attention, broadcast to token resolution.
|
|
|
|
Args:
|
|
q, k, v: [B, H, L, D] POST-RoPE queries/keys/values.
|
|
Returns:
|
|
[B, H, L, D] compressed global context per token.
|
|
'''
|
|
B, H, L, D = q.shape
|
|
pq = self._pool_flat(q, self.BLKQ)
|
|
pk = self._pool_flat(k, self.BLKK)
|
|
pv = self._pool_flat(v, self.BLKK)
|
|
scale = D ** -0.5
|
|
scores = (pq.float() @ pk.float().transpose(-1, -2)) * scale
|
|
attn = torch.softmax(scores, dim=-1).to(pv.dtype)
|
|
oc_blocks = attn @ pv # [B,H,Lq,D]
|
|
# Broadcast each query block's context to its BLKQ tokens, then trim to L.
|
|
oc = oc_blocks.repeat_interleave(self.BLKQ, dim=2)[:, :, :L]
|
|
return oc
|
|
|
|
def forward_h3(self, q, k, v, x, segments, return_sparsity=False):
|
|
R'''Segment-aware H3 VSA -- the official FastVideo segment/tile mask semantics.
|
|
|
|
Args:
|
|
q, k, v: [B, L, H, D] POST-RoPE queries/keys/values in PACKED order.
|
|
x: [B, L, C] attention-input hidden states (drives the compress gate).
|
|
segments: packed layout -- dense row counts (``int``) and video region
|
|
``(t, h, w)`` token grid, in packed order. Built by the model forward.
|
|
Returns:
|
|
[B, L, H, D] attention output (before the output projection).
|
|
|
|
Tiles the packed sequence into [prefix tiles][video cubic tiles] of ``self.tile_shape``,
|
|
builds the exempt block mask on pooled tile scores, expands that logical selection down to
|
|
the 64-block granularity the kernels run at, runs the block-sparse Triton kernel for the
|
|
fine branch, adds the gated compressed branch on the same logical geometry, then untiles
|
|
back to packed order.
|
|
'''
|
|
B, L, H, D = q.shape
|
|
device = q.device
|
|
(tile_partition_indices, variable_block_sizes, untile_combined_index,
|
|
num_prefix_tiles, num_video_tiles, video_tile_spans, tile_elems) = \
|
|
build_h3_tile_geometry(tuple(segments), device, self.tile_shape)
|
|
n_tiles = num_prefix_tiles + num_video_tiles
|
|
n_regions = len(video_tile_spans)
|
|
# Logical tiles are laid end to end with no spare: the kernels take a per-64-block valid
|
|
# count, so a partial tile needs no zero dummy partner tile any more.
|
|
L_buf = n_tiles * tile_elems
|
|
sparsity = 1.0 - self.topk # self.topk is the KEEP ratio
|
|
vbs = variable_block_sizes.to(device)
|
|
|
|
reuse = not torch.is_grad_enabled()
|
|
buf_shape = (B, H, L_buf, D)
|
|
|
|
def _tile(t, dtype, role):
|
|
# [B,L,H,D] packed -> [B,H,L_buf,D] tile order (pad slots stay zero); under no-grad the
|
|
# buffer is shared across layers/steps (see _h3_tile_into_buf).
|
|
return _h3_tile_into_buf(t, dtype, untile_combined_index, buf_shape,
|
|
device, role, reuse)
|
|
|
|
qb = _tile(q, self.dtype, 'q')
|
|
kb = _tile(k, self.dtype, 'k')
|
|
vb = _tile(v, self.dtype, 'v')
|
|
gate = self.to_gate_compress(x).view(B, L, H, D)
|
|
gb = _tile(gate, gate.dtype, 'g')
|
|
|
|
pq = _h3_pool_tiles(qb, vbs, tile_elems, n_tiles) # [B,H,n_tiles,D] fp32
|
|
pk = _h3_pool_tiles(kb, vbs, tile_elems, n_tiles)
|
|
pv = _h3_pool_tiles(vb, vbs, tile_elems, n_tiles)
|
|
|
|
# Block selection is a non-differentiable constant (exactly like the flat path's
|
|
# get_block_map): build the mask under no_grad on detached pooled scores so no graph is
|
|
# retained for the top-k.
|
|
with torch.no_grad():
|
|
scores = torch.matmul(pq.detach(), pk.detach().transpose(-2, -1)) / (D ** 0.5)
|
|
mask = _h3_build_block_mask(scores, num_prefix_tiles, sparsity, True,
|
|
video_tile_spans, (sparsity,) * n_regions)
|
|
# Logical (tile_elems) selection -> the physical 64-block map + per-block counts.
|
|
# Identity when tile_elems == 64.
|
|
block_mask, block_sizes = _h3_expand_mask_and_sizes_to_64(mask, vbs, tile_elems)
|
|
o_s = sparse_linear_attention_varlen(
|
|
qb, kb, vb, block_mask, block_sizes, qk_scale=D ** -0.5)
|
|
|
|
# gated compressed branch on the same logical tile geometry
|
|
cs = torch.matmul(pq, pk.transpose(-2, -1)) * (D ** -0.5)
|
|
attn = torch.softmax(cs, dim=-1).to(pv.dtype)
|
|
oc = (attn @ pv).repeat_interleave(tile_elems, dim=2) # [B,H,L_buf,D]
|
|
out_buf = o_s.to(gb.dtype) + gb * oc.to(gb.dtype)
|
|
|
|
out = out_buf[:, :, untile_combined_index].transpose(1, 2) # [B,L,H,D] packed
|
|
out = out.to(q.dtype)
|
|
if return_sparsity:
|
|
return out, self.topk
|
|
return out
|
|
|
|
def forward(self, q, k, v, x, return_sparsity=False, segments=None):
|
|
R'''
|
|
Args:
|
|
q, k, v: [B, L, H, D] POST-RoPE queries/keys/values.
|
|
x: [B, L, C] attention-input hidden states (drives the compress gate).
|
|
return_sparsity: also return the sparse-branch keep ratio.
|
|
Returns:
|
|
[B, L, H, D] attention output (before the output projection).
|
|
'''
|
|
if segments is not None:
|
|
# Segment-aware tile geometry with a differentiable variable-length block-sparse
|
|
# kernel (forward + backward), so training and inference share one selection
|
|
# semantics. Falls back to the flat 1-D block map only when the caller cannot
|
|
# supply a clean segment layout (segments is None).
|
|
return self.forward_h3(q, k, v, x, segments, return_sparsity=return_sparsity)
|
|
B, L, H, D = q.shape
|
|
qb = q.transpose(1, 2).contiguous() # [B,H,L,D]
|
|
kb = k.transpose(1, 2).contiguous()
|
|
vb = v.transpose(1, 2).contiguous()
|
|
|
|
o_s = self.sparse(qb, kb, vb)
|
|
if o_s is None:
|
|
# Shorter than one block: exact dense attention (sparsity gives no benefit).
|
|
o_s = F.scaled_dot_product_attention(
|
|
qb.float(), kb.float(), vb.float()).to(qb.dtype)
|
|
gate = self.to_gate_compress(x).view(B, L, H, D).transpose(1, 2) # [B,H,L,D]
|
|
o_c = self.compressed(qb, kb, vb)
|
|
out = o_s + gate.to(o_s.dtype) * o_c.to(o_s.dtype)
|
|
out = out.transpose(1, 2) # [B,L,H,D]
|
|
if return_sparsity:
|
|
return out, self.topk
|
|
return out |