Blame for DecMem/wan/modules/ltm_processor.py

"""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)