[feat]: build VSA sparse-attention metadata in the Regenerate-2K refine loop

fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py

@@ -14,6 +14,10 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout, build_row_timesteps
+from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import (
+    _h3_vsa_metadata_builder,
+    _h3_vsa_single_region_segments,
+)
 from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
 from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import refine_sigmas
 from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches
@@ -62,18 +66,38 @@ class Regenerate2KDenoisingStage(PipelineStage):
         device = get_local_torch_device()
         sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma)
         state = batch.latents.to(device)
+        vsa = self._vsa_plan(batch, layout, transformer, fastvideo_args, device)
         with self._attention_kernel(device):
-            state = self._euler_loop(batch, layout, state, sigmas, device, transformer)
+            state = self._euler_loop(batch, layout, state, sigmas, device, transformer, vsa)
         batch.latents = state
         return batch
 
+    @staticmethod
+    def _vsa_plan(batch: ForwardBatch, layout: MiniMaxH3PackedLayout, transformer: Any, fastvideo_args: FastVideoArgs,
+                  device: torch.device) -> tuple[Any, ...] | None:
+        """Per-step VSA metadata inputs, or None for a dense backend.
+
+        The builder and the packed segments come from MiniMax-H3's own denoising stage so the refine sees the same
+        sparse layout its t2v runs do; the refine layout is a single generated-video region (no reference rows).
+        """
+        builder = _h3_vsa_metadata_builder(transformer, fastvideo_args)
+        if builder is None:
+            return None
+        patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
+        sparsity = fastvideo_args.VSA_sparsity
+        if sparsity is None:
+            sparsity = getattr(batch, "VSA_sparsity", 0.0)
+        tile_size = fastvideo_args.VSA_tile_size
+        return (builder, patch_size, _h3_vsa_single_region_segments(layout,
+                                                                    patch_size), float(sparsity), tile_size, device)
+
     @staticmethod
     def _attention_kernel(device: torch.device) -> contextlib.AbstractContextManager[None]:
         """cuDNN attention: at 133k tokens PyTorch's default SDPA dispatch picks its flash kernel, 3x slower on B200."""
         return sdpa_kernel(SDPBackend.CUDNN_ATTENTION) if device.type == "cuda" else contextlib.nullcontext()
 
     def _euler_loop(self, batch: ForwardBatch, layout: MiniMaxH3PackedLayout, state: torch.Tensor, sigmas: list[float],
-                    device: torch.device, transformer: Any) -> torch.Tensor:
+                    device: torch.device, transformer: Any, vsa: Any) -> torch.Tensor:
         position_ids = layout.position_ids.to(device)
         token_tags = layout.token_tags.to(device)
         video_indices = layout.video_indices.to(device)
@@ -89,8 +113,19 @@ class Regenerate2KDenoisingStage(PipelineStage):
                                                   condition_audio_timestep=1.0)
             unique, inverse = unique.to(device), inverse.to(device)
             self.curve.set_timesteps(unique)
+            attn_metadata = None
+            if vsa is not None:
+                builder, patch_size, segments, sparsity, tile_size, vsa_device = vsa
+                attn_metadata = builder.build(current_timestep=index,
+                                              patch_size=patch_size,
+                                              VSA_sparsity=sparsity,
+                                              packed_segments=segments,
+                                              device=vsa_device,
+                                              exempt=True,
+                                              dense_layers=(),
+                                              tile_size=tile_size)
             torch.compiler.cudagraph_mark_step_begin()
-            with set_forward_context(current_timestep=index, attn_metadata=None, forward_batch=batch):
+            with set_forward_context(current_timestep=index, attn_metadata=attn_metadata, forward_batch=batch):
                 velocity, _ = transformer(hidden_states=state.to(torch.bfloat16)[None],
                                           audio_hidden_states=audio,
                                           encoder_hidden_states=prompt_embeds,

fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py

@@ -186,9 +186,13 @@ class _ConstantVelocityTransformer:
     def __init__(self, velocity):
         self.velocity = velocity
         self.calls = []
+        self.attn_metadata = []
 
     def __call__(self, **kwargs):
+        from fastvideo.forward_context import get_forward_context
+
         self.calls.append(kwargs)
+        self.attn_metadata.append(get_forward_context().attn_metadata)
         rows = kwargs["hidden_states"].shape[1]
         audio_rows = kwargs["audio_hidden_states"].shape[1]
         return torch.full((1, rows, 8), self.velocity), torch.zeros(1, audio_rows, 4)
@@ -207,13 +211,24 @@ class _RecordingCurve:
         self.timesteps.append(unique.clone())
 
 
-def _denoise(monkeypatch, steps, velocity=0.5):
+def _refine_args(vsa_builder=None):
+    from types import SimpleNamespace
+
+    return SimpleNamespace(pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 2, 2))),
+                           VSA_sparsity=0.8,
+                           VSA_tile_size=64,
+                           vsa_builder=vsa_builder)
+
+
+def _denoise(monkeypatch, steps, velocity=0.5, vsa_builder=None):
     from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence
     from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY
     from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising
     from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 
     monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(denoising, "_h3_vsa_metadata_builder", lambda *args: vsa_builder)
+    monkeypatch.setattr(denoising, "_h3_vsa_single_region_segments", lambda layout, patch: ("segments", ))
     layout = build_packed_sequence(torch.ones(3, dtype=torch.long), 2, 2, 2, 2, (1, 2, 2))
     batch = ForwardBatch(data_type="video")
     batch.num_inference_steps = steps
@@ -223,7 +238,7 @@ def _denoise(monkeypatch, steps, velocity=0.5):
     batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout
     transformer, curve = _ConstantVelocityTransformer(velocity), _RecordingCurve()
     stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None)
-    return stage.forward(batch, None), transformer, curve
+    return stage.forward(batch, _refine_args(vsa_builder)), transformer, curve
 
 
 def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch):
@@ -323,3 +338,27 @@ def test_saved_draft_latent_rejects_the_wrong_rank(tmp_path):
     torch.save(torch.zeros(24, 31, 48), path)
     with pytest.raises(ValueError, match=r"\[C, T, h, w\]"):
         _preparation_stage(tmp_path, str(path))._load_draft_latent(str(path), torch.device("cpu"))
+
+
+class _RecordingVSABuilder:
+
+    def __init__(self):
+        self.builds = []
+
+    def build(self, **kwargs):
+        self.builds.append(kwargs)
+        return f"metadata-{len(self.builds)}"
+
+
+def test_vsa_backend_builds_one_sparse_metadata_per_refine_step(monkeypatch):
+    builder = _RecordingVSABuilder()
+    _, transformer, _ = _denoise(monkeypatch, steps=3, vsa_builder=builder)
+    assert [build["VSA_sparsity"] for build in builder.builds] == [0.8, 0.8, 0.8]
+    assert all(build["tile_size"] == 64 and build["exempt"] for build in builder.builds)
+    assert [build["current_timestep"] for build in builder.builds] == [0, 1, 2]
+    assert transformer.attn_metadata == ["metadata-1", "metadata-2", "metadata-3"]
+
+
+def test_dense_backend_passes_no_attention_metadata(monkeypatch):
+    _, transformer, _ = _denoise(monkeypatch, steps=3)
+    assert transformer.attn_metadata == [None, None, None]

scripts/regen2k_fasth3/score_all.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+# usage: score_all.sh DIR...
+REF=/mnt/nfs/vlm-aryan/regen2k/out/parity_fox/ref_frames_u8.pt
+source /mnt/nfs/vlm-aryan/regen2k-fast/scripts/common.sh
+for d in "$@"; do
+  echo "== $d"
+  $PY $W/scripts/score_run.py --ref $REF --dir $d
+done

scripts/regen2k_fasth3/score_run.py

@@ -16,7 +16,7 @@ import torch
 import torch.nn.functional as F
 
 
-def gaussian_window(size: int = 11, sigma: float = 1.5, device: str = "cuda") -> torch.Tensor:
+def gaussian_window(size: int = 11, sigma: float = 1.5, device: str = "cpu") -> torch.Tensor:
     coords = torch.arange(size, dtype=torch.float32, device=device) - size // 2
     kernel = torch.exp(-(coords**2) / (2 * sigma**2))
     kernel = kernel / kernel.sum()
@@ -29,8 +29,8 @@ def psnr(x: torch.Tensor, y: torch.Tensor) -> float:
 
 
 def ssim(x: torch.Tensor, y: torch.Tensor, window: torch.Tensor) -> float:
-    x = x.float().mean(-1)[None, None].cuda()
-    y = y.float().mean(-1)[None, None].cuda()
+    x = x.float().mean(-1)[None, None]
+    y = y.float().mean(-1)[None, None]
     c1, c2 = (0.01 * 255)**2, (0.03 * 255)**2
     mu_x, mu_y = F.conv2d(x, window), F.conv2d(y, window)
     sxx = F.conv2d(x * x, window) - mu_x * mu_x