exp10a: AOTI diag endpoint

aoti_attention.py

@@ -0,0 +1,133 @@
+"""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_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.
+    """
+    batch, heads, padded_len, dim = query.shape
+    n_tiles = variable_block_sizes.numel()
+    n_real = untile_index.numel()
+    num_q_video = n_tiles - num_prefix_tiles
+
+    # Tile: scatter the packed rows into the kernels' `[B, H, S_pad, D]` layout, pad slots zero.
+    query_tiled = torch.index_copy(
+        torch.zeros((batch, heads, padded_len, dim), dtype=query.dtype, device=query.device),
+        2, untile_index, query.transpose(1, 2),
+    )
+    key_tiled = torch.index_copy(
+        torch.zeros_like(query_tiled), 2, untile_index, key.transpose(1, 2),
+    )
+    value_tiled = torch.index_copy(
+        torch.zeros_like(query_tiled), 2, untile_index, value.transpose(1, 2),
+    )
+
+    # 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_scores = scores[..., num_prefix_tiles:, num_prefix_tiles:]
+    indices = video_scores.topk(topk, dim=-1, sorted=False).indices
+    mask_video = torch.zeros_like(video_scores, dtype=torch.bool).scatter(-1, indices, True)
+    prefix_cols = torch.zeros_like(video_scores, dtype=torch.bool)
+    prefix_cols[..., :num_prefix_tiles] = True
+    mask_video = mask_video | prefix_cols
+    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)

app.py

@@ -427,6 +427,129 @@ def _generate(prompt_embeds, text_token_tags, height: int, width: int, num_frame
     )
 
 
+@spaces.GPU(duration=300, size=GPU_SIZE)
+def aoti_diag() -> str:
+    """Phase A of the AOTI experiment: validate the traceable functional reimplementation of the VSA
+    sparse-attention soup, and (when `H3_AOTI=1`) export + AOTI-compile it for this GPU and check the
+    artifact against eager, with timings.
+
+    Uses the *real bench geometry* (1344x768, 124 frames, 203 text + 414 audio + 37 296 video rows),
+    synthesized directly — no generation, no weights.
+    """
+    import torch
+
+    import aoti_attention
+
+    device = torch.device("cuda")
+    heads, dim = 56, 128
+    text_rows, audio_rows = 203, 414
+    grid_t, grid_h, grid_w = 37, 24, 42
+    n_video = grid_t * grid_h * grid_w
+    seq_len = text_rows + audio_rows + n_video
+
+    tags = torch.tensor([1] * text_rows + [2] * audio_rows + [0] * n_video, dtype=torch.long)
+    position_ids = torch.zeros(seq_len, 3, dtype=torch.float64)
+    video_start = text_rows + audio_rows
+    frame = torch.cartesian_prod(torch.arange(float(grid_h)), torch.arange(float(grid_w)))
+    position_ids[video_start:, 0] = torch.arange(float(grid_t)).repeat_interleave(grid_h * grid_w)
+    position_ids[video_start:, 1:] = frame.repeat(grid_t, 1)
+
+    import vsa_h3
+
+    geometry = vsa_h3.geometry_from_layout(tags, position_ids, 0.9)
+    if geometry is None:
+        return "**FAILED**: synthetic bench layout did not yield VSA geometry."
+
+    def timed(fn, *args, repeats=3):
+        fn(*args)  # warm (JIT/compile)
+        torch.cuda.synchronize()
+        import time as _t
+
+        started = _t.perf_counter()
+        for _ in range(repeats):
+            fn(*args)
+        torch.cuda.synchronize()
+        return (_t.perf_counter() - started) / repeats
+
+    generator = torch.Generator(device=device).manual_seed(0)
+    shape = (1, seq_len, heads, dim)
+    q, k, v = (torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) for _ in range(3))
+    gate = torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) * 0.05
+
+    lines = [f"`{torch.cuda.get_device_name()}` · seq {seq_len}, padded {geometry.padded_len}, "
+             f"tiles {geometry.n_tiles}, topk {geometry.topk}", ""]
+
+    def eager(gate_h):
+        return vsa_h3.sparse_attention(q, k, v, gate_h.transpose(1, 2), geometry)
+
+    def functional(gate_h):
+        return aoti_attention.sparse_attention_functional(
+            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate_h.transpose(1, 2),
+            geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
+            geometry.topk, geometry.num_prefix_tiles,
+        )
+
+    for label, gate_h in (("gate=None", None), ("gate=rand*0.05", gate)):
+        ref, got = eager(gate_h), functional(gate_h)
+        error = (ref.float() - got.float()).abs().max().item()
+        scale = ref.float().abs().max().item()
+        cosine = torch.nn.functional.cosine_similarity(
+            ref.float().flatten(), got.float().flatten(), dim=0
+        ).item()
+        t_eager = timed(eager, gate_h)
+        t_fn = timed(functional, gate_h)
+        lines.append(
+            f"| {label} | eager {t_eager * 1000:.0f} ms | functional {t_fn * 1000:.0f} ms | "
+            f"rel err {error / scale:.2e} | cosine {cosine:.6f} |"
+        )
+
+    report = (
+        "VSA soup: eager vs functional (op-for-op reimplementation)\n\n"
+        "| case | eager | functional | rel err | cosine |\n|---|---|---|---|---|\n" + "\n".join(lines[2:])
+    )
+
+    if os.environ.get("H3_AOTI") == "1":
+        try:
+            from torch._inductor import aoti_compile_and_package, aoti_load_package
+
+            with torch.no_grad():
+                exported = torch.export.export(
+                    aoti_attention.sparse_attention_functional,
+                    args=(
+                        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
+                        gate.transpose(1, 2), geometry.untile_index,
+                        geometry.variable_block_sizes, geometry.tile_divisor,
+                    ),
+                    kwargs={"topk": geometry.topk, "num_prefix_tiles": geometry.num_prefix_tiles},
+                )
+            package = "/tmp/vsa_sparse_aoti.pt2"
+            aoti_compile_and_package(exported, package_path=package)
+            compiled = aoti_load_package(package)
+            got = compiled(
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate.transpose(1, 2),
+                geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
+            )
+            ref = eager(gate)
+            error = (ref.float() - got.float()).abs().max().item()
+            scale = ref.float().abs().max().item()
+            cosine = torch.nn.functional.cosine_similarity(
+                ref.float().flatten(), got.float().flatten(), dim=0
+            ).item()
+            t_eager = timed(eager, gate)
+            t_compiled = timed(lambda *a: compiled(*a), (
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate.transpose(1, 2),
+                geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
+            ))
+            report += (
+                f"\n\nAOTI artifact: `{package}`\n\n| eager | compiled | rel err | cosine |\n|---|---|---|---|\n"
+                f"| {t_eager * 1000:.0f} ms | {t_compiled * 1000:.0f} ms | {error / scale:.2e} | {cosine:.6f} |"
+            )
+        except Exception as error:  # noqa: BLE001 - surfaced verbatim for the log
+            report += f"\n\n**AOTI FAILED**: `{type(error).__name__}: {error}`"
+
+    return report
+
+
 @spaces.GPU(duration=120, size=GPU_SIZE)
 def selftest() -> str:
     """Check the vendored VSA-H3 kernels against dense attention on this GPU.
@@ -710,9 +833,11 @@ with gr.Blocks(title="FastH3 v1 (VSA)") as demo:
     demo.load(status, None, banner, api_name="status")
 
     # No UI, API only: the sparse-attention equivalence check, so the kernel can be verified on this pool without
-    # spending a full generation.
+    # spending a full generation; and the AOTI Phase A validation.
     diagnose = gr.Button(visible=False)
     diagnose.click(selftest, None, gr.Markdown(visible=False), api_name="selftest")
+    diagnose2 = gr.Button(visible=False)
+    diagnose2.click(aoti_diag, None, gr.Markdown(visible=False), api_name="aoti_diag")
 
 
 if __name__ == "__main__":