Upload folder using huggingface_hub

README.md

@@ -1,13 +1,166 @@
 ---
-title: Fasth3 4step Preview Demo
-emoji: 🐠
-colorFrom: pink
-colorTo: purple
+title: FastH3 4-step Preview
+emoji: 🎬
+colorFrom: red
+colorTo: gray
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: 4-step MiniMax-H3 with sparse attention — video + audio
+python_version: "3.12"
+startup_duration_timeout: 1h
+suggested_hardware: zero-a10g
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FastH3 4-step Preview (VSA, data-free) — MiniMax-H3 in four forward passes
+
+[`FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree`](https://huggingface.co/FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree)
+is a **data-free DMD2 distillation** of [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3), the
+33B dual-modality DiT that emits a video *and* its synchronized soundtrack from one denoising pass. The student
+keeps the teacher's architecture exactly and only replaces `transformer/`, so it is a drop-in for the same
+`diffusers` modular pipeline — but it needs **four** DiT forwards instead of thirty, and it was distilled **with
+Video Sparse Attention on**.
+
+Everything runs unquantized at **bfloat16**.
+
+## The sampling contract: five grid points, four forwards
+
+`MiniMaxH3Scheduler.set_timesteps(n)` builds `linspace(1, 0, n)`, applies the shift
+`σ' = s·σ / (1 + (s−1)·σ)`, and then drops the trailing zero when it forms the timesteps — so **`n` sigma grid
+points drive `n − 1` model evaluations**. The distilled ladder is `t = 999, 749, 500, 250 → 0`: five points, four
+forwards.
+
+This Space therefore fixes `num_inference_steps = 5`, matching the checkpoint's own
+`fastvideo_inference.json` (`num_inference_steps: 5`, `transformer_forwards: 4`,
+`dmd_denoising_steps: [999, 749, 500, 250]`, `guidance_scale: 1.0`). It is not a knob. There is no CFG and no
+negative prompt — the teacher is guidance-distilled and the student inherits that.
+
+## Video Sparse Attention is not optional here
+
+This checkpoint is the **VSA** variant. Its `fastvideo_inference.json` pins
+`attention_backend: VIDEO_SPARSE_ATTN_H3`, `vsa_sparsity: 0.9`, `vsa_tile_size: 64`, and the transformer ships 50
+trained `attn.to_gate_compress` tensors (~3.6 GiB) that only the sparse path consumes. FastVideo publishes a
+separate `…-Dense-DataFree` checkpoint for people who want dense — running *this* one dense is running it
+off-distribution.
+
+The published kernel is `vsa_kernel: sm100a`, which is GB200-only, and the ZeroGPU pool is **sm120** (RTX PRO 6000
+Blackwell). FastVideo's other officially supported route is `--vsa-kernel triton`, which is pure Triton and
+architecture-agnostic — so this Space vendors those two files verbatim from
+[FastVideo](https://github.com/hao-ai-lab/FastVideo) (Apache-2.0) into `vsa_kernel/` and ports the H3 backend on top
+of them:
+
+| File | What it is |
+|---|---|
+| `vsa_kernel/block_sparse_attn_triton.py` | FastVideo's Triton block-sparse attention, verbatim except that the autotune sweep is collapsed to the single config it lands on for Blackwell (re-enable it with `H3_VSA_AUTOTUNE=1`). |
+| `vsa_kernel/index.py` | FastVideo's `map_to_index` / `topk_index_to_map`, verbatim. |
+| `vsa_h3.py` | The port: `MiniMaxH3VSAAttnProcessor`, a `diffusers` attention processor reproducing `MiniMaxH3VSABackend`. |
+
+`vsa_h3.py` follows FastVideo's `video_sparse_attn_h3.py` step for step: 64-token `(4, 4, 4)` tiles over the
+post-patchify video grid, segment-pure prefix tiles, per-head fp32 pooled tile scores,
+`topk = max(1, min(⌈(1 − sparsity)·n_video_tiles⌉, n_video_tiles))`, prefix keys exempt (always selected) and prefix
+queries always dense, plus the compression branch `softmax(scores) @ pool(v)` broadcast back over each tile row and
+scaled by the trained gate with **no** activation. The tile geometry is derived per-forward from the pipeline's own
+`token_tags` / `position_ids`, so the `[text | cond | audio | video]` packing stays authoritative.
+
+`diffusers` 0.40.0 has no `to_gate_compress`, so `vsa_h3.add_gate_compress_modules()` patches
+`MiniMaxH3TransformerBlock.__init__` **before** the pipeline loads; otherwise the 50 gate tensors load as
+"unexpected keys" and are silently dropped. Blocks whose gate is all-zero have it removed again after load, exactly
+as FastVideo does.
+
+A hidden `/selftest` API endpoint runs the ported kernel at sparsity 0 against `F.scaled_dot_product_attention` on
+a real packed layout, so a wrong tile order or transpose is caught without spending a generation.
+
+## Split across two Spaces
+
+MiniMax-H3 is ~196 GiB in bfloat16 and a ZeroGPU Space is evicted at **150 GB of storage**, so no single
+unquantized Space can hold it. Cutting `MiniMaxH3Blocks` at its `text_encoder` step splits it in two, and both
+halves fit:
+
+| Space | Subfolders | Download |
+|---|---|---|
+| [`multimodalart/qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner) | `text_encoder/` + `tokenizer/` + `processor/` | 66.7 GB → 62.15 GiB bf16 |
+| this one | `transformer/` + `vae/` + `audio_vae/` | ~85 GB → 75.7 GiB resident |
+
+The conditioner is a **public** Space and is unchanged by the distillation — the student's
+`modular_model_index.json` points its `text_encoder` at the same Qwen3-VL weights — so this Space reuses it as-is
+and calls it over the gradio API for every request.
+
+`h3_split_blocks.py` subclasses the pipeline's blocks with the `text_encoder` step removed. Dropping the step drops
+the components it declares, so `load_components` resolves only `transformer` / `vae` / `audio_vae` / the two
+schedulers, and `prompt_embeds` + `text_token_tags` become ordinary required pipeline inputs. The wire format is
+those two tensors — `(1, num_text_tokens, 5120)` bfloat16 and `(num_text_tokens,)` int64 — in one safetensors file
+with the resolved `height` / `width` / `num_frames` in its metadata header.
+
+## Text-to-video+audio only
+
+The preview distills **the T2VA path only**, so this Space exposes no keyframe / reference inputs — the student's
+`transformer_ref` tower is not packaged with the checkpoint. Image conditioning is what
+[`multimodalart/minimax-h3`](https://huggingface.co/spaces/multimodalart/minimax-h3) (the undistilled 30-step
+teacher) is for.
+
+## Prompt format
+
+MiniMax-H3 is trained on a structured multimodal caption, not a bare sentence:
+
+```
+integrated_multimodal_description: <shots, camera, subjects, action, lighting>
+overall_soundscape: <diegetic sound>
+non_diegetic_music: <score>
+```
+
+**Expand prompt** (on by default) sends a short prompt through the conditioner's Qwen3-VL prompt rewriter, which
+writes that structure for you and returns it. Turn it off when you have already written a full-format prompt — the
+examples that carry MiniMax's own official prompts do exactly that.
+
+## Examples
+
+The two long examples are MiniMax's own published prompts, taken verbatim from the base model's repo (Apache-2.0
+code / docs):
+
+- the starship-bridge two-shot from `scripts/readme/reproducible-768p-t2va-request.sh`
+- the bakery two-shot, Case 1 of `docs/VIDEO_PROMPT_WRITING_GUIDE_base_en.md`
+
+## Generation constraints
+
+Fixed by the checkpoint: 24 fps, a 768 pixel short edge at the training canvas, `num_frames` snapped up to the next
+`17·n + 5`. The distillation's operating point is **1344×768 × 124 frames (≈5 s)** — that is the default, and it is
+exactly the layout VSA was tuned on (`grid = 37×24×42`, 37 296 video rows, 672 tiles, 66 selected). The duration
+slider reaches 8 s and the canvas dropdown offers smaller/faster grids, both outside the distilled operating point,
+so quality degrades gracefully rather than being guaranteed. The 8 s ceiling is a VRAM limit, not a model limit:
+75.7 GiB of weights sit resident on a 95.0 GiB card and the sparse working set grows with sequence length.
+
+The checkpoint itself is a **preview**.
+
+## Placement
+
+`H3_PLACEMENT=lazy`: the weights move onto the card on the first GPU call and stay there. `spaces`' startup
+`torch.pack()` would write a second on-disk copy of every resident CUDA tensor, and 85 + 75 GB exceeds the 150 GB
+quota, so packing is not an option here. The one-time `.to("cuda")` plus the Triton JIT of the block-sparse kernels
+lands inside the first request of a cold worker; after that there is **no offloading in the request path at all**.
+
+## Space variables
+
+| Variable | Default | Meaning |
+|---|---|---|
+| `H3_MODEL_REPO` | `FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` | The distilled diffusers-layout checkpoint. Public. |
+| `H3_CONDITIONER` | `multimodalart/qwen3vl-conditioner` | The public Space this one asks for embeddings. |
+| `H3_ATTENTION` | `vsa` | The ported VSA-H3 backend. Any other value is passed to `set_attention_backend` as a dense escape hatch (e.g. `_native_cudnn`). |
+| `H3_VSA_SPARSITY` | `0.9` | The checkpoint's trained sparsity. |
+| `H3_VSA_AUTOTUNE` | unset | Set to `1` to restore FastVideo's full Triton autotune sweep instead of the pinned Blackwell config. |
+| `H3_PLACEMENT` | `lazy` | `lazy` moves all weights onto the card on the first GPU call; `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload`. |
+| `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` (48 GB) does not fit 75 GiB of weights. |
+
+## Secrets
+
+None. Every weight this Space downloads is public, and the conditioner is a public Space called on the requesting
+user's own ZeroGPU token.
+
+## License
+
+The weights are under the **MiniMax H3 Community License**, inherited from the base model — it carries territory and
+acceptable-use restrictions. Read
+[the base model's license](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/LICENSE) before using outputs. The
+vendored kernels in `vsa_kernel/` are Apache-2.0, from
+[hao-ai-lab/FastVideo](https://github.com/hao-ai-lab/FastVideo).
+</content>
+</invoke>

app.py

@@ -0,0 +1,579 @@
+"""FastH3 v1 (VSA) — the 4-step DMD2 distillation of MiniMax-H3, text to video + synchronized audio.
+
+`FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` replaces only the `transformer/` of the MiniMax-H3 release
+with a data-free DMD2 student. Everything else in the repo (Qwen3-VL conditioner, both autoencoders, both schedulers)
+is an unmodified copy of the base checkpoint, so it runs on the released `diffusers` modular pipeline — the two
+differences at inference time are the step count and the **attention backend**.
+
+**The sampling contract.** `num_inference_steps` counts sigma *grid points*, and `N` points drive `N - 1` transformer
+forwards. The checkpoint's own `fastvideo_inference.json` states it exactly: `num_inference_steps: 5`,
+`transformer_forwards: 4`, `dmd_denoising_steps: [999, 749, 500, 250]`, `guidance_scale: 1.0`. It is fixed here.
+
+**VSA is not optional.** That same file records `attention_backend: VIDEO_SPARSE_ATTN_H3`, `vsa_tile_size: 64`,
+`vsa_sparsity: 0.9`. This student was distilled *under* block-sparse attention and ships 50 trained
+`attn.to_gate_compress` tensors that only the sparse path reads, so `vsa_h3.py` ports FastVideo's VSA-H3 backend onto
+`MiniMaxH3Attention` and runs it, on FastVideo's own Triton kernels (vendored under `vsa_kernel/`). The checkpoint's
+`vsa_kernel: sm100a` is the GB200-only fast path for the same mask semantics; this pool is sm120.
+
+**Why the Space is split.** MiniMax-H3 is ~196 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage.
+This half holds the distilled transformer and the two autoencoders (81 GB); the 62.15 GiB Qwen3-VL conditioner runs in
+`multimodalart/qwen3vl-conditioner`, which this Space calls over the gradio API for every request. FastH3 ships the
+base release's conditioner verbatim, so that Space encodes this checkpoint exactly. Nothing is quantized anywhere.
+"""
+
+from __future__ import annotations
+
+import os
+import tempfile
+import time
+import traceback
+from functools import cache
+
+# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 75 GiB load can happen at
+# startup rather than on GPU time.
+import spaces
+import gradio as gr
+
+MODEL_REPO = os.environ.get("H3_MODEL_REPO", "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree")
+BASE_REPO = "MiniMaxAI/MiniMax-H3"
+CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
+# `lazy` moves all weights onto the card on the first GPU call and leaves them there; `offload` hands placement to
+# `ComponentsManager.enable_auto_cpu_offload`. Packing at startup is not an option here: `spaces` writes every
+# startup-resident CUDA tensor to a second on-disk copy, and this checkpoint's transformer (70.1 GB on disk) would
+# bust the 150 GB quota.
+PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
+# `vsa` is the trained route (see the module docstring). `dense` is the escape hatch: it runs the released dense
+# `diffusers` path with cuDNN's fused kernel, which is off-distribution for this student but useful to bisect against.
+ATTENTION = os.environ.get("H3_ATTENTION", "vsa").lower()
+# The checkpoint's own `vsa_sparsity`. Only read when `H3_ATTENTION=vsa`.
+VSA_SPARSITY = float(os.environ.get("H3_VSA_SPARSITY", "0.9"))
+# 75.7 GiB of weights plus activations does not fit a `large` (48 GiB) allocation.
+GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
+
+# The distilled ladder, as sigma grid points. 5 points -> 4 transformer forwards at t = 1000, 750, 500, 250.
+SIGMA_GRID_POINTS = 5
+NUM_FORWARDS = SIGMA_GRID_POINTS - 1
+
+# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
+# is rejected there and surfaces as a failure here.
+CANVASES = {
+    # 16:9
+    "960x544 · 16:9 fast": (544, 960),
+    "1024x576 · 16:9 fast": (576, 1024),
+    "1152x640 · 16:9": (640, 1152),
+    "1280x704 · 16:9": (704, 1280),
+    "1344x768 · 16:9 full": (768, 1344),
+    # 9:16
+    "544x960 · 9:16 fast": (960, 544),
+    "640x1152 · 9:16": (1152, 640),
+    "768x1344 · 9:16 full": (1344, 768),
+    # 1:1
+    "544x544 · 1:1 fast": (544, 544),
+    "768x768 · 1:1 full": (768, 768),
+    "1024x1024 · 1:1 max": (1024, 1024),
+    # 4:3 / 3:4
+    "768x576 · 4:3 fast": (576, 768),
+    "1024x768 · 4:3 full": (768, 1024),
+    "576x768 · 3:4 fast": (768, 576),
+    "768x1024 · 3:4 full": (1024, 768),
+    # 21:9
+    "1152x512 · 21:9 fast": (512, 1152),
+    "1536x672 · 21:9 full": (672, 1536),
+}
+# The distillation's own operating point: 768x1344, 124 frames, 24 fps.
+DEFAULT_CANVAS = "1344x768 · 16:9 full"
+DEFAULT_DURATION = 5
+FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
+# The ceiling holds for the *snapped* frame count. 75.74 GiB of weights are resident on a 95.0 GiB card, and the
+# sparse path's own working set grows with the packed sequence, so this is a memory ceiling, not a policy one.
+MIN_UI_DURATION, MAX_UI_DURATION = 2, 8
+
+
+def snap_frames(seconds: float) -> int:
+    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
+    frames = max(1, round(float(seconds) * FPS))
+    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
+        frames += 1
+    return frames
+
+
+def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
+    """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the released checkpoint."""
+    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+
+    MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
+
+
+PIPE = None
+MANAGER = None
+LOAD_ERROR: str | None = None
+LOADED_IN: float | None = None
+VSA_BLOCKS = 0
+VSA_GATES = 0
+
+
+def status() -> str:
+    if LOAD_ERROR:
+        return LOAD_ERROR
+    if PIPE is None:
+        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 81 GB). Watch the Space logs."
+    if ATTENTION == "vsa":
+        attention = (
+            f"**VSA-H3** block-sparse, tile 64 / sparsity {VSA_SPARSITY:g} on {VSA_BLOCKS} blocks "
+            f"({VSA_GATES} trained compression gates live)"
+        )
+    else:
+        attention = f"dense `{ATTENTION}` (off-distribution for this student)"
+    return (
+        f"Ready · distilled transformer + VAEs **bfloat16, unquantized** · {NUM_FORWARDS} transformer forwards "
+        f"({SIGMA_GRID_POINTS}-point sigma grid) · attention {attention} · placement `{PLACEMENT}` · "
+        f"loaded in {LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}`"
+    )
+
+
+def load_models() -> str | None:
+    """Load the denoising half at startup.
+
+    `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and `video_processor`,
+    so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` are never touched.
+    Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a bfloat16 audio VAE decodes
+    the soundtrack roughly 20 dB too quiet.
+
+    `add_gate_compress_modules()` has to run *before* the transformer is instantiated. The checkpoint carries 50
+    `transformer_blocks.*.attn.to_gate_compress.weight` tensors — the trained VSA compression gate — and stock
+    `MiniMaxH3Attention` does not declare the module, so `from_pretrained` would report them as unexpected and drop
+    them. Declaring it first is what makes them load.
+    """
+    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, VSA_BLOCKS, VSA_GATES
+
+    if PIPE is not None or LOAD_ERROR is not None:
+        return LOAD_ERROR
+
+    started = time.time()
+    try:
+        import torch
+        from diffusers import ComponentsManager
+
+        from h3_split_blocks import MiniMaxH3GeneratorBlocks
+
+        lower_duration_floor()
+        if ATTENTION == "vsa":
+            import vsa_h3
+
+            vsa_h3.add_gate_compress_modules()
+
+        manager = ComponentsManager()
+        blocks = MiniMaxH3GeneratorBlocks()
+        print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
+        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="fasth3")
+        pipe.load_components(dtype=torch.bfloat16)
+
+        if ATTENTION == "vsa":
+            VSA_BLOCKS, VSA_GATES = vsa_h3.install(pipe.transformer, sparsity=VSA_SPARSITY)
+            print(f"[gen] VSA-H3 on {VSA_BLOCKS} blocks, {VSA_GATES} trained gates", flush=True)
+        else:
+            pipe.transformer.set_attention_backend(ATTENTION)
+
+        if PLACEMENT == "offload":
+            manager.enable_auto_cpu_offload(device="cuda")
+            _arm_decode_hooks(pipe)
+
+        PIPE, MANAGER = pipe, manager
+        LOADED_IN = time.time() - started
+        print(f"[gen] ready in {LOADED_IN:.0f}s", flush=True)
+    except Exception as error:
+        traceback.print_exc()
+        LOAD_ERROR = (
+            f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: "
+            f"`{type(error).__name__}: {error}`"
+        )
+    return LOAD_ERROR
+
+
+def _arm_decode_hooks(pipe):
+    """Make the offload hooks fire for the two VAEs.
+
+    `enable_auto_cpu_offload` wraps `forward`, and the decode blocks call `vae.decode(...)` directly, so the hook
+    never runs and the VAE is still on the host when the latents arrive on the card.
+    """
+    for name in ("vae", "audio_vae"):
+        module = getattr(pipe, name)
+        inner = module.decode
+
+        def armed(*args, _module=module, _decode=inner, **kwargs):
+            hook = getattr(_module, "_hf_hook", None)
+            if hook is not None:
+                hook.pre_forward(_module)
+            return _decode(*args, **kwargs)
+
+        module.decode = armed
+
+
+@cache
+def conditioner():
+    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, so the
+    conditioner's booking is billed to whoever asked for the video."""
+    from gradio_client import Client
+
+    return Client(CONDITIONER_SPACE)
+
+
+def encode_remote(prompt: str, canvas: str, num_frames: int, rewrite_prompt: bool = False):
+    """`/encode` on the conditioner Space: a safetensors file holding `prompt_embeds` + `text_token_tags`, with the
+    resolved `height` / `width` / `num_frames` in its metadata, plus the plan. `canvas` is the label."""
+    from safetensors import safe_open
+
+    path, plan = conditioner().predict(
+        prompt=prompt,
+        image_path=None,
+        last_image_path=None,
+        canvas=canvas,
+        num_frames=num_frames,
+        rewrite_prompt=bool(rewrite_prompt),
+        api_name="/encode",
+    )
+    with safe_open(path, framework="pt") as handle:
+        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan
+
+
+# Seconds of GPU one request needs, from the packed rows it is about to denoise: linear in the rows for the matmuls
+# and the two decoders, quadratic for the selection scores (which stay dense over tiles even though attention does
+# not). Refit on this Space's own VSA measurements.
+_DUR_A, _DUR_B, _DUR_BASE = 7.413e-4, 1.7507e-8, 5.0
+# `lazy` placement: 75.74 GiB crosses PCIe once per cold worker. Plus the Triton kernels' one-time JIT compile, which
+# is per new packed length and lands on whoever asks for that canvas first on a cold worker.
+_PLACEMENT_ALLOWANCE = 30
+# Booked over the estimate. Keep it small: an inflated duration burns the visitor's quota and drops queue priority.
+_MARGIN = 1.15
+
+
+def get_duration(prompt_embeds, text_token_tags, height, width, num_frames, seed, *a, **k):
+    height, width, num_frames = int(height), int(width), int(num_frames)
+    latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
+    rows = latent_frames * (height // 32) * (width // 32)
+    estimate = _DUR_A * rows + _DUR_B * rows**2 + _DUR_BASE + _PLACEMENT_ALLOWANCE
+    return max(60, int(estimate * _MARGIN) + 2)
+
+
+@spaces.GPU(duration=get_duration, size=GPU_SIZE)
+def _generate(prompt_embeds, text_token_tags, height: int, width: int, num_frames: int, seed: int):
+    """The only thing on GPU time: the four-forward packed-sequence denoise loop and the two decoders.
+
+    Only the three generated outputs come back — a `@spaces.GPU` return crosses a process boundary by pickling, and
+    the full `PipelineState` still holds the packed latents, the rotary grid and the row indices on the card.
+    """
+    import torch
+
+    if PLACEMENT == "lazy":
+        PIPE.to("cuda")
+
+    torch.cuda.reset_peak_memory_stats()
+    state = PIPE(
+        prompt_embeds=prompt_embeds.to("cuda"),
+        text_token_tags=text_token_tags,
+        height=height,
+        width=width,
+        num_frames=num_frames,
+        num_inference_steps=SIGMA_GRID_POINTS,
+        generator=torch.Generator("cpu").manual_seed(int(seed)),
+    )
+    peak = torch.cuda.max_memory_allocated() / 1024**3
+    print(f"[gen] peak allocated {peak:.2f} GiB", flush=True)
+    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
+
+
+@spaces.GPU(duration=120, size=GPU_SIZE)
+def selftest() -> str:
+    """Check the vendored VSA-H3 kernels against dense attention on this GPU.
+
+    At `sparsity = 0` the block map is all-true, so VSA-H3 has to reproduce full attention exactly (up to the tile
+    padding and the fp32 pooled selection, which cannot change an all-true mask). That is the one assertion that
+    catches a wrong tile order, a wrong `variable_block_sizes`, or a mis-transposed buffer — all of which would
+    otherwise show up only as a subtly wrong video. Runs on random tensors; no weights are touched.
+
+    Returns:
+        A markdown report: the dense-equivalence error, and how much of the dense output the 90%-sparse path keeps.
+    """
+    import torch
+    import torch.nn.functional as F
+
+    from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLayoutStep
+
+    import vsa_h3
+
+    device = torch.device("cuda")
+    heads, dim = 8, 128
+    # A real packed layout, just a small one: 300 text rows, 5 latent frames of 8x12 video, its soundtrack.
+    _, token_tags, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence(
+        torch.ones(300, dtype=torch.long), 5, 16, 24, 50, (1, 2, 2), 2, 2, 0, ()
+    )
+    position_ids = torch.zeros(token_tags.numel(), 3, dtype=torch.float64)
+    video_start = int((token_tags == 0).nonzero()[0])
+    frame = torch.cartesian_prod(torch.arange(8.0), torch.arange(12.0))
+    position_ids[video_start:, 0] = torch.arange(5.0).repeat_interleave(96)
+    position_ids[video_start:, 1:] = frame.repeat(5, 1)
+    token_tags, position_ids = token_tags.to(device), position_ids.to(device)
+
+    lines = []
+    generator = torch.Generator(device=device).manual_seed(0)
+    shape = (1, token_tags.numel(), heads, dim)
+    query, key, value = (
+        torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) for _ in range(3)
+    )
+    reference = F.scaled_dot_product_attention(
+        query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+    ).transpose(1, 2)
+
+    for sparsity in (0.0, 0.9):
+        vsa_h3.reset_tile_buffers()
+        geometry = vsa_h3.geometry_from_layout(token_tags, position_ids, sparsity)
+        if geometry is None:
+            return "**FAILED**: `geometry_from_layout` did not recognize the standard packed layout."
+        out = vsa_h3.sparse_attention(query, key, value, None, geometry)
+        error = (out.float() - reference.float()).abs().max().item()
+        scale = reference.float().abs().max().item()
+        similarity = F.cosine_similarity(out.float().flatten(), reference.float().flatten(), dim=0).item()
+        lines.append(
+            f"| {sparsity:g} | {geometry.topk}/{geometry.num_video_tiles} | {error:.4f} | "
+            f"{error / scale:.2e} | {similarity:.6f} |"
+        )
+        if sparsity == 0.0 and error / scale > 0.02:
+            lines.append(f"\n**FAILED**: dense-equivalent VSA differs from SDPA by {error / scale:.3f} relative.")
+
+    return (
+        f"VSA-H3 on `{torch.cuda.get_device_name()}`, {token_tags.numel()} packed rows, "
+        f"{heads} heads x {dim}.\n\n"
+        "| sparsity | tiles kept | max abs err | relative | cosine |\n|---|---|---|---|---|\n" + "\n".join(lines)
+    )
+
+
+def generate(
+    prompt: str,
+    canvas: str = DEFAULT_CANVAS,
+    duration: float = DEFAULT_DURATION,
+    upsample: bool = True,
+    seed: int = 42,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate a video with a synchronized soundtrack from a text prompt, in four transformer forwards.
+
+    Args:
+        prompt: The request. MiniMax-H3 was trained on a structured format
+            (`integrated_multimodal_description: ... overall_soundscape: ... non_diegetic_music: ...`); leave
+            `upsample` on to have the conditioner rewrite a plain sentence into it first.
+        canvas: One of the released canvases, as a `WIDTHxHEIGHT · ratio` label. The distillation's own operating
+            point is `1344x768 · 16:9 full`.
+        duration: Length in seconds, rounded up to the next frame count the video VAE can decode (`17 * n + 5`).
+        upsample: Rewrite the prompt into MiniMax-H3's trained format before encoding it.
+        seed: Random seed.
+
+    Returns:
+        The path of an mp4 holding h264 video and AAC audio, a one-line report of what ran, and the rewritten
+        prompt when there was one.
+    """
+    if LOAD_ERROR:
+        raise gr.Error(LOAD_ERROR)
+    if PIPE is None:
+        raise gr.Error("The denoiser is still loading.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("MiniMax-H3 always takes a prompt.")
+
+    from diffusers.utils import encode_video
+
+    num_frames = snap_frames(duration)
+
+    progress(0.0, desc=f"{'Rewriting the prompt' if upsample else 'Conditioning'} on {CONDITIONER_SPACE} ...")
+    conditioned = time.time()
+    prompt_embeds, text_token_tags, metadata, plan = encode_remote(
+        prompt, canvas, num_frames, rewrite_prompt=bool(upsample)
+    )
+    condition_seconds = time.time() - conditioned
+    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
+    refined = plan.get("refined_prompt") or ""
+
+    progress(0.2, desc=f"{NUM_FORWARDS} transformer forwards at {width}x{height}, {num_frames} frames ...")
+    started = time.time()
+    frames, audio, sampling_rate = _generate(prompt_embeds, text_token_tags, height, width, num_frames, seed)
+    generate_seconds = time.time() - started
+
+    directory = os.path.join(tempfile.gettempdir(), "fasth3-outputs")
+    os.makedirs(directory, exist_ok=True)
+    path = os.path.join(directory, f"fasth3-{int(time.time() * 1000)}.mp4")
+    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)
+
+    report = (
+        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {NUM_FORWARDS} transformer forwards · "
+        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
+        f"{', rewritten' if refined else ''}) · denoise + decode {generate_seconds:.0f}s "
+        f"({generate_seconds / NUM_FORWARDS:.1f} s/forward) · seed {int(seed)}"
+    )
+    print(f"[gen] {report}", flush=True)
+    return path, report, refined
+
+
+load_models()
+
+INTRO = f"""# FastH3 v1 (VSA) — MiniMax-H3 in 4 steps, sparse
+
+<div>
+  <a href="https://huggingface.co/{MODEL_REPO}" target="_blank" rel="noopener"><strong>[ model ]</strong></a> &nbsp;
+  <a href="https://github.com/hao-ai-lab/FastVideo" target="_blank" rel="noopener"><strong>[ FastVideo ]</strong></a> &nbsp;
+  <a href="https://huggingface.co/{BASE_REPO}" target="_blank" rel="noopener"><strong>[ base model ]</strong></a>
+</div>
+
+[`{MODEL_REPO}`](https://huggingface.co/{MODEL_REPO}) is a **data-free DMD2 distillation** of
+[MiniMax-H3](https://huggingface.co/{BASE_REPO}), the 33B dual-modality transformer that generates video **and** a
+fully synchronized soundtrack (ambience, foley, speech) in one denoising pass. The base model samples in 50 steps;
+this student walks a trained 4-step ladder — `t = 999, 749, 500, 250` — for **{NUM_FORWARDS} transformer forwards**
+per video.
+
+It is also distilled **under Video Sparse Attention**: 64-token tiles at 90% sparsity, with a trained per-head
+compression gate. This Space runs that sparse path, on FastVideo's own Triton kernels — not a dense substitute.
+"""
+
+FORMAT_NOTE = """MiniMax-H3 was trained on a structured prompt, not a caption:
+
+```text
+integrated_multimodal_description: [Shot 1] ... <d>[English] spoken line.</d> [Shot 2] At 00:04.500, ...
+overall_soundscape: ...
+non_diegetic_music: ...
+```
+
+**Expand prompt** (on by default) sends a plain sentence through the Qwen3-VL conditioner's own language-model head
+first, which writes that format with the same weights that are about to encode it. Turn it off when the prompt is
+already written out — as the last two examples below are. See the base model's
+[prompt writing guide](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/docs/VIDEO_PROMPT_WRITING_GUIDE_base_en.md).
+"""
+
+# Both long examples are the MiniMax-H3 authors' own published T2VA prompts: the first is Case 1 of the prompt writing
+# guide, the second is the reproducible 768p T2VA case from the model card
+# (`scripts/readme/reproducible-768p-t2va-request.sh`). Both are documentation of the Apache-2.0 base repo.
+GUIDE_CASE_1 = (
+    "integrated_multimodal_description: [Shot 1] Live-action, cinematic, a medium-wide shot frames a baker opening "
+    "the shutters of a small street bakery before sunrise. The camera pushes in with small amplitude at slow speed "
+    "as the middle-aged baker with a calm, slightly raspy voice (S1) places a fresh loaf on the wooden counter and "
+    "says: <d>[English] First batch of the morning.</d> [Shot 2] At 00:05.000, the camera cuts to a close-up of "
+    "steam rising from the sliced bread while the baker's final words carry over from the previous shot.\n\n"
+    "overall_soundscape: Wooden shutters scrape open over a quiet street as trays clink softly inside the bakery. "
+    "The doorbell rings once, followed by light footsteps and the crisp sound of bread being sliced.\n\n"
+    "non_diegetic_music: A soft acoustic-guitar pattern at a moderate tempo, joined by sparse upright-bass notes and "
+    "a gentle fade at the end."
+)
+OFFICIAL_T2VA = (
+    "integrated_multimodal_description: [Shot 1] Cinematic, medium wide shot, pushing in slowly. In the cavernous, "
+    "dimly lit bridge of a starship, sleek metallic consoles with glowing amber displays flank a massive, curved "
+    "observation window. A female captain, in her late 40s with an athletic build and short silver-streaked black "
+    "hair, stands in the center midground. She wears a structured, high-collared dark navy military tunic with "
+    "silver chest insignias. Her back is to the camera, silhouetted against the cool, ambient starlight pouring "
+    "through the thick glass. She stands perfectly still with her hands clasped tightly behind her back. Outside the "
+    "window, a massive armada of jagged, dark grey dreadnoughts hovers in tight formation against a deep purple "
+    "space nebula. The fleet's massive rear thrusters begin to glow with an intense, escalating bright blue light. "
+    "[Shot 2] At 00:04.500, the camera cuts to a close-up of the captain's face and shakes strongly. The brilliant "
+    "blue-white light from the fleet's gathering energy reflects vividly in her dark eyes. Suddenly, a blinding "
+    "white flash floods through the window, completely washing out the background as the fleet jumps to hyperspace. "
+    "The sheer spatial force violently jolts the bridge, causing the captain from Shot 1 to stagger slightly "
+    "forward, her shoulders tensing as she visibly braces herself against the physical tremors. As the intense "
+    "white light fades abruptly, leaving only the dim, empty expanse of the purple nebula reflected on her starkly "
+    "lit skin, her jaw clenches, and she slowly closes her eyes in the newly emptied space.\n"
+    "overall_soundscape: A low, resonant hum of the ship's ambient life support systems serves as the baseline, soon "
+    "drowned out by an audible, escalating, high-pitched electronic whine as the fleet outside charges its "
+    "hyperdrives. A massive, deafening, bass-heavy boom and sharp crackle erupts during the blinding flash, "
+    "accompanied by the loud metallic creaking, rattling, and deep thuds of the bridge's bulkheads vibrating under "
+    "immense physical stress. The intense roaring impact then cuts abruptly back to a hollow, echoing room tone, "
+    "leaving only the faint, steady hum of the isolated bridge.\n"
+    "non_diegetic_music: Cinematic space-opera orchestral score, slow tempo, featuring a solitary, mournful French "
+    "horn melody over deep, sustained string dissonances that build rapidly in volume and intensity, swelling to a "
+    "massive orchestral peak before snapping immediately into silence right after the jump."
+)
+
+CSS = """
+.main.fillable {max-width: 1250px !important}
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(title="FastH3 v1 (VSA)") as demo:
+    gr.Markdown(INTRO)
+
+    with gr.Row():
+        with gr.Column():
+            prompt = gr.Textbox(
+                label="Prompt",
+                lines=5,
+                placeholder="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
+                value="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
+            )
+            upsample = gr.Checkbox(
+                label="Expand prompt into MiniMax-H3's trained format",
+                value=True,
+                info="Runs on the conditioner Space before encoding. Turn off for a prompt already in that format.",
+            )
+            run = gr.Button("Generate", variant="primary")
+            with gr.Accordion("Advanced options", open=False):
+                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
+                duration = gr.Slider(
+                    label="Duration (s)",
+                    minimum=MIN_UI_DURATION,
+                    maximum=MAX_UI_DURATION,
+                    step=1,
+                    value=DEFAULT_DURATION,
+                )
+                seed = gr.Number(label="Seed", value=42, precision=0)
+                gr.Markdown(
+                    f"Steps are fixed at the trained ladder — a {SIGMA_GRID_POINTS}-point sigma grid, "
+                    f"{NUM_FORWARDS} transformer forwards, exactly what the checkpoint's own "
+                    "`fastvideo_inference.json` specifies. There is no guidance scale and no negative prompt: the "
+                    "base model is guidance-distilled."
+                )
+
+        with gr.Column():
+            video = gr.Video(label="Video + soundtrack")
+            report = gr.Markdown()
+            with gr.Accordion("Expanded prompt", open=False):
+                upsampled = gr.Textbox(show_label=False, lines=10, interactive=False)
+
+    with gr.Accordion("Prompt format", open=False):
+        gr.Markdown(FORMAT_NOTE)
+    banner = gr.Markdown()
+
+    gr.Examples(
+        examples=[
+            [
+                "A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
+                "960x544 · 16:9 fast",
+                5,
+                True,
+            ],
+            [
+                "A cellist playing a slow, low melody alone in an empty concert hall",
+                "1344x768 · 16:9 full",
+                5,
+                True,
+            ],
+            [GUIDE_CASE_1, "1344x768 · 16:9 full", 5, False],
+            [OFFICIAL_T2VA, "1344x768 · 16:9 full", 5, False],
+        ],
+        inputs=[prompt, canvas, duration, upsample],
+        outputs=[video, report, upsampled],
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
+        label="Examples — the last two are the MiniMax-H3 authors' own published T2VA prompts",
+    )
+
+    run.click(
+        generate,
+        [prompt, canvas, duration, upsample, seed],
+        [video, report, upsampled],
+        api_name="generate",
+    )
+    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.
+    diagnose = gr.Button(visible=False)
+    diagnose.click(selftest, None, gr.Markdown(visible=False), api_name="selftest")
+
+
+if __name__ == "__main__":
+    # Gradio 6 moved `theme` and `css` off the `Blocks` constructor onto `launch`.
+    demo.launch(theme=gr.themes.Citrus(), css=CSS, show_error=True, max_threads=1000, mcp_server=True)

h3_split_blocks.py

@@ -0,0 +1,51 @@
+"""The denoising half of a **split** MiniMax-H3 deployment.
+
+MiniMax-H3 is ~196 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage, so `MiniMaxH3Blocks` is cut at
+its `text_encoder` step: the 62.14 GiB Qwen3-VL conditioner runs in its own Space and everything else — the distilled
+transformer and the two autoencoders — runs here. `prompt_embeds` + `text_token_tags` is the whole wire format between
+the two halves.
+
+For a text-only (`t2va`) request the full `MiniMaxH3Blocks` sequence is `before_encode -> text_encoder -> vae_encoder
+-> denoise -> decode`, and the two `Auto*` encoder steps select nothing without a keyframe or a reference. What is
+left once the text encoder goes is exactly `MiniMaxH3CoreDenoiseStep -> MiniMaxH3DecodeStep`, and those two declare
+only `transformer`, the two schedulers, both autoencoders and `video_processor` — so `load_components` resolves those
+out of the checkpoint's `modular_model_index.json` and never fetches the conditioner or the `ref2va` partition.
+
+Adapted from `multimodalart/minimax-h3`, trimmed to the `t2va` (text -> video + audio) branch this Space serves.
+"""
+
+import torch
+from diffusers.modular_pipelines.minimax_h3.modular_blocks_minimax_h3 import (
+    MiniMaxH3CoreDenoiseStep,
+    MiniMaxH3DecodeStep,
+)
+from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
+
+
+class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
+    """The denoising half of a split MiniMax-H3: the `t2va` branch of `MiniMaxH3Blocks` without its text encoder."""
+
+    model_name = "minimax-h3"
+    block_classes = [MiniMaxH3CoreDenoiseStep, MiniMaxH3DecodeStep]
+    block_names = ["denoise", "decode"]
+
+    @property
+    def description(self):
+        return (
+            "The denoising half of a split MiniMax-H3 deployment: the `t2va` branch of `MiniMaxH3Blocks` without its "
+            "text-encoder step, so `prompt_embeds` and `text_token_tags` come in as inputs and the 62.14 GiB Qwen3-VL "
+            "conditioner is never loaded here."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("videos", description="The generated video."),
+            OutputParam(
+                "audio",
+                type_hint=torch.Tensor,
+                description="The generated soundtrack, of shape `(1, 2, num_samples)`.",
+            ),
+            OutputParam("sampling_rate", type_hint=int, description="Sample rate of the generated soundtrack in Hz."),
+        ]

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,13 @@
+# MiniMax-H3 landed in diffusers 0.40.0 (modular pipelines only). Pinned because `h3_split_blocks.py` subclasses its
+# block classes and `vsa_h3.py` patches `MiniMaxH3TransformerBlock` / `MiniMaxH3Attention` by name.
+diffusers==0.40.0
+# torch is deliberately unpinned: the ZeroGPU runtime preinstalls a supported build, and its bundled `pytorch-triton`
+# is what compiles the vendored VSA block-sparse kernels. torchvision is not needed anywhere in this Space.
+transformers
+accelerate
+# PyAV muxes the generated soundtrack onto the frames (`diffusers.utils.encode_video`).
+av
+pillow
+numpy
+safetensors>=0.8.0
+# huggingface-hub is deliberately absent: it is platform-managed, and diffusers already constrains it to <2.

vsa_h3.py

@@ -0,0 +1,420 @@
+"""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
+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,
+    )
+
+
+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.
+    """
+    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
+
+    return build_geometry(tuple(prefix_segments), (grid_t, grid_h, grid_w), token_tags.device, sparsity)
+
+
+# --------------------------------------------------------------------------------------------------------------
+# 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)
+    indices = scores[..., prefix:].topk(geometry.topk, dim=-1).indices + prefix
+    mask.scatter_(-1, indices, True)
+    mask[..., :prefix] = True
+    mask[:, :, :prefix, :] = True
+    return mask
+
+
+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)
+    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}
+
+
+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]:
+            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)
+        _ACTIVE["geometry"] = geometry
+        try:
+            return original_forward(*args, **kwargs)
+        finally:
+            _ACTIVE["geometry"] = None
+
+    transformer.forward = forward
+    transformer._vsa_forward_wrapped = True
+    return installed, gates

vsa_kernel/__init__.py

@@ -0,0 +1,16 @@
+"""Vendored FastVideo block-sparse attention Triton kernels.
+
+Copied verbatim from `fastvideo-kernel` (https://github.com/hao-ai-lab/FastVideo,
+`fastvideo-kernel/python/fastvideo_kernel/triton_kernels/`), Apache License 2.0.
+
+Only the two pure-Triton modules are vendored: the package's default route is a
+CUDA extension that has to be compiled per architecture (and whose fastest entry,
+`block_sparse_attn_sm100a`, is GB200-only), while FastVideo's own
+`--vsa-kernel triton` route -- these two files -- runs anywhere Triton does and
+computes exactly the same masked attention.
+"""
+
+from .block_sparse_attn_triton import triton_block_sparse_attn_forward
+from .index import map_to_index
+
+__all__ = ["triton_block_sparse_attn_forward", "map_to_index"]

vsa_kernel/block_sparse_attn_triton.py

@@ -0,0 +1,879 @@
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+
+Credits: OpenAI kernel team
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+import math  # small utility needed by the sparse wrapper
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+# BLOCK_M / BLOCK_N are fixed at 64 because they are structural, not tunable:
+# the kernel indexes the top-k list per BLOCK_M q-tile and addresses keys as
+# kv_idx * BLOCK_N, so both must match the granularity q2k_index and
+# variable_block_sizes were built at.
+#
+# num_stages / num_warps ARE free, and the previous {3, 4, 7} was inherited from
+# the upstream tutorial rather than tuned here. It skips 5 and 6; on Blackwell
+# (sm_121) the optimum is num_stages=5, so the search could not reach it. Both
+# block paths independently select 5 once it is available. Autotune still picks
+# per architecture, so other GPUs re-tune rather than inheriting this choice.
+#
+# VENDORING NOTE (the only edit made to this file): Triton re-benchmarks the sweep below on every new `N_CTX_Q` --
+# i.e. on every canvas and duration a visitor picks, and again in each fresh ZeroGPU worker -- and at these sequence
+# lengths one bench round costs seconds of the visitor's GPU quota. Triton skips benchmarking altogether when a
+# single config is offered, so this pins the point the comment above says the search lands on for Blackwell.
+# `H3_VSA_AUTOTUNE=1` restores the upstream sweep.
+import os as _os
+
+if _os.environ.get("H3_VSA_AUTOTUNE", "0") == "1":
+    configs = [
+        triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \
+        for BM in [64]\
+        for BN in [64]\
+        for s in [2, 3, 4, 5, 6, 7]\
+        for w in [4, 8]\
+    ]
+else:
+    configs = [
+        triton.Config({
+            'BLOCK_M': 64,
+            'BLOCK_N': 64
+        },
+                      num_stages=int(_os.environ.get("H3_VSA_STAGES", "5")),
+                      num_warps=int(_os.environ.get("H3_VSA_WARPS", "4")))
+    ]
+
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+@triton.autotune(configs, key=["N_CTX_Q", "HEAD_DIM"])
+@triton.jit
+def _attn_fwd_sparse(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        q2k_index,
+        q2k_num,
+        max_kv_blks,  #
+        variable_block_sizes,
+        M,
+        Out,  #
+        stride_qz,
+        stride_qh,
+        stride_qm,
+        stride_qk,
+        stride_kz,
+        stride_kh,
+        stride_kn,
+        stride_kk,
+        stride_vz,
+        stride_vh,
+        stride_vk,
+        stride_vn,
+        stride_oz,
+        stride_oh,
+        stride_om,
+        stride_on,
+        Z,
+        H,
+        N_CTX_Q,  #
+        N_CTX_KV,  #
+        HEAD_DIM: tl.constexpr,  #
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+        STAGE: tl.constexpr):
+    """
+    64×64 **block-sparse** forward kernel. Back-prop kernels remain dense
+    (32×64 and 64×32) – memory footprint unchanged.
+    """
+
+    # ----- program-id mapping -----
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(1)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    q_tiles = N_CTX_Q // BLOCK_M
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+
+    # ----- base pointers -----
+    # Note: when q and kv have different sequence lengths, their per-(batch,head)
+    # strides differ, so we must compute separate base offsets.
+    q_off = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    k_off = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    v_off = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    o_off = (b.to(tl.int64) * stride_oz + h.to(tl.int64) * stride_oh)
+
+    Q_ptr = tl.make_block_ptr(base=Q + q_off,
+                              shape=(N_CTX_Q, HEAD_DIM),
+                              strides=(stride_qm, stride_qk),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
+
+    K_base = tl.make_block_ptr(base=K + k_off,
+                               shape=(HEAD_DIM, N_CTX_KV),
+                               strides=(stride_kk, stride_kn),
+                               offsets=(0, 0),
+                               block_shape=(HEAD_DIM, BLOCK_N),
+                               order=(0, 1))
+
+    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)
+    V_base = tl.make_block_ptr(base=V + v_off,
+                               shape=(N_CTX_KV, HEAD_DIM),
+                               strides=(stride_vk, stride_vn),
+                               offsets=(0, 0),
+                               block_shape=(BLOCK_N, HEAD_DIM),
+                               order=v_order)
+
+    O_ptr = tl.make_block_ptr(base=Out + o_off,
+                              shape=(N_CTX_Q, HEAD_DIM),
+                              strides=(stride_om, stride_on),
+                              offsets=(q_blk * BLOCK_M, 0),
+                              block_shape=(BLOCK_M, HEAD_DIM),
+                              order=(1, 0))
+
+    # ----- accumulators -----
+    offs_m = q_blk * BLOCK_M + tl.arange(0, BLOCK_M)
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
+    qk_scale = sm_scale * 1.44269504  # 1/ln2
+    q = tl.load(Q_ptr)
+
+    # ----- sparse loop over valid K/V tiles -----
+    for i in range(0, kv_blocks):
+        kv_idx = tl.load(kv_ptr + i).to(tl.int32)
+        block_size = tl.load(variable_block_sizes + kv_idx)
+        K_ptr = tl.advance(K_base, (0, kv_idx * BLOCK_N))
+        V_ptr = tl.advance(V_base, (kv_idx * BLOCK_N, 0))
+
+        k = tl.load(K_ptr)
+        qk = tl.dot(q, k)
+        # mask out invalid columns
+        mask = tl.arange(0, BLOCK_N) < block_size
+        qk = tl.where(mask[None, :], qk, -float("inf"))
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale)
+        p = tl.math.exp2(qk * qk_scale - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+
+        alpha = tl.math.exp2(m_i - m_ij)
+        l_i = l_i * alpha + l_ij
+        acc = acc * alpha[:, None]
+
+        v = tl.load(V_ptr)
+        acc = tl.dot(p.to(tl.bfloat16), v, acc)
+        m_i = m_ij
+
+    # ----- epilogue -----
+    m_i += tl.math.log2(l_i)
+    acc = acc / l_i[:, None]
+    tl.store(M + off_hz * N_CTX_Q + offs_m, m_i)
+    tl.store(O_ptr, acc.to(Out.type.element_ty))
+
+
+# ──────────────────────────── SPARSE ADDITION END ─────────────────────────────
+
+
+@triton.jit
+def _attn_bwd_preprocess(
+        O,
+        DO,  #
+        Delta,  #
+        Z,
+        H,
+        N_CTX,  #
+        BLOCK_M: tl.constexpr,
+        HEAD_DIM: tl.constexpr  #
+):
+    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_hz = tl.program_id(1)
+    off_n = tl.arange(0, HEAD_DIM)
+    # load
+    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
+    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
+    delta = tl.sum(o * do, axis=1)
+    # write-back
+    tl.store(Delta + off_hz * N_CTX + off_m, delta)
+
+
+# The main inner-loop logic for computing dK and dV.
+@triton.jit
+def _attn_bwd_dkdv(
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        sm_scale,  #
+        DO,  #
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX_KV,
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,  #
+        # Filled in by the wrapper.
+    start_n,
+        start_m,
+        num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M1)
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+    offs_k = tl.arange(0, HEAD_DIM)
+    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
+    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
+    step_m = BLOCK_M1
+    kv_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    kv_tiles = N_CTX_KV // BLOCK_N1
+    meta_base = ((b * H + h) * kv_tiles + kv_blk)
+
+    q_blocks = tl.load(k2q_num + meta_base)  # int32
+    q_ptr = k2q_index + meta_base * max_q_blks  # ptr to list
+    block_size = tl.load(variable_block_sizes + kv_blk)
+
+    for blk_idx in range(q_blocks * 2):
+        block_sparse_offset = (tl.load(q_ptr + blk_idx // 2).to(tl.int32) * 2 + blk_idx % 2) * step_m
+        qT = tl.load(qT_ptrs + block_sparse_offset * stride_tok)
+        # Load m before computing qk to reduce pipeline stall.
+        offs_m = start_m + block_sparse_offset + tl.arange(0, BLOCK_M1)
+        m = tl.load(M + offs_m)
+        # Recompute logits exactly as the forward does: raw bf16 operands into
+        # the dot, fp32 scale after accumulation. A bf16 pre-scaled K perturbs
+        # the recomputed logits relative to the saved M by an error
+        # proportional to |logit|, which exp2 amplifies into arbitrarily wrong
+        # probabilities at large activations.
+        qkT = tl.dot(k, qT) * (sm_scale * 1.4426950408889634)
+        pT = tl.math.exp2(qkT - m[None, :])
+        mask = tl.arange(0, BLOCK_N1) < block_size
+        pT = tl.where(mask[:, None], pT, 0.0)
+
+        do = tl.load(do_ptrs + block_sparse_offset * stride_tok)
+        # Compute dV.
+        ppT = pT
+        ppT = ppT.to(tl.bfloat16)
+        dv += tl.dot(ppT, do)
+        # D (= delta) is pre-divided by ds_scale.
+        Di = tl.load(D + offs_m)
+        # Compute dP and dS.
+        dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
+        dsT = pT * (dpT - Di[None, :])
+        dsT = dsT.to(tl.bfloat16)
+        dk += tl.dot(dsT, tl.trans(qT))
+        # Increment pointers.
+    return dk, dv
+
+
+# the main inner-loop logic for computing dQ
+@triton.jit
+def _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        do,
+        m,
+        D,
+        sm_scale,
+        # shared by Q/K/V/DO.
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr,
+        # Filled in by the wrapper.
+        start_m,
+        start_n,
+        num_steps):
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+    offs_n = start_n + tl.arange(0, BLOCK_N2)
+    offs_k = tl.arange(0, HEAD_DIM)
+    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
+    # D (= delta) is pre-divided by ds_scale.
+    Di = tl.load(D + offs_m)
+    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
+    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
+    step_n = BLOCK_N2
+
+    q_blk = tl.program_id(0)  # Q-tile index
+    off_hz = tl.program_id(2)  # fused (batch, head)
+    b = off_hz // H
+    h = off_hz % H
+    q_tiles = N_CTX // BLOCK_M2
+    meta_base = ((b * H + h) * q_tiles + q_blk)
+
+    kv_blocks = tl.load(q2k_num + meta_base)  # int32
+    kv_ptr = q2k_index + meta_base * max_kv_blks  # ptr to list
+
+    for blk_idx in range(kv_blocks * 2):
+        kv_idx = tl.load(kv_ptr + blk_idx // 2).to(tl.int32)
+        # variable_block_sizes is defined per KV block (tile). Mask must therefore
+        # use kv_idx (not q_blk). Also, because we split each 64-token block into
+        # two 32-token halves, the mask must account for the half-block offset.
+        block_size = tl.load(variable_block_sizes + kv_idx).to(tl.int32)
+        half = (blk_idx % 2).to(tl.int32)
+        block_sparse_offset = (kv_idx * 2 + half) * step_n * stride_tok
+        kT = tl.load(kT_ptrs + block_sparse_offset)
+        vT = tl.load(vT_ptrs + block_sparse_offset)
+        qk = tl.dot(q, kT) * (sm_scale * 1.4426950408889634)
+        p = tl.math.exp2(qk - m)
+        offs_in_block = half * step_n + tl.arange(0, BLOCK_N2)
+        mask = offs_in_block < block_size
+        p = tl.where(mask[None, :], p, 0.0)
+        # Compute dP and dS.
+        dp = tl.dot(do, vT).to(tl.float32)
+        ds = p * (dp - Di[:, None])
+        ds = ds.to(tl.bfloat16)
+        # Compute dQ (kT is raw; the caller applies sm_scale once at the end).
+        dq += tl.dot(ds, tl.trans(kT))
+        # Increment pointers.
+    return dq
+
+
+@triton.jit
+def _attn_bwd(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        DO,  #
+        DQ,
+        DK,
+        DV,  #
+        M,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared by Q/K/V/DO.
+        stride_z,
+        stride_h,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr):
+    LN2 = 0.6931471824645996  # = ln(2)
+
+    bhid = tl.program_id(2)
+    off_chz = (bhid * N_CTX).to(tl.int64)
+    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
+    pid = tl.program_id(0)
+
+    # offset pointers for batch/head
+    Q += adj
+    K += adj
+    V += adj
+    DO += adj
+    DQ += adj
+    DK += adj
+    DV += adj
+    M += off_chz
+    D += off_chz
+
+    # load scales
+    offs_k = tl.arange(0, HEAD_DIM)
+
+    start_n = pid * BLOCK_N1
+    start_m = 0
+
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    num_steps = N_CTX // BLOCK_M1
+
+    dk, dv = _attn_bwd_dkdv(  #
+        dk,
+        dv,  #
+        Q,
+        k,
+        v,
+        sm_scale,  #
+        DO,  #
+        M,
+        D,  #
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M1,
+        BLOCK_N1,
+        HEAD_DIM,  #
+        start_n,
+        start_m,
+        num_steps  #
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv)
+
+    # Write back dK.
+    dk *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk)
+
+    # THIS BLOCK DOES DQ:
+    start_m = pid * BLOCK_M2
+    end_n = 0
+
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    m = tl.load(M + offs_m)
+    m = m[:, None]
+
+    num_steps = N_CTX // BLOCK_N2
+    dq = _attn_bwd_dq(
+        dq,
+        q,
+        K,
+        V,  #
+        do,
+        m,
+        D,  #
+        sm_scale,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,  #
+        H,
+        N_CTX,  #
+        BLOCK_M2,
+        BLOCK_N2,
+        HEAD_DIM,  #
+        start_m,
+        end_n,
+        num_steps  #
+    )
+    # Write back dQ.
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq *= sm_scale
+    tl.store(dq_ptrs, dq)
+
+
+@triton.jit
+def _attn_bwd_dkdv_kernel(
+        Q,
+        K,
+        V,
+        sm_scale,  #
+        DO,  #
+        DK,
+        DV,  #
+        M,
+        D,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        # shared token/dim strides (assumed contiguous along token and dim)
+        stride_tok,
+        stride_d,  #
+        # batch/head strides (may differ between Q and KV)
+    stride_qz,
+        stride_qh,
+        stride_kz,
+        stride_kh,
+        stride_vz,
+        stride_vh,
+        stride_doz,
+        stride_doh,
+        stride_dkz,
+        stride_dkh,
+        stride_dvz,
+        stride_dvh,
+        H,
+        N_CTX_Q,
+        N_CTX_KV,
+        BLOCK_M1: tl.constexpr,  #
+        BLOCK_N1: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr):
+    """
+    Backward kernel that computes dK and dV for each KV block (64 tokens).
+    Grid:
+      pid0: kv_blk in [0, N_CTX_KV/BLOCK_N1)
+      pid2: fused (batch, head) in [0, B*H)
+    """
+    bhid = tl.program_id(2)
+    b = bhid // H
+    h = bhid % H
+    kv_blk = tl.program_id(0)
+
+    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
+    dk_adj = (b.to(tl.int64) * stride_dkz + h.to(tl.int64) * stride_dkh)
+    dv_adj = (b.to(tl.int64) * stride_dvz + h.to(tl.int64) * stride_dvh)
+
+    Q = Q + q_adj
+    K = K + kv_adj_k
+    V = V + kv_adj_v
+    DO = DO + do_adj
+    DK = DK + dk_adj
+    DV = DV + dv_adj
+
+    # M and D (delta) are always sized by Q length.
+    M = M + (bhid * N_CTX_Q).to(tl.int64)
+    D = D + (bhid * N_CTX_Q).to(tl.int64)
+
+    offs_k = tl.arange(0, HEAD_DIM)
+    start_n = kv_blk * BLOCK_N1
+    offs_n = start_n + tl.arange(0, BLOCK_N1)
+
+    # load K and V: they stay in SRAM throughout the inner loop.
+    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
+
+    dv_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+    dk_acc = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
+
+    num_steps = N_CTX_Q // BLOCK_M1
+    dk_acc, dv_acc = _attn_bwd_dkdv(
+        dk_acc,
+        dv_acc,
+        Q,
+        k,
+        v,
+        sm_scale,
+        DO,
+        M,
+        D,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,
+        H,
+        N_CTX_KV,
+        BLOCK_M1=BLOCK_M1,
+        BLOCK_N1=BLOCK_N1,
+        HEAD_DIM=HEAD_DIM,
+        start_n=start_n,
+        start_m=0,
+        num_steps=num_steps,
+    )
+
+    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dv_ptrs, dv_acc)
+
+    dk_acc *= sm_scale
+    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
+    tl.store(dk_ptrs, dk_acc)
+
+
+@triton.jit
+def _attn_bwd_dq_kernel(
+        Q,
+        K,
+        V,
+        sm_scale,
+        DO,  #
+        DQ,
+        M,
+        D,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        # shared token/dim strides (assumed contiguous along token and dim)
+        stride_tok,
+        stride_d,  #
+        # batch/head strides (may differ between Q and KV)
+    stride_qz,
+        stride_qh,
+        stride_kz,
+        stride_kh,
+        stride_vz,
+        stride_vh,
+        stride_doz,
+        stride_doh,
+        stride_dqz,
+        stride_dqh,
+        H,
+        N_CTX_Q,
+        BLOCK_M2: tl.constexpr,  #
+        BLOCK_N2: tl.constexpr,  #
+        HEAD_DIM: tl.constexpr):
+    """
+    Backward kernel that computes dQ for each Q block (64 tokens).
+    Grid:
+      pid0: q_blk in [0, N_CTX_Q/BLOCK_M2)
+      pid2: fused (batch, head) in [0, B*H)
+    """
+    LN2 = 0.6931471824645996  # = ln(2)
+    bhid = tl.program_id(2)
+    b = bhid // H
+    h = bhid % H
+    q_blk = tl.program_id(0)
+
+    q_adj = (b.to(tl.int64) * stride_qz + h.to(tl.int64) * stride_qh)
+    kv_adj_k = (b.to(tl.int64) * stride_kz + h.to(tl.int64) * stride_kh)
+    kv_adj_v = (b.to(tl.int64) * stride_vz + h.to(tl.int64) * stride_vh)
+    do_adj = (b.to(tl.int64) * stride_doz + h.to(tl.int64) * stride_doh)
+    dq_adj = (b.to(tl.int64) * stride_dqz + h.to(tl.int64) * stride_dqh)
+
+    Q = Q + q_adj
+    K = K + kv_adj_k
+    V = V + kv_adj_v
+    DO = DO + do_adj
+    DQ = DQ + dq_adj
+
+    M = M + (bhid * N_CTX_Q).to(tl.int64)
+    D = D + (bhid * N_CTX_Q).to(tl.int64)
+
+    offs_k = tl.arange(0, HEAD_DIM)
+    start_m = q_blk * BLOCK_M2
+    offs_m = start_m + tl.arange(0, BLOCK_M2)
+
+    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
+    m = tl.load(M + offs_m)[:, None]
+
+    dq_acc = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
+    num_steps = 0  # unused in _attn_bwd_dq
+    dq_acc = _attn_bwd_dq(
+        dq_acc,
+        q,
+        K,
+        V,
+        do,
+        m,
+        D,
+        sm_scale,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        stride_tok,
+        stride_d,
+        H,
+        N_CTX_Q,
+        BLOCK_M2=BLOCK_M2,
+        BLOCK_N2=BLOCK_N2,
+        HEAD_DIM=HEAD_DIM,
+        start_m=start_m,
+        start_n=0,
+        num_steps=num_steps,
+    )
+
+    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
+    dq_acc *= sm_scale
+    tl.store(dq_ptrs, dq_acc)
+
+
+# ──────────────────────────── SPARSE ADDITION BEGIN ───────────────────────────
+def triton_block_sparse_attn_forward(q, k, v, q2k_index, q2k_num, variable_block_sizes):
+    B, H, Tq, D = q.shape
+    Tkv = k.shape[2]
+    sm_scale = 1.0 / math.sqrt(D)
+    max_kv_blks = q2k_index.shape[-1]
+    assert Tq % 64 == 0, f"q length must be a multiple of 64, but got {Tq}"
+    assert Tkv % 64 == 0, f"kv length must be a multiple of 64, but got {Tkv}"
+    assert q2k_num.shape[
+        -1] == Tq // 64, f"shape mismatch, Tq // 64 = {Tq // 64}, q2k_num.shape[-2] = {q2k_num.shape[-2]}"
+    assert variable_block_sizes.numel() == Tkv // 64, (
+        f"shape mismatch, variable_block_sizes must have length {Tkv // 64}, "
+        f"got {variable_block_sizes.numel()}")
+    o = torch.empty_like(q)
+    M = torch.empty((B, H, Tq), dtype=torch.float32, device=q.device)
+
+    grid = lambda _: (triton.cdiv(Tq, 64), B * H, 1)
+    _attn_fwd_sparse[grid](q,
+                           k,
+                           v,
+                           sm_scale,
+                           q2k_index,
+                           q2k_num,
+                           max_kv_blks,
+                           variable_block_sizes,
+                           M,
+                           o,
+                           q.stride(0),
+                           q.stride(1),
+                           q.stride(2),
+                           q.stride(3),
+                           k.stride(0),
+                           k.stride(1),
+                           k.stride(2),
+                           k.stride(3),
+                           v.stride(0),
+                           v.stride(1),
+                           v.stride(2),
+                           v.stride(3),
+                           o.stride(0),
+                           o.stride(1),
+                           o.stride(2),
+                           o.stride(3),
+                           B,
+                           H,
+                           Tq,
+                           Tkv,
+                           HEAD_DIM=D,
+                           STAGE=3)
+
+    return o, M
+
+
+def triton_block_sparse_attn_backward(do, q, k, v, o, M, q2k_index, q2k_num, k2q_index, k2q_num, variable_block_sizes):
+    assert do.is_contiguous()
+
+    B, H, Tq, D = q.shape
+    Tkv = k.shape[2]
+    sm_scale = 1.0 / math.sqrt(D)
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    BATCH, N_HEAD = q.shape[:2]
+    BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 64, 64, 32
+    # K stays raw: the backward kernels apply sm_scale in fp32 after the dot,
+    # matching the forward's rounding exactly. (A bf16 pre-scaled K perturbs
+    # the recomputed logits vs the saved M; exp2 turns that into unboundedly
+    # wrong probabilities at large activations.)
+    arg_k = k
+    PRE_BLOCK = 64
+    assert Tq % PRE_BLOCK == 0
+    pre_grid = (Tq // PRE_BLOCK, BATCH * N_HEAD)
+    delta = torch.empty_like(M)
+    _attn_bwd_preprocess[pre_grid](
+        o,
+        do,  #
+        delta,  #
+        BATCH,
+        N_HEAD,
+        Tq,  #
+        BLOCK_M=PRE_BLOCK,
+        HEAD_DIM=D  #
+    )
+
+    max_q_blks = k2q_index.shape[-1]
+    max_kv_blks = q2k_index.shape[-1]
+
+    # dK/dV kernel: grid over KV blocks
+    grid_kv = (Tkv // BLOCK_N1, 1, BATCH * N_HEAD)
+    _attn_bwd_dkdv_kernel[grid_kv](
+        q,
+        arg_k,
+        v,
+        sm_scale,
+        do,
+        dk,
+        dv,
+        M,
+        delta,
+        k2q_index,
+        k2q_num,
+        max_q_blks,
+        variable_block_sizes,
+        q.stride(2),
+        q.stride(3),
+        q.stride(0),
+        q.stride(1),
+        arg_k.stride(0),
+        arg_k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        do.stride(0),
+        do.stride(1),
+        dk.stride(0),
+        dk.stride(1),
+        dv.stride(0),
+        dv.stride(1),
+        N_HEAD,
+        Tq,
+        Tkv,
+        BLOCK_M1=BLOCK_M1,
+        BLOCK_N1=BLOCK_N1,
+        HEAD_DIM=D,
+    )
+
+    # dQ kernel: grid over Q blocks
+    grid_q = (Tq // BLOCK_M2, 1, BATCH * N_HEAD)
+    _attn_bwd_dq_kernel[grid_q](
+        q,
+        arg_k,
+        v,
+        sm_scale,
+        do,
+        dq,
+        M,
+        delta,
+        q2k_index,
+        q2k_num,
+        max_kv_blks,
+        variable_block_sizes,
+        q.stride(2),
+        q.stride(3),
+        q.stride(0),
+        q.stride(1),
+        arg_k.stride(0),
+        arg_k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        do.stride(0),
+        do.stride(1),
+        dq.stride(0),
+        dq.stride(1),
+        N_HEAD,
+        Tq,
+        BLOCK_M2=BLOCK_M2,
+        BLOCK_N2=BLOCK_N2,
+        HEAD_DIM=D,
+    )
+
+    return dq, dk, dv

vsa_kernel/index.py

@@ -0,0 +1,250 @@
+## pytorch sdpa version of block sparse ##
+from typing import Tuple
+
+import triton
+import triton.language as tl
+import torch
+
+
+@triton.jit
+def topk_index_to_map_kernel(
+    map_ptr,
+    index_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    topk,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    for i in tl.static_range(topk):
+        index = tl.load(index_ptr_base + i * index_kv_stride)
+        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
+
+
+@triton.jit
+def map_to_index_kernel(
+    map_ptr,
+    index_ptr,
+    index_num_ptr,
+    map_bs_stride,
+    map_h_stride,
+    map_q_stride,
+    map_kv_stride,
+    index_bs_stride,
+    index_h_stride,
+    index_q_stride,
+    index_kv_stride,
+    index_num_bs_stride,
+    index_num_h_stride,
+    index_num_q_stride,
+    num_kv_blocks,
+):
+    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
+    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
+
+    num = 0
+    for i in tl.range(num_kv_blocks):
+        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
+        if map_entry:
+            tl.store(index_ptr_base + num * index_kv_stride, i)
+            num += 1
+
+    tl.store(index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride + q * index_num_q_stride, num)
+
+
+def topk_index_to_map(index: torch.Tensor, num_kv_blocks: int, transpose_map: bool = False):
+    """
+    Convert topk indices to a map.
+    
+    Args:
+        index: [bs, h, num_q_blocks, topk]
+            The topk indices tensor.
+        num_kv_blocks: int
+            The number of key-value blocks in the block_map returned
+        transpose_map: bool
+            If True, the block_map will be transposed on the final two dimensions.
+    
+    Returns:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            A binary map where 1 indicates that the q block attends to the kv block.
+    """
+    bs, h, num_q_blocks, topk = index.shape
+
+    if transpose_map is False:
+        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks), dtype=torch.bool, device=index.device)
+    else:
+        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks), dtype=torch.bool, device=index.device)
+        block_map = block_map.transpose(2, 3)
+
+    grid = (bs, h, num_q_blocks)
+    topk_index_to_map_kernel[grid](
+        block_map,
+        index,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        topk=topk,
+    )
+
+    return block_map
+
+
+def map_to_index(block_map: torch.Tensor):
+    """
+    Convert a block map to indices and counts.
+    
+    Args:
+        block_map: [bs, h, num_q_blocks, num_kv_blocks]
+            The block map tensor.
+    
+    Returns:
+        index: [bs, h, num_q_blocks, num_kv_blocks]
+            The indices of the blocks.
+        index_num: [bs, h, num_q_blocks]
+            The number of blocks for each q block.
+    """
+    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
+
+    index = torch.full((block_map.shape), -1, dtype=torch.int32, device=block_map.device)
+    index_num = torch.empty((bs, h, num_q_blocks), dtype=torch.int32, device=block_map.device)
+
+    grid = (bs, h, num_q_blocks)
+    map_to_index_kernel[grid](
+        block_map,
+        index,
+        index_num,
+        block_map.stride(0),
+        block_map.stride(1),
+        block_map.stride(2),
+        block_map.stride(3),
+        index.stride(0),
+        index.stride(1),
+        index.stride(2),
+        index.stride(3),
+        index_num.stride(0),
+        index_num.stride(1),
+        index_num.stride(2),
+        num_kv_blocks=num_kv_blocks,
+    )
+
+    return index, index_num
+
+
+@triton.jit
+def _invert_indices_kernel(
+    q2k_idx_ptr,
+    q2k_num_ptr,
+    k2q_idx_ptr,
+    k2q_num_ptr,
+    q2k_idx_b,
+    q2k_idx_h,
+    q2k_idx_q,
+    q2k_idx_k,
+    q2k_num_b,
+    q2k_num_h,
+    q2k_num_q,
+    k2q_idx_b,
+    k2q_idx_h,
+    k2q_idx_k,
+    k2q_idx_q,
+    k2q_num_b,
+    k2q_num_h,
+    k2q_num_k,
+    MAX_KV_PER_Q: tl.constexpr,
+):
+    # One program per (b, h, q): reserve a slot in k2q via atomicAdd, write q.
+    pid_b = tl.program_id(0)
+    pid_h = tl.program_id(1)
+    pid_q = tl.program_id(2)
+
+    n = tl.load(q2k_num_ptr + pid_b * q2k_num_b + pid_h * q2k_num_h + pid_q * q2k_num_q)
+
+    q2k_row = (q2k_idx_ptr + pid_b * q2k_idx_b + pid_h * q2k_idx_h + pid_q * q2k_idx_q)
+
+    for i in tl.range(0, MAX_KV_PER_Q):
+        if i < n:
+            kv = tl.load(q2k_row + i * q2k_idx_k)
+            count_ptr = (k2q_num_ptr + pid_b * k2q_num_b + pid_h * k2q_num_h + kv * k2q_num_k)
+            pos = tl.atomic_add(count_ptr, 1)
+            tl.store(
+                k2q_idx_ptr + pid_b * k2q_idx_b + pid_h * k2q_idx_h + kv * k2q_idx_k + pos * k2q_idx_q,
+                pid_q,
+            )
+
+
+def invert_indices(
+    q2k_idx: torch.Tensor,
+    q2k_num: torch.Tensor,
+    num_kv_blocks: int,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Transpose a Q->KV index list into a K->Q one via atomic compaction (GPU)."""
+    if q2k_idx.dim() != 4:
+        raise ValueError(f"q2k_idx must be [B, H, Nq, Mk], got shape={tuple(q2k_idx.shape)}")
+    if q2k_num.dim() != 3:
+        raise ValueError(f"q2k_num must be [B, H, Nq], got shape={tuple(q2k_num.shape)}")
+    if not q2k_idx.is_cuda or not q2k_num.is_cuda:
+        raise RuntimeError("invert_indices requires CUDA tensors.")
+
+    B, H, Nq, Mk = q2k_idx.shape
+    if q2k_num.shape != (B, H, Nq):
+        raise ValueError(f"q2k_num shape {tuple(q2k_num.shape)} does not match q2k_idx "
+                         f"[B, H, Nq] = {(B, H, Nq)}")
+
+    q2k_idx = q2k_idx.contiguous()
+    q2k_num = q2k_num.contiguous()
+    if q2k_idx.dtype != torch.int32:
+        q2k_idx = q2k_idx.to(torch.int32)
+    if q2k_num.dtype != torch.int32:
+        q2k_num = q2k_num.to(torch.int32)
+
+    # Any KV block is attended by at most Nq Q blocks (one per Q row), so
+    # `Nq` is a tight upper bound on the compacted K->Q slots.
+    k2q_idx = torch.empty(
+        (B, H, num_kv_blocks, Nq),
+        dtype=torch.int32,
+        device=q2k_idx.device,
+    )
+    k2q_num = torch.zeros(
+        (B, H, num_kv_blocks),
+        dtype=torch.int32,
+        device=q2k_idx.device,
+    )
+
+    grid = (B, H, Nq)
+    _invert_indices_kernel[grid](
+        q2k_idx,
+        q2k_num,
+        k2q_idx,
+        k2q_num,
+        q2k_idx.stride(0),
+        q2k_idx.stride(1),
+        q2k_idx.stride(2),
+        q2k_idx.stride(3),
+        q2k_num.stride(0),
+        q2k_num.stride(1),
+        q2k_num.stride(2),
+        k2q_idx.stride(0),
+        k2q_idx.stride(1),
+        k2q_idx.stride(2),
+        k2q_idx.stride(3),
+        k2q_num.stride(0),
+        k2q_num.stride(1),
+        k2q_num.stride(2),
+        MAX_KV_PER_Q=Mk,
+    )
+
+    return k2q_idx, k2q_num