import torch
import torch.nn as nn
from typing import List, Tuple
import math
class LayerWiseBlockPartitioner:
"""
Cycles through 1D (temporal), 2D (spatial), and 3D (spatio-temporal)
block partitioning schemes across diffusion layers.
"""
def __init__(self, num_layers=24, frame_count=16, spatial_size=32):
self.num_layers = num_layers
self.frame_count = frame_count
self.spatial_h = spatial_size
self.spatial_w = spatial_size
self.partition_cycle = ['1d', '2d', '3d']
def get_partition_type(self, layer_idx: int) -> str:
"""Determine partition type for layer (cycles 1D→2D→3D)."""
cycle_position = layer_idx % 3
return self.partition_cycle[cycle_position]
def partition_keys_1d(self, keys: torch.Tensor) -> List[torch.Tensor]:
"""
Temporal partitioning: group keys by frame.
Args:
keys: (batch, seq_len, dim) where seq_len = frames * height * width
Returns:
blocks: List of (batch, block_size, dim) tensors, one per frame
"""
seq_len = keys.shape[1]
spatial_tokens_per_frame = seq_len // self.frame_count
blocks = []
for frame_idx in range(self.frame_count):
start = frame_idx * spatial_tokens_per_frame
end = start + spatial_tokens_per_frame
frame_block = keys[:, start:end, :]
blocks.append(frame_block)
return blocks
def partition_keys_2d(self, keys: torch.Tensor) -> List[torch.Tensor]:
"""
Spatial partitioning: group keys by spatial regions within each frame.
"""
batch_size, seq_len, dim = keys.shape
spatial_tokens_per_frame = seq_len // self.frame_count
blocks = []
patch_size = 4
for frame_idx in range(self.frame_count):
frame_start = frame_idx * spatial_tokens_per_frame
frame_keys = keys[:, frame_start:frame_start + spatial_tokens_per_frame, :]
spatial_grid = frame_keys.view(
batch_size, self.spatial_h, self.spatial_w, dim
)
for patch_y in range(0, self.spatial_h, patch_size):
for patch_x in range(0, self.spatial_w, patch_size):
patch = spatial_grid[
:, patch_y:patch_y+patch_size, patch_x:patch_x+patch_size, :
]
patch_flat = patch.reshape(batch_size, -1, dim)
blocks.append(patch_flat)
return blocks
def partition_keys_3d(self, keys: torch.Tensor) -> List[torch.Tensor]:
"""
Spatio-temporal partitioning: 3D volumes combining frames and spatial.
"""
batch_size, seq_len, dim = keys.shape
spatial_tokens_per_frame = seq_len // self.frame_count
blocks = []
frame_block_size = 4
spatial_block_size = 4
for frame_block_idx in range(0, self.frame_count, frame_block_size):
frames_in_block = min(
frame_block_size, self.frame_count - frame_block_idx
)
frame_end = frame_block_idx + frames_in_block
temporal_keys = keys[
:, frame_block_idx*spatial_tokens_per_frame:frame_end*spatial_tokens_per_frame, :
]
for patch_y in range(0, self.spatial_h, spatial_block_size):
for patch_x in range(0, self.spatial_w, spatial_block_size):
block_3d = temporal_keys
blocks.append(block_3d)
return blocks
class GlobalBlockSelector:
"""
Selects important blocks based on aggregated query-key similarities
across all queries, reducing per-query overhead.
"""
def __init__(self, similarity_threshold=0.1):
self.threshold = similarity_threshold
def select_blocks(
self,
queries: torch.Tensor,
blocks: List[torch.Tensor],
num_heads: int = 12
) -> Tuple[List[torch.Tensor], torch.Tensor]:
"""
Select blocks by aggregating similarities across queries.
Returns:
selected_blocks: List of selected block tensors
selection_mask: (batch, num_blocks) binary mask
"""
batch_size = queries.shape[0]
num_blocks = len(blocks)
dim = queries.shape[-1]
head_dim = dim // num_heads
similarities = torch.zeros(batch_size, num_blocks)
for block_idx, block in enumerate(blocks):
block_sim = torch.einsum('bqd,bkd->bq', queries, block) / math.sqrt(head_dim)
similarities[:, block_idx] = block_sim.max(dim=1)[0]
selection_mask = (similarities > self.threshold).float()
min_blocks = max(1, int(num_blocks * 0.2))
for b in range(batch_size):
num_selected = selection_mask[b].sum()
if num_selected < min_blocks:
top_k = torch.topk(similarities[b], min_blocks)[1]
selection_mask[b, top_k] = 1.0
selected_blocks = [
block for block_idx, block in enumerate(blocks)
if selection_mask[0, block_idx] > 0
]
return selected_blocks, selection_mask
class VMoBAAttentionLayer(nn.Module):
"""
Sparse attention layer combining cyclic block partitioning
and global block selection for video diffusion.
"""
def __init__(self, dim=768, num_heads=12, layer_idx=0, num_layers=24):
super().__init__()
self.dim = dim
self.num_heads = num_heads
self.layer_idx = layer_idx
self.partitioner = LayerWiseBlockPartitioner(
num_layers=num_layers, frame_count=16, spatial_size=32
)
self.selector = GlobalBlockSelector(similarity_threshold=0.1)
self.q_proj = nn.Linear(dim, dim)
self.k_proj = nn.Linear(dim, dim)
self.v_proj = nn.Linear(dim, dim)
self.out_proj = nn.Linear(dim, dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Apply VMoBA sparse attention.
Args:
x: (batch, seq_len, dim)
Returns:
attended: (batch, seq_len, dim)
"""
batch_size, seq_len, dim = x.shape
queries = self.q_proj(x)
keys = self.k_proj(x)
values = self.v_proj(x)
queries = queries.view(batch_size, seq_len, self.num_heads, dim // self.num_heads)
queries = queries.transpose(1, 2)
partition_type = self.partitioner.get_partition_type(self.layer_idx)
if partition_type == '1d':
blocks = self.partitioner.partition_keys_1d(keys)
elif partition_type == '2d':
blocks = self.partitioner.partition_keys_2d(keys)
else:
blocks = self.partitioner.partition_keys_3d(keys)
selected_blocks, _ = self.selector.select_blocks(queries, blocks)
attended = self._sparse_attention(
queries, selected_blocks, values, self.num_heads
)
attended = attended.transpose(1, 2).contiguous()
attended = attended.view(batch_size, seq_len, dim)
output = self.out_proj(attended)
return output
def _sparse_attention(self, queries, selected_blocks, values, num_heads):
"""Compute attention only over selected sparse blocks."""
selected_keys = torch.cat(selected_blocks, dim=1)
attn_weights = torch.softmax(
torch.bmm(queries, selected_keys.transpose(1, 2)) / math.sqrt(self.dim),
dim=-1
)
attended = torch.bmm(attn_weights, selected_keys)
return attended
