"""VSA-H3 — MiniMax-H3's Video Sparse Attention, as a `diffusers` attention processor. `FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` is distilled **under** block-sparse attention (`attention_backend: VIDEO_SPARSE_ATTN_H3`, `vsa_tile_size: 64`, `vsa_sparsity: 0.9` in the checkpoint's own `fastvideo_inference.json`), and FastVideo's release notes are explicit that dense attention is *not* a drop-in substitute for a VSA-trained student: the student learned to attend to the top-10% tiles its selector picks, and it also ships 50 trained `attn.to_gate_compress` tensors that only the sparse path reads. So this Space runs the sparse path. `diffusers` has no VSA backend, so this module ports FastVideo's `MiniMaxH3VSABackend` (`fastvideo/attention/backends/video_sparse_attn_h3.py`) onto `MiniMaxH3Attention`: 1. **Tiling.** The packed sequence `[text | audio | video]` is cut into 64-token tiles: segment-pure prefix chunks first, then `(4, 4, 4)` cubes of the post-patchify `(t, h, w)` video grid. Tiles are zero-padded to 64 and `variable_block_sizes` carries each tile's true occupancy. 2. **Selection.** Per head, tiles are mean-pooled in fp32, `scores = q_pooled @ k_pooledᵀ / √d`, and each query tile keeps the top `ceil((1 - 0.9) * num_video_tiles)` video tiles. Prefix (text/audio) keys are *exempt* — always selected — and prefix queries are always dense, which is FastVideo's default `vsa_mode`. 3. **Kernel.** The resulting bool block map is compacted with FastVideo's `map_to_index` Triton kernel and consumed by its `triton_block_sparse_attn_forward`, both vendored verbatim under `vsa_kernel/`. This is FastVideo's own `--vsa-kernel triton` route; the checkpoint's `vsa_kernel: sm100a` is the GB200-only fast path for the *same* mask semantics, and this Space's Blackwell RTX PRO 6000 is sm120. 4. **Compression branch.** `out_c = softmax(scores) @ v_pooled` broadcast back over each tile's rows and scaled by the trained per-row `to_gate_compress` gate, added to the sparse output. The base MiniMax-H3 release zero-initializes this gate (branch inert); this student ships it trained. The one deliberate deviation from the reference is memory layout, not math: buffers are tiled straight into the kernel's `[B, H, S_pad, D]` layout and reused across the 50 blocks, and the gate is applied *after* untiling, so the sparse path costs one persistent tile buffer per projection instead of four plus three transposed copies. """ from __future__ import annotations import functools import math import os from dataclasses import dataclass import torch # 64-token tiles: `(4, 4, 4)` over the post-patchify video grid. This is the checkpoint's `vsa_tile_size`, and it is # the Triton kernels' native block size, so the block map needs no expansion. TILE_ELEMS = 64 TILE_SHAPE = (4, 4, 4) DEFAULT_SPARSITY = 0.9 # -------------------------------------------------------------------------------------------------------------- # Geometry — ported from `fastvideo.attention.backends.video_sparse_attn{,_h3}` # -------------------------------------------------------------------------------------------------------------- def _tile_partition_indices(dit_seq_shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: """Row indices of the video grid in `(4, 4, 4)` tile order.""" grid_t, grid_h, grid_w = dit_seq_shape ts, hs, ws = TILE_SHAPE indices = torch.arange(grid_t * grid_h * grid_w, device=device, dtype=torch.long).reshape(grid_t, grid_h, grid_w) chunks = [] for t in range(math.ceil(grid_t / ts)): for h in range(math.ceil(grid_h / hs)): for w in range(math.ceil(grid_w / ws)): chunks.append( indices[t * ts : min(t * ts + ts, grid_t), h * hs : min(h * hs + hs, grid_h), w * ws : min(w * ws + ws, grid_w)].flatten() ) return torch.cat(chunks, dim=0) def _video_block_sizes(dit_seq_shape: tuple[int, int, int], device: torch.device) -> torch.Tensor: """Valid (non-padded) token count of every video tile, in the same tile order.""" ts, hs, ws = TILE_SHAPE def sizes(length: int, tile: int) -> torch.Tensor: count = math.ceil(length / tile) out = torch.full((count,), tile, dtype=torch.long, device=device) remainder = length - (count - 1) * tile out[-1] = remainder if remainder > 0 else tile return out t_sizes, h_sizes, w_sizes = (sizes(length, tile) for length, tile in zip(dit_seq_shape, TILE_SHAPE)) return (t_sizes[:, None, None] * h_sizes[None, :, None] * w_sizes[None, None, :]).reshape(-1) def _non_pad_index(variable_block_sizes: torch.Tensor) -> torch.Tensor: """Padded-buffer slot of every real token, in tile order.""" device = variable_block_sizes.device starts = torch.arange(variable_block_sizes.shape[0], device=device) * TILE_ELEMS slots = starts[:, None] + torch.arange(TILE_ELEMS, device=device)[None, :] keep = torch.arange(TILE_ELEMS, device=device)[None, :] < variable_block_sizes[:, None] return slots[keep] @dataclass(frozen=True, eq=False) class VSAGeometry: """Everything the sparse path needs about one packed layout. Cached per `(prefix_segments, grid)`.""" seq_len: int padded_len: int n_tiles: int num_prefix_tiles: int num_video_tiles: int topk: int # int32 tile occupancies, as the Triton kernels take them. variable_block_sizes: torch.Tensor # fp32 tile occupancies, for the pooled mean. tile_divisor: torch.Tensor # packed row -> padded tile-buffer slot (scatter on the way in, gather on the way out) untile_index: torch.Tensor # packed row -> its tile, so the compression branch can be applied after untiling row_tile_index: torch.Tensor @property def dense(self) -> bool: return self.topk >= self.num_video_tiles @functools.lru_cache(maxsize=8) def build_geometry( prefix_segments: tuple[int, ...], dit_seq_shape: tuple[int, int, int], device: torch.device, sparsity: float = DEFAULT_SPARSITY, ) -> VSAGeometry: prefix_len = sum(prefix_segments) prefix_sizes: list[int] = [] for segment in prefix_segments: full, remainder = divmod(segment, TILE_ELEMS) prefix_sizes.extend([TILE_ELEMS] * full) if remainder: prefix_sizes.append(remainder) video_sizes = _video_block_sizes(dit_seq_shape, device) variable_block_sizes = torch.cat( [torch.tensor(prefix_sizes, dtype=torch.long, device=device), video_sizes] ) partition = torch.cat( [ torch.arange(prefix_len, device=device, dtype=torch.long), _tile_partition_indices(dit_seq_shape, device) + prefix_len, ] ) untile_index = _non_pad_index(variable_block_sizes)[torch.argsort(partition)] n_tiles = int(variable_block_sizes.numel()) num_video_tiles = int(video_sizes.numel()) return VSAGeometry( seq_len=int(partition.numel()), padded_len=n_tiles * TILE_ELEMS, n_tiles=n_tiles, num_prefix_tiles=len(prefix_sizes), num_video_tiles=num_video_tiles, # FastVideo's `compute_topk`, clamped to [1, num_video_tiles]. topk=max(1, min(math.ceil((1.0 - sparsity) * num_video_tiles), num_video_tiles)), variable_block_sizes=variable_block_sizes.to(torch.int32).contiguous(), tile_divisor=variable_block_sizes.to(torch.float32).view(1, 1, -1, 1), untile_index=untile_index, row_tile_index=untile_index // TILE_ELEMS, ) # Request-scoped geometry reuse. The layout of a packed sequence — the prefix segment run-lengths and the video # grid — is identical across all forwards of one request (same canvas, same frames, same token counts), and # `geometry_from_layout` normally recovers it from `token_tags` / `position_ids` with host reads (`tolist`, several # `int(...)`) that each drain the GPU queue. The app calls `begin_request()` once per request; within a request the # first forward derives the geometry and the rest reuse it with zero syncs. _REQUEST_STATE: dict = {"nonce": 0, "geometry": None} _REQUEST_COUNTER = 0 def begin_request() -> None: """Open a new packed-sequence request: drop the reused geometry so the next forward re-derives it.""" global _REQUEST_COUNTER _REQUEST_COUNTER += 1 _REQUEST_STATE["nonce"] = _REQUEST_COUNTER _REQUEST_STATE["geometry"] = None def geometry_from_layout( token_tags: torch.Tensor, position_ids: torch.Tensor, sparsity: float = DEFAULT_SPARSITY, ) -> VSAGeometry | None: """Recover the VSA geometry from what the transformer is actually given. The packed sequence a `t2va` request builds is `[text | audio | video]`, but nothing downstream is told that, so the layout is read back off the two per-row descriptions the transformer already takes: `token_tags` (0 video, 1 text, 2 audio) gives the segment boundaries, and the `(t, h, w)` rotary grid of the video rows gives the shape of the video block. Returns `None` for any layout the sparse path does not cover, so the caller can stay dense. """ state = _REQUEST_STATE cached = state["geometry"] if cached is not None and state["nonce"] == _REQUEST_COUNTER and cached.seq_len == token_tags.shape[0]: return cached tags = token_tags.tolist() seq_len = len(tags) if seq_len == 0 or tags[-1] != 0: return None video_start = seq_len while video_start > 0 and tags[video_start - 1] == 0: video_start -= 1 if video_start == 0: return None prefix_segments: list[int] = [] previous = None for tag in tags[:video_start]: if tag == previous: prefix_segments[-1] += 1 else: prefix_segments.append(1) previous = tag # The video rows are `torch.meshgrid(height_grid, width_grid, indexing="ij")` repeated once per latent frame, so # the grid falls out of the run lengths of the leading rows. grid = position_ids[video_start:] rows_per_frame = int((grid[:, 0] == grid[0, 0]).sum()) grid_w = int((grid[:rows_per_frame, 1] == grid[0, 1]).sum()) if rows_per_frame == 0 or grid_w == 0 or rows_per_frame % grid_w: return None grid_h = rows_per_frame // grid_w num_video_rows = seq_len - video_start if num_video_rows % rows_per_frame: return None grid_t = num_video_rows // rows_per_frame geometry = build_geometry(tuple(prefix_segments), (grid_t, grid_h, grid_w), token_tags.device, sparsity) state["geometry"] = geometry return geometry # -------------------------------------------------------------------------------------------------------------- # The sparse attention itself # -------------------------------------------------------------------------------------------------------------- # One reusable padded tile buffer per projection. Tiles are written by scatter and pad slots are never touched, so a # buffer only has to be re-zeroed when the geometry behind it changes. _TILE_BUFFERS: dict[str, tuple[torch.Tensor, int]] = {} def reset_tile_buffers() -> None: _TILE_BUFFERS.clear() def _tile(x: torch.Tensor, geometry: VSAGeometry, slot: str) -> torch.Tensor: """`[B, S, H, D]` -> the kernel's `[B, H, S_pad, D]`, pad slots zero.""" batch, _, heads, dim = x.shape shape = (batch, heads, geometry.padded_len, dim) cached = _TILE_BUFFERS.get(slot) if cached is None or cached[0].shape != shape or cached[0].dtype != x.dtype or cached[0].device != x.device: buffer = torch.zeros(shape, dtype=x.dtype, device=x.device) else: buffer = cached[0] if cached[1] != id(geometry): buffer.zero_() buffer.index_copy_(2, geometry.untile_index, x.transpose(1, 2)) _TILE_BUFFERS[slot] = (buffer, id(geometry)) return buffer def _pool(tiled: torch.Tensor, geometry: VSAGeometry) -> torch.Tensor: """fp32 masked mean over each 64-token tile. `[B, H, S_pad, D]` -> `[B, H, n_tiles, D]`.""" batch, heads, _, dim = tiled.shape pooled = tiled.view(batch, heads, geometry.n_tiles, TILE_ELEMS, dim).sum(dim=3, dtype=torch.float32) return pooled / geometry.tile_divisor def _block_mask(scores: torch.Tensor, geometry: VSAGeometry) -> torch.Tensor: """Top-k video tiles per query tile, with the prefix exempt and prefix queries dense.""" prefix = geometry.num_prefix_tiles mask = torch.zeros_like(scores, dtype=torch.bool) # sorted=False: the block map only needs the selected SET, not the ranking, and torch.topk's sorted path # full-sorts every row of the `[B, H, tiles, tiles]` score matrix for nothing. indices = scores[..., prefix:].topk(geometry.topk, dim=-1, sorted=False).indices + prefix mask.scatter_(-1, indices, True) mask[..., :prefix] = True mask[:, :, :prefix, :] = True return mask # The sparse-attention kernel, resolved once per process. FastVideo's published CUDA wheel carries the # block-sparse forward as sm_100a source with sm_120a cubins in the same fatbin — this pool's GPU is sm_120 — # but its `is_supported` gate hard-checks capability == (10, 0) and would never select it here. We call the # public `block_sparse_attn_sm100a` entry directly (need_lse=False, the production inference path) and treat # the dense-equivalence `/selftest` as the numerical gate. Any import or launch failure falls back to the # vendored Triton kernels; `H3_VSA_CUDA=0` forces Triton up front. _CUDA_SPARSE: dict = {"resolved": False, "op": None} def _resolve_cuda_sparse_op(): if _CUDA_SPARSE["resolved"]: return _CUDA_SPARSE["op"] _CUDA_SPARSE["resolved"] = True if os.environ.get("H3_VSA_CUDA", "1") != "1": return None try: from fastvideo_kernel import block_sparse_attn_sm100a as sm100a if not sm100a._HAS_VSA_SM100A: # Surface what the extension actually exposes so the logs show which pybind name this wheel has. from fastvideo_kernel._C import fastvideo_kernel_ops as _ops available = sorted(a for a in dir(_ops) if "sparse" in a.lower()) raise ImportError(f"sm100a forward flag unset; pybind exposes {available}") op = sm100a.block_sparse_attn_sm100a _CUDA_SPARSE["op"] = op print("[vsa] CUDA sparse-attention kernel: fastvideo_kernel (sm100a fatbin on sm120)", flush=True) except Exception as error: # noqa: BLE001 - any failure means Triton print(f"[vsa] fastvideo_kernel unavailable ({type(error).__name__}: {error}); using vendored Triton", flush=True) _CUDA_SPARSE["op"] = None return _CUDA_SPARSE["op"] def sparse_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, gate_compress: torch.Tensor | None, geometry: VSAGeometry, ) -> torch.Tensor: """VSA-H3 over one packed sequence. All tensors are `[B, S, H, D]`; the result is too.""" from vsa_kernel import map_to_index, triton_block_sparse_attn_forward query_tiled = _tile(query, geometry, "q") key_tiled = _tile(key, geometry, "k") value_tiled = _tile(value, geometry, "v") scores = torch.matmul(_pool(query_tiled, geometry), _pool(key_tiled, geometry).transpose(-2, -1)) scores = scores / math.sqrt(query.shape[-1]) mask = _block_mask(scores, geometry) q2k_index, q2k_num = map_to_index(mask) # The CUDA path only handles an even tile count (one CTA owns an adjacent pair of query blocks); anything # else — or any launch failure — falls back to the vendored Triton kernels for the rest of the process. out_tiled = None cuda_op = _resolve_cuda_sparse_op() if cuda_op is not None and geometry.n_tiles % 2 == 0: try: out_tiled, _ = cuda_op( query_tiled, key_tiled, value_tiled, q2k_index, q2k_num, geometry.variable_block_sizes, need_lse=False, ) except Exception as error: # noqa: BLE001 - a refused launch must not kill the request print(f"[vsa] CUDA kernel failed ({type(error).__name__}: {error}); falling back to Triton", flush=True) _CUDA_SPARSE["op"] = None out_tiled = None if out_tiled is None: out_tiled, _ = triton_block_sparse_attn_forward( query_tiled, key_tiled, value_tiled, q2k_index, q2k_num, geometry.variable_block_sizes, ) out = out_tiled.index_select(2, geometry.untile_index) if gate_compress is not None: # The compression branch: dense attention over the pooled tiles, broadcast back to every row of its tile and # scaled by the trained gate. Applied here rather than on the tile buffer so the gate never needs one. pooled = torch.matmul(torch.softmax(scores, dim=-1), _pool(value_tiled, geometry)).to(out.dtype) contribution = pooled.index_select(2, geometry.row_tile_index) del pooled contribution.mul_(gate_compress.transpose(1, 2)) out.add_(contribution) del contribution return out.transpose(1, 2) # -------------------------------------------------------------------------------------------------------------- # The processor, and installing it on a loaded transformer # -------------------------------------------------------------------------------------------------------------- # `forward` fills this in from the layout it was handed; the 50 block processors read it. One request at a time — # `@spaces.GPU` serializes them anyway. _ACTIVE: dict[str, VSAGeometry | None] = {"geometry": None} # What the last transformer forward actually ran on. A silent fall back to dense would otherwise be invisible — the # video still comes out, just off-distribution — so every distinct layout logs one line and the Space surfaces the # most recent one in its status banner. LAST_LAYOUT = "no forward yet" _REPORTED: set[str] = set() def _report(summary: str) -> None: global LAST_LAYOUT LAST_LAYOUT = summary if summary not in _REPORTED: _REPORTED.add(summary) print(f"[vsa] {summary}", flush=True) class MiniMaxH3VSAAttnProcessor: """`MiniMaxH3AttnProcessor` with `dispatch_attention_fn` replaced by VSA-H3.""" _attention_backend = None _parallel_config = None def __call__(self, attn, hidden_states, rotary_emb=None, attention_mask=None): from diffusers.models.transformers.transformer_minimax_h3 import _apply_rotary_emb geometry = _ACTIVE["geometry"] query = attn.to_q(hidden_states).unflatten(-1, (attn.heads, -1)) key = attn.to_k(hidden_states).unflatten(-1, (attn.heads, -1)) value = attn.to_v(hidden_states).unflatten(-1, (attn.heads, -1)) query = attn.norm_q(query) key = attn.norm_k(key) if rotary_emb is not None: query = _apply_rotary_emb(query, *rotary_emb) key = _apply_rotary_emb(key, *rotary_emb) if geometry is None or geometry.seq_len != hidden_states.shape[1]: _report(f"DENSE FALLBACK in the attention processor at seq_len={hidden_states.shape[1]}") from diffusers.models.attention_dispatch import dispatch_attention_fn out = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False, backend=self._attention_backend, parallel_config=self._parallel_config, ) else: gate_compress = None gate = getattr(attn, "to_gate_compress", None) if gate is not None: gate_compress = gate(hidden_states).unflatten(-1, (attn.heads, -1)) out = sparse_attention(query, key, value, gate_compress, geometry) out = out.flatten(2, 3).type_as(query) out = attn.to_out[0](out) return attn.to_out[1](out) def add_gate_compress_modules() -> None: """Give every block's attention the `to_gate_compress` projection the base `diffusers` port has no use for. The trained gate lives in the checkpoint as 50 `transformer_blocks.*.attn.to_gate_compress.weight` tensors, but `MiniMaxH3Attention` does not declare the module, so `from_pretrained` reports them as unexpected and drops them. Declaring it before the transformer is instantiated is what makes them load. """ from diffusers.models.transformers import transformer_minimax_h3 as module block_cls = module.MiniMaxH3TransformerBlock if getattr(block_cls, "_vsa_gate_patched", False): return original_init = block_cls.__init__ def patched_init(self, hidden_size, num_attention_heads, attention_head_dim, *args, **kwargs): original_init(self, hidden_size, num_attention_heads, attention_head_dim, *args, **kwargs) self.attn.to_gate_compress = torch.nn.Linear( hidden_size, num_attention_heads * attention_head_dim, bias=False ) block_cls.__init__ = patched_init block_cls._vsa_gate_patched = True def install(transformer, sparsity: float = DEFAULT_SPARSITY) -> tuple[int, int]: """Put the VSA processor on the 50 packed-sequence blocks and wrap `forward` to publish the layout. The token refiner has its own block class, over the *text* stream rather than the packed sequence, so it is untouched. Returns `(blocks, live gates)`. """ installed = 0 gates = 0 for block in transformer.transformer_blocks: gate = getattr(block.attn, "to_gate_compress", None) if gate is not None and not bool((gate.weight != 0).any()): # FastVideo's `_gate_active`: a zero (not-yet-finetuned) gate makes the branch a guaranteed zero — a full # GEMM plus a pooled attention per layer for nothing — so drop it rather than pay for it. block.attn.to_gate_compress = None elif gate is not None: gates += 1 block.attn.set_processor(MiniMaxH3VSAAttnProcessor()) installed += 1 if getattr(transformer, "_vsa_forward_wrapped", False): return installed, gates original_forward = transformer.forward @functools.wraps(original_forward) def forward(*args, **kwargs): token_tags = kwargs.get("token_tags") position_ids = kwargs.get("position_ids") if token_tags is None or position_ids is None: # `MiniMaxH3LoopDenoiser` passes the layout by keyword; bind positionally only as a fallback. import inspect bound = inspect.signature(original_forward).bind_partial(*args, **kwargs).arguments token_tags = token_tags if token_tags is not None else bound.get("token_tags") position_ids = position_ids if position_ids is not None else bound.get("position_ids") geometry = None if token_tags is not None and position_ids is not None: geometry = geometry_from_layout(token_tags, position_ids, sparsity) if geometry is None: _report("DENSE FALLBACK: no VSA geometry could be derived from this forward's layout") else: _report( f"sparse: seq={geometry.seq_len} padded={geometry.padded_len} tiles={geometry.n_tiles} " f"(prefix {geometry.num_prefix_tiles} + video {geometry.num_video_tiles}) " f"topk={geometry.topk} ({geometry.topk / geometry.num_video_tiles:.1%} of video tiles kept)" ) _ACTIVE["geometry"] = geometry try: return original_forward(*args, **kwargs) finally: _ACTIVE["geometry"] = None transformer.forward = forward transformer._vsa_forward_wrapped = True return installed, gates