"""AOTI (Ahead-Of-Time Inductor) support for the VSA-H3 sparse attention soup. Phase A of the AOTI experiment (torch 2.11.0+cu130 on ZeroGPU, where AOTI is officially supported): 1. The two raw Triton kernels the sparse path depends on — `map_to_index` and `triton_block_sparse_attn_forward` — are wrapped as `torch.library` custom ops with fake implementations, which is what lets `torch.export` trace through them. 2. `sparse_attention_functional` reimplements `vsa_h3.sparse_attention` as a *pure tensor function*: op-for-op the same math (fresh tile buffers, fp32 pooled scores, prefix-exempt top-k mask, sparse kernel, untile, compression branch with the trained per-row gate), but the cached `VSAGeometry` tensors arrive as plain arguments. No weights, no globals — small artifacts. 3. `/aoti_diag` on the Space (a) checks the functional reimplementation against the eager reference on the real bench geometry, and (b) when `H3_AOTI=1`, exports + AOTI-compiles the function for this GPU and checks the compiled artifact against eager, with timings. The compiled function is only valid for one static layout (the frozen benchmark spec's). Any other canvas falls back to the eager path. """ from __future__ import annotations import math import torch _LIB = None # registered-once torch.library fragment def _ns() -> torch.library.Library: """Register the custom ops once and return the library fragment.""" global _LIB if _LIB is not None: return _LIB from vsa_kernel import map_to_index as _map_to_index from vsa_kernel import triton_block_sparse_attn_forward as _fwd lib = torch.library.Library("fasth3", "FRAGMENT") def map_fake(block_map: torch.Tensor): return ( torch.empty_like(block_map, dtype=torch.int32), torch.empty(block_map.shape[:-1], dtype=torch.int32, device=block_map.device), ) def fwd_fake(q, k, v, q2k_index, q2k_num, vbs): return ( torch.empty_like(q), torch.empty(q.shape[:3], dtype=torch.float32, device=q.device), ) lib.define("map_to_index(Tensor block_map) -> (Tensor, Tensor)") lib.impl("map_to_index", _map_to_index, "CUDA") lib._register_fake("map_to_index", map_fake) lib.define( "block_sparse_fwd(Tensor q, Tensor k, Tensor v, Tensor q2k_index, Tensor q2k_num, " "Tensor vbs) -> (Tensor, Tensor)" ) lib.impl("block_sparse_fwd", _fwd, "CUDA") lib._register_fake("block_sparse_fwd", fwd_fake) _LIB = lib return lib def sparse_module(topk: int, num_prefix_tiles: int) -> torch.nn.Module: """`torch.export.export` requires an `nn.Module`; wrap the functional with baked scalars.""" class _SparseModule(torch.nn.Module): def __init__(self): super().__init__() self.topk = topk self.num_prefix_tiles = num_prefix_tiles def forward(self, q, k, v, gate, untile_index, variable_block_sizes, tile_divisor): return sparse_attention_functional( q, k, v, gate, untile_index, variable_block_sizes, tile_divisor, self.topk, self.num_prefix_tiles, ) return _SparseModule() def sparse_attention_functional( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, gate: torch.Tensor | None, untile_index: torch.Tensor, variable_block_sizes: torch.Tensor, tile_divisor: torch.Tensor, topk: int, num_prefix_tiles: int, ) -> torch.Tensor: """`vsa_h3.sparse_attention` with every cached buffer inlined as an argument. Op-for-op the same math as the eager path: fp32 pooled tile scores, top-k video tiles with the prefix exempt (prefix key tiles always selected, prefix query tiles dense), the sparse kernel, untiling, then the compression branch scaled by the trained per-row gate. `gate` is `[B, H, S_real, D]` (the eager path's `gate.transpose(1, 2)`), applied after untililing. """ _ns() # ensure the fasth3 custom ops exist before any torch.ops.fasth3.* call batch, heads, seq_len, dim = query.shape # inputs are `[B, H, S_real, D]` (BHSD) n_tiles = variable_block_sizes.numel() padded_len = n_tiles * 64 # Tile: scatter the packed rows into the kernels' `[B, H, S_pad, D]` layout, pad slots zero. query_tiled = torch.zeros((batch, heads, padded_len, dim), dtype=query.dtype, device=query.device) query_tiled.index_copy_(2, untile_index, query) key_tiled = torch.zeros_like(query_tiled) key_tiled.index_copy_(2, untile_index, key) value_tiled = torch.zeros_like(query_tiled) value_tiled.index_copy_(2, untile_index, value) # Per-head fp32 pooled tile scores. q_pool = query_tiled.view(batch, heads, n_tiles, 64, dim).sum(dim=3, dtype=torch.float32) / tile_divisor k_pool = key_tiled.view(batch, heads, n_tiles, 64, dim).sum(dim=3, dtype=torch.float32) / tile_divisor scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) / math.sqrt(dim) # Prefix-exempt top-k mask, assembled from fresh tensors (functional, export-traceable): # video query rows keep the top-k VIDEO kv tiles, prefix kv tiles are always selected, # prefix query rows are dense over everything. num_q_video = n_tiles - num_prefix_tiles video_kv_scores = scores[..., num_prefix_tiles:, num_prefix_tiles:] # [B, H, nq_video, kv_video] indices = video_kv_scores.topk(topk, dim=-1, sorted=False).indices mask_video_kv = torch.zeros_like(video_kv_scores, dtype=torch.bool).scatter(-1, indices, True) prefix_cols = torch.ones( (batch, heads, num_q_video, num_prefix_tiles), dtype=torch.bool, device=query.device ) mask_video = torch.cat([prefix_cols, mask_video_kv], dim=-1) # [B, H, nq_video, n_tiles] mask = torch.cat( [ torch.ones((batch, heads, num_prefix_tiles, n_tiles), dtype=torch.bool, device=query.device), mask_video, ], dim=2, ) q2k_index, q2k_num = torch.ops.fasth3.map_to_index(mask) out_tiled, _ = torch.ops.fasth3.block_sparse_fwd( query_tiled, key_tiled, value_tiled, q2k_index, q2k_num, variable_block_sizes ) out = out_tiled.index_select(2, untile_index) if gate is not None: v_pool = value_tiled.view(batch, heads, n_tiles, 64, dim).sum(dim=3, dtype=torch.float32) / tile_divisor pooled = torch.matmul(torch.softmax(scores, dim=-1), v_pool).to(out.dtype) out = out + pooled.index_select(2, untile_index // 64) * gate return out.transpose(1, 2)