"""Long-Term Memory (LTM) sparse attention processor. Ported from ar_mem/framear/src/models/attention_processor.py with the STM (sliding window dense) path removed: in Self-Forcing, STM is computed with flex_attention outside this module, and only the LTM block-sparse output (``out_v``) flows through ``LTMAttentionProcessor``. The H100 SM90 block-sparse kernel is required at import time; non-H100 GPUs will raise. """ import functools import math from dataclasses import dataclass from typing import List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F from einops import rearrange major, minor = torch.cuda.get_device_capability(0) if major == 9 and minor == 0: from vsa.block_sparse_wrapper import block_sparse_attn_SM90 block_sparse_attn = block_sparse_attn_SM90 else: raise RuntimeError( f"LTMAttentionProcessor requires H100 (SM 9.0); got SM {major}.{minor}" ) @dataclass class LTMAttentionMetadata: dit_seq_shape: List[int] num_tiles: List[int] total_seq_length: int tile_partition_indices: torch.LongTensor reverse_tile_partition_indices: torch.LongTensor variable_block_sizes: torch.LongTensor non_pad_index: torch.LongTensor ltm_topk: int @functools.lru_cache(maxsize=10) def get_tile_partition_indices( dit_seq_shape: Tuple[int, int, int], tile_size: Tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: T, H, W = dit_seq_shape ts, hs, ws = tile_size indices = torch.arange(T * H * W, device=device, dtype=torch.long).reshape(T, H, W) ls = [] for t in range(math.ceil(T / ts)): for h in range(math.ceil(H / hs)): for w in range(math.ceil(W / ws)): ls.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(ls, dim=0) @functools.lru_cache(maxsize=10) def get_reverse_tile_partition_indices( dit_seq_shape: Tuple[int, int, int], tile_size: Tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: return torch.argsort(get_tile_partition_indices(dit_seq_shape, tile_size, device)) @functools.lru_cache(maxsize=10) def construct_variable_block_sizes( dit_seq_shape: Tuple[int, int, int], num_tiles: Tuple[int, int, int], ltm_tile_size: Tuple[int, int, int], device: torch.device, ) -> torch.LongTensor: t, h, w = dit_seq_shape ts_t, ts_h, ts_w = ltm_tile_size n_t, n_h, n_w = num_tiles def _sizes(dim_len: int, tile: int, n_tiles: int) -> torch.LongTensor: sizes = torch.full((n_tiles,), tile, dtype=torch.int, device=device) remainder = dim_len - (n_tiles - 1) * tile sizes[-1] = remainder if remainder > 0 else tile return sizes t_sizes = _sizes(t, ts_t, n_t) h_sizes = _sizes(h, ts_h, n_h) w_sizes = _sizes(w, ts_w, n_w) block_sizes = ( t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :] ).reshape(-1) return block_sizes @functools.lru_cache(maxsize=10) def get_non_pad_index(variable_block_sizes: torch.LongTensor, max_block_size: int): n_win = variable_block_sizes.shape[0] device = variable_block_sizes.device starts_pad = torch.arange(n_win, device=device) * max_block_size index_pad = starts_pad[:, None] + torch.arange(max_block_size, device=device)[None, :] index_mask = ( torch.arange(max_block_size, device=device)[None, :] < variable_block_sizes[:, None] ) return index_pad[index_mask] @functools.lru_cache(maxsize=10) def build_ltm_metadata( dit_seq_shape: Tuple[int, int, int], ltm_topk: int, ltm_tile_size: Tuple[int, int, int], device: torch.device, is_tf: bool = False ) -> LTMAttentionMetadata: num_tiles = ( math.ceil(dit_seq_shape[0] / ltm_tile_size[0]), math.ceil(dit_seq_shape[1] / ltm_tile_size[1]), math.ceil(dit_seq_shape[2] / ltm_tile_size[2]), ) total_seq_length = math.prod(dit_seq_shape) tile_partition_indices = get_tile_partition_indices(dit_seq_shape, ltm_tile_size, device) reverse_tile_partition_indices = get_reverse_tile_partition_indices( dit_seq_shape, ltm_tile_size, device ) variable_block_sizes = construct_variable_block_sizes( dit_seq_shape, num_tiles, ltm_tile_size, device ) non_pad_index = get_non_pad_index(variable_block_sizes, math.prod(ltm_tile_size)) if is_tf: # offsets are DIFFERENT for the input-space vs output-space index tables: # - tile_partition_indices / reverse_tile_partition_indices live in the # unpadded INPUT space, whose second half starts at total_seq_length. # - non_pad_index lives in the padded OUTPUT space (after tile()'s # internal zero pad), whose second half starts at padded_total_seq_len. single_side_unpadded_len = total_seq_length single_side_padded_len = math.prod(num_tiles) * math.prod(ltm_tile_size) num_tiles = (num_tiles[0] * 2, num_tiles[1], num_tiles[2]) total_seq_length = total_seq_length * 2 tile_partition_indices = torch.cat( [tile_partition_indices, tile_partition_indices + single_side_unpadded_len], dim=0 ) reverse_tile_partition_indices = torch.cat( [reverse_tile_partition_indices, reverse_tile_partition_indices + single_side_unpadded_len], dim=0 ) variable_block_sizes = torch.cat([variable_block_sizes, variable_block_sizes], dim=0) non_pad_index = torch.cat( [non_pad_index, non_pad_index + single_side_padded_len], dim=0 ) assert ltm_topk != -1, "ltm_topk must be set" return LTMAttentionMetadata( dit_seq_shape=dit_seq_shape, num_tiles=num_tiles, total_seq_length=total_seq_length, tile_partition_indices=tile_partition_indices, reverse_tile_partition_indices=reverse_tile_partition_indices, variable_block_sizes=variable_block_sizes, non_pad_index=non_pad_index, ltm_topk=ltm_topk, ) def generate_clean_noise_mask(clean_mask: torch.Tensor, num_frames: int) -> torch.Tensor: """Expand a (1,1,l,l) clean-only block mask to a (1,1,2l,2l) interleaved clean/noisy mask matching ar_mem's teacher-forcing convention. """ clean_mask_extend = torch.cat( [clean_mask, torch.zeros_like(clean_mask)], dim=-1 ) _, _, seq_len = clean_mask.shape token_per_frame = seq_len // num_frames idx = torch.arange(seq_len, device=clean_mask.device) row_idx = idx.unsqueeze(1) col_idx = idx.unsqueeze(0) lower = (row_idx // token_per_frame > col_idx // token_per_frame) diag = (row_idx // token_per_frame == col_idx // token_per_frame) noise_mask_right = torch.where(diag, clean_mask, torch.zeros_like(clean_mask)) noise_mask_left = torch.where(lower, clean_mask, torch.zeros_like(clean_mask)) noise_mask_extend = torch.cat([noise_mask_left, noise_mask_right], dim=-1) return torch.cat([clean_mask_extend, noise_mask_extend], dim=1) @functools.lru_cache(maxsize=10) def prepare_block_causal_mask( q_block: int, kv_block: int, num_tiles_per_block: int, enable_clean_noise_mask: bool, device_str: str, dtype: torch.dtype = torch.bfloat16, ) -> torch.Tensor: """LTM coarse-block mask. Allowed (=0) positions: ego block (same frame) or kv frame at least ``sliding_window`` frames earlier than q frame. Everything else is -10000. Output shape: (1, 1, L, L) where L is the number of coarse blocks. """ device = torch.device(device_str) q_frames = q_block // num_tiles_per_block frame_indices = torch.arange(q_frames, device=device).repeat_interleave(num_tiles_per_block) diff = frame_indices.unsqueeze(1) - frame_indices.unsqueeze(0) block_mask = (diff >= 0).float() if enable_clean_noise_mask: block_mask = generate_clean_noise_mask(block_mask.unsqueeze(0), q_frames) block_mask = (1 - block_mask) * -10000.0 return block_mask.unsqueeze(0).to(dtype) # @torch.compile def compress_block(x: torch.Tensor, block_size: int, variable_sizes: torch.Tensor) -> torch.Tensor: B, H, L, D = x.shape x_reshaped = x.view(B, H, L // block_size, block_size, D) return (x_reshaped.float().sum(dim=3) / variable_sizes.view(1, 1, -1, 1)).to(x.dtype) def _coarse_attention(q, k, v, attn_mask) -> Tuple[torch.Tensor, torch.Tensor]: QK = torch.matmul(q, k.transpose(-2, -1)) / (q.size(-1) ** 0.5) if attn_mask is not None: QK = QK + attn_mask QK = torch.nn.functional.softmax(QK, dim=-1) return torch.matmul(QK, v), QK class LTMAttentionProcessor(nn.Module): """LTM sparse attention processor. Forward returns the sparse (long-term) attention output ``out_v`` shaped ``(B, H, L, D)`` in the original (untiled) spatial order. The caller is responsible for computing the STM/dense path separately and combining the two (e.g., ``out_s + mem_proj(x) * out_v``). """ def __init__(self, ltm_tile_size: Tuple[int, int, int] = (1, 8, 8)): super().__init__() self.ltm_tile_size = ltm_tile_size def tile(self, x, num_tiles, tile_partition_indices, non_pad_index): t_pad = num_tiles[0] * self.ltm_tile_size[0] h_pad = num_tiles[1] * self.ltm_tile_size[1] w_pad = num_tiles[2] * self.ltm_tile_size[2] x_padded = torch.zeros( (x.shape[0], x.shape[1], t_pad * h_pad * w_pad, x.shape[-1]), device=x.device, dtype=x.dtype, ) x_padded[:, :, non_pad_index] = x[:, :, tile_partition_indices] return x_padded def untile(self, x, reverse_tile_partition_indices, non_pad_index): return x[:, :, non_pad_index][:, :, reverse_tile_partition_indices] def preprocess_qkv(self, qkv, ltm_metadata): return self.tile( qkv, ltm_metadata.num_tiles, ltm_metadata.tile_partition_indices, ltm_metadata.non_pad_index, ) def postprocess_output(self, output, ltm_metadata): return self.untile( output, ltm_metadata.reverse_tile_partition_indices, ltm_metadata.non_pad_index, ) def _merge_current_kv_into_cache( self, key: torch.Tensor, value: torch.Tensor, ltm_metadata: LTMAttentionMetadata, kv_cache: Optional[dict], write_cache: bool, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Tile/compress the current chunk only, then concat with cached KV.""" block_elements = math.prod(self.ltm_tile_size) variable_block_sizes_cur = ltm_metadata.variable_block_sizes k_cur = self.tile( key, ltm_metadata.num_tiles, ltm_metadata.tile_partition_indices, ltm_metadata.non_pad_index, ) v_cur = self.tile( value, ltm_metadata.num_tiles, ltm_metadata.tile_partition_indices, ltm_metadata.non_pad_index, ) k_comp_cur = compress_block(k_cur, block_elements, variable_block_sizes_cur) v_comp_cur = compress_block(v_cur, block_elements, variable_block_sizes_cur) if kv_cache is not None and kv_cache.get("k_tiled") is not None: k_tiled = torch.cat([kv_cache["k_tiled"], k_cur], dim=2) v_tiled = torch.cat([kv_cache["v_tiled"], v_cur], dim=2) k_compress = torch.cat([kv_cache["k_compress"], k_comp_cur], dim=2) v_compress = torch.cat([kv_cache["v_compress"], v_comp_cur], dim=2) variable_block_sizes_kv = torch.cat( [kv_cache["variable_block_sizes"], variable_block_sizes_cur], dim=0 ) else: k_tiled, v_tiled = k_cur, v_cur k_compress, v_compress = k_comp_cur, v_comp_cur variable_block_sizes_kv = variable_block_sizes_cur if write_cache and kv_cache is not None: kv_cache["k_tiled"] = k_tiled kv_cache["v_tiled"] = v_tiled kv_cache["k_compress"] = k_compress kv_cache["v_compress"] = v_compress kv_cache["variable_block_sizes"] = variable_block_sizes_kv return k_tiled, v_tiled, k_compress, v_compress, variable_block_sizes_kv def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, ltm_metadata: LTMAttentionMetadata, enable_clean_noise_mask: bool, num_frame_per_block: int, kv_cache: Optional[dict] = None, write_cache: bool = True, apply_block_causal_mask: bool = True, ) -> torch.Tensor: """Args ---- query, key, value : (B, H, L, D) with L == prod(dit_seq_shape). In teacher-forcing mode they should arrive batch-concatenated as ``(2B, H, L, D)`` (clean half then noisy half); when ``enable_clean_noise_mask=True`` they are internally reshuffled into ``(B, H, 2L, D)``. key, value : at inference, only the **current** block (not full history). History is read from ``kv_cache`` (tiled + compressed tokens). kv_cache : optional per-layer cache with keys ``k_tiled``, ``v_tiled``, ``k_compress``, ``v_compress``, ``variable_block_sizes``. Updated when ``write_cache=True``. apply_block_causal_mask : if False, coarse top-k uses no block mask (inference; causality is already enforced by the KV cache schedule). Returns ------- out_v : same shape as ``query`` (after the optional TF unshuffle), in the original spatial order (un-tiled). """ block_elements = math.prod(self.ltm_tile_size) # Spatial tiles per video frame, INDEPENDENT of TF doubling (num_tiles[0] # may have been doubled in build_ltm_metadata(is_tf=True), but [1:] is not). tiles_per_frame = math.prod(ltm_metadata.num_tiles[1:]) num_tiles_per_block = tiles_per_frame * num_frame_per_block # 1) Tile Q (current block only) q_tiled = self.tile( query, ltm_metadata.num_tiles, ltm_metadata.tile_partition_indices, ltm_metadata.non_pad_index, ) # 2) Tile/compress current K/V and merge with cached history k_tiled, v_tiled, k_compress, v_compress, variable_block_sizes_kv = ( self._merge_current_kv_into_cache( key, value, ltm_metadata, kv_cache, write_cache ) ) B, H, _, D = q_tiled.shape q_blocks_total = q_tiled.shape[2] // block_elements k_blocks_total = k_tiled.shape[2] // block_elements if enable_clean_noise_mask: assert q_blocks_total % 2 == 0 and k_blocks_total % 2 == 0 q_blocks = q_blocks_total // 2 k_blocks = k_blocks_total // 2 else: q_blocks = q_blocks_total k_blocks = k_blocks_total variable_block_sizes_q = ltm_metadata.variable_block_sizes # 3) Compress Q blocks (K/V compressed cache built incrementally above) q_compress = compress_block(q_tiled, block_elements, variable_block_sizes_q) # 4) Coarse block attention (+ optional LTM block-causal mask for training) if apply_block_causal_mask: block_causal_mask = prepare_block_causal_mask( q_block=q_blocks, kv_block=k_blocks, num_tiles_per_block=num_tiles_per_block, enable_clean_noise_mask=enable_clean_noise_mask, device_str=str(query.device), dtype=query.dtype, ) else: block_causal_mask = None _, block_attn_score = _coarse_attention(q_compress, k_compress, v_compress, block_causal_mask) # 5) Top-k block selection + causality fix-up topk = min(ltm_metadata.ltm_topk, block_attn_score.shape[-1]) if block_causal_mask is not None: scored = block_attn_score + block_causal_mask topk_indices = torch.topk(scored, topk, dim=-1).indices block_mask = torch.zeros_like(scored, dtype=torch.bool).scatter_(-1, topk_indices, True) block_mask &= (block_causal_mask >= -1e-3) else: topk_indices = torch.topk(block_attn_score, topk, dim=-1).indices block_mask = torch.zeros_like(block_attn_score, dtype=torch.bool).scatter_( -1, topk_indices, True ) # 6) Block-sparse fine attention (H100 SM90 kernel) if not apply_block_causal_mask and q_tiled.shape[0] == 2: output_select = torch.empty_like(q_tiled, dtype=q_tiled.dtype, device=q_tiled.device) for i in range(2): # batch size must be 2 output_select[i:i+1] = block_sparse_attn(q_tiled[i:i+1], k_tiled[i:i+1], v_tiled[i:i+1], block_mask[i:i+1], \ variable_block_sizes_kv, variable_block_sizes_q)[0] else: output_select, _ = block_sparse_attn( q_tiled, k_tiled, v_tiled, block_mask, variable_block_sizes_kv, variable_block_sizes_q, ) return self.postprocess_output(output_select, ltm_metadata)