[bugfix] Fix FastH3 regional compile recompiling for every prompt length (#1942)

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -157,6 +157,21 @@ def _h3_vsa_ptx_blackwell_from_mask_compat_fake(
     return torch.empty_like(q)
 
 
+@torch.library.custom_op("fastvideo::h3_vsa_block_mask", mutates_args=())
+def _h3_vsa_block_mask(scores: torch.Tensor) -> torch.Tensor:
+    """``_build_block_mask`` reading the per-request geometry at run time, so compiled graphs don't guard on it."""
+    from fastvideo.forward_context import get_forward_context
+
+    metadata = get_forward_context().attn_metadata
+    return _build_block_mask(scores, metadata.num_prefix_tiles, metadata.VSA_sparsity, metadata.exempt,
+                             metadata.video_tile_spans, metadata.span_sparsities)
+
+
+@torch.library.register_fake("fastvideo::h3_vsa_block_mask")
+def _h3_vsa_block_mask_fake(scores: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(scores, dtype=torch.bool)
+
+
 def _ptx_blackwell_has_compile_safe_mask_route(ptx_blackwell_mod: Any) -> bool:
     return (callable(getattr(ptx_blackwell_mod, "block_sparse_attn_ptx_blackwell_from_mask", None))
             or (callable(getattr(ptx_blackwell_mod, "block_sparse_attn_ptx_blackwell", None))
@@ -638,14 +653,14 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         additional all-zero tile internally; metadata and all observable
         outputs retain the logical geometry.
         """
-        if x.shape[1] != attn_metadata.total_seq_length:
+        compiling = torch.compiler.is_compiling()
+        regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True
+        if not regional_compiling and x.shape[1] != attn_metadata.total_seq_length:
             raise ValueError(f"VSA-H3 metadata was built for sequence length {attn_metadata.total_seq_length}, "
                              f"got {x.shape[1]}. A non-packed sequence (e.g. the token refiner) is "
                              "routed to the VSA-H3 backend; exclude it from the supported backends.")
         n_tiles = attn_metadata.variable_block_sizes.numel()
         grad_mode = torch.is_grad_enabled() and x.requires_grad
-        compiling = torch.compiler.is_compiling()
-        regional_compiling = compiling and self._regional_compile_ptx_blackwell_enabled is True
         if regional_compiling:
             ptx_blackwell_requested = bool(self._regional_compile_ptx_blackwell_enabled)
         elif compiling:
@@ -656,8 +671,12 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         # Tile 128 has no route other than the ptx_blackwell CUDA kernel.
         ptx_blackwell_route = attn_metadata.tile_elems == 128 or (attn_metadata.tile_elems == 64
                                                                   and ptx_blackwell_requested)
-        needs_ptx_blackwell_pair = n_tiles % 2 != 0 and not grad_mode and ptx_blackwell_route
-        kernel_tiles = n_tiles + int(needs_ptx_blackwell_pair)
+        if regional_compiling:
+            # Pad arithmetically so odd and even tile counts share one graph.
+            needs_ptx_blackwell_pair, kernel_tiles = False, n_tiles + n_tiles % 2
+        else:
+            needs_ptx_blackwell_pair = n_tiles % 2 != 0 and not grad_mode and ptx_blackwell_route
+            kernel_tiles = n_tiles + int(needs_ptx_blackwell_pair)
         target_shape = (x.shape[0], kernel_tiles * attn_metadata.tile_elems, x.shape[-2], x.shape[-1])
 
         # A grad-tracking forward must not reuse the builder-owned buffer. The
@@ -668,7 +687,7 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         # autograd edges, and through them the activations they saved, for the
         # rest of training. This is the VSA-H3 counterpart of the Wan tile-cache
         # OOM (#1423), which training fixed with ``vsa_cache_tile_buf=False``.
-        if grad_mode:
+        if grad_mode or regional_compiling:
             return scatter_into_tile_buf(x, target_shape, attn_metadata.untile_combined_index, None)
 
         # ``untile_combined_index`` maps each packed row to a logical tile
@@ -753,14 +772,19 @@ class MiniMaxH3VSAImpl(AttentionImpl):
         # the logical prefix, and reject every other shape before a kernel sees it.
         n_tiles = attn_metadata.variable_block_sizes.numel()
         logical_seq_len = n_tiles * tile_elems
-        pair_pad_seq_len = logical_seq_len + tile_elems
-        pair_pad_is_valid = tile_elems in _PTX_BLACKWELL_TILE_ELEMS and n_tiles % 2 != 0
-        allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, )
-        if query.shape[1] not in allowed_seq_lengths:
-            expected = (f"the logical length {logical_seq_len} or one ptx_blackwell partner tile "
-                        f"({pair_pad_seq_len})" if pair_pad_is_valid else f"the logical length {logical_seq_len}")
-            raise ValueError(f"VSA-H3 tiled query has length {query.shape[1]}, expected {expected}.")
-        has_ptx_blackwell_pair = query.shape[1] == pair_pad_seq_len
+        if regional_compiling:
+            pair_tiles = query.shape[1] // tile_elems - n_tiles
+            has_ptx_blackwell_pair = True
+        else:
+            pair_pad_seq_len = logical_seq_len + tile_elems
+            pair_pad_is_valid = tile_elems in _PTX_BLACKWELL_TILE_ELEMS and n_tiles % 2 != 0
+            allowed_seq_lengths = (logical_seq_len, pair_pad_seq_len) if pair_pad_is_valid else (logical_seq_len, )
+            if query.shape[1] not in allowed_seq_lengths:
+                expected = (f"the logical length {logical_seq_len} or one ptx_blackwell partner tile "
+                            f"({pair_pad_seq_len})" if pair_pad_is_valid else f"the logical length {logical_seq_len}")
+                raise ValueError(f"VSA-H3 tiled query has length {query.shape[1]}, expected {expected}.")
+            has_ptx_blackwell_pair = query.shape[1] == pair_pad_seq_len
+            pair_tiles = int(has_ptx_blackwell_pair)
         for name, tensor in (("key", key), ("value", value)):
             if tensor.shape[1] != query.shape[1]:
                 raise ValueError(f"VSA-H3 tiled {name} length {tensor.shape[1]} does not match query "
@@ -798,6 +822,8 @@ class MiniMaxH3VSAImpl(AttentionImpl):
 
         if scores is None:
             mask = torch.ones(query.shape[0], query.shape[2], n_tiles, n_tiles, dtype=torch.bool, device=query.device)
+        elif regional_compiling:
+            mask = _h3_vsa_block_mask(scores)
         else:
             mask = _build_block_mask(
                 scores,
@@ -837,10 +863,10 @@ class MiniMaxH3VSAImpl(AttentionImpl):
                 # The synthetic tile is neither a logical query nor key. Its
                 # all-False row yields q2k_num=0, the all-False column keeps it
                 # out of real rows, and vbs=0 masks all of its key slots.
-                ptx_blackwell_mask = torch.nn.functional.pad(mask, (0, 1, 0, 1), value=False)
+                ptx_blackwell_mask = torch.nn.functional.pad(mask, (0, pair_tiles, 0, pair_tiles), value=False)
                 ptx_blackwell_variable_block_sizes = torch.nn.functional.pad(
                     attn_metadata.variable_block_sizes,
-                    (0, 1),
+                    (0, pair_tiles),
                     value=0,
                 )
 

fastvideo/models/dits/minimax_h3.py

@@ -1035,8 +1035,15 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
                         prepared_vsa_impls)
         if unsupported_reasons:
             return "; ".join(sorted(unsupported_reasons))
+        self._regional_compile_inputs = True
         return None
 
+    def _block_input(self, shard: torch.Tensor) -> torch.Tensor:
+        """Give a sequence-parallel shard standard strides under regional compile, so block 0 shares one graph."""
+        if getattr(self, "_regional_compile_inputs", False):
+            return shard.clone(memory_format=torch.contiguous_format)
+        return shard
+
     def materialize_non_persistent_buffers(
         self,
         device: torch.device,
@@ -1136,7 +1143,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         text_original_seq_len = text_embeds.shape[1]
         sp_world_size = get_sp_world_size() if model_parallel_is_initialized() else 1
         if sp_world_size > 1:
-            text_embeds, _ = sequence_model_parallel_shard(text_embeds, dim=1)
+            text_embeds = self._block_input(sequence_model_parallel_shard(text_embeds, dim=1)[0])
         text_embeds = self.token_refiner(text_embeds, text_original_seq_len)
         if sp_world_size > 1:
             text_embeds = sequence_model_parallel_all_gather_with_unpad(
@@ -1201,7 +1208,7 @@ class MiniMaxH3Transformer3DModel(BaseDiT):
         original_seq_len = sequence_length
 
         if sp_world_size > 1:
-            packed_hidden_states, _ = sequence_model_parallel_shard(packed_hidden_states, dim=1)
+            packed_hidden_states = self._block_input(sequence_model_parallel_shard(packed_hidden_states, dim=1)[0])
             rotary_cos, _ = sequence_model_parallel_shard(rotary_emb[0], dim=0)
             rotary_sin, _ = sequence_model_parallel_shard(rotary_emb[1], dim=0)
             adaln_indices, _ = sequence_model_parallel_shard(adaln_indices, dim=0)

fastvideo/tests/attention/test_vsa_h3_blackwell_route.py

@@ -296,10 +296,12 @@ def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_ov
     monkeypatch.setattr(vsa_h3, "_ptx_blackwell_unavailable_reason", fail_eager_dispatch)
     monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch)
 
+    from fastvideo.forward_context import set_forward_context
+
     torch._dynamo.reset()
     try:
         compiled = [torch.compile(impl.forward, backend="eager", fullgraph=True) for impl in implementations]
-        with torch.inference_mode():
+        with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta):
             for layer_idx, run in enumerate(compiled):
                 actual = run(q, k, v, None, meta)
                 expected_delta = 1.0 if layer_idx in meta.dense_layers else 0.0
@@ -664,9 +666,11 @@ def test_real_ptx_blackwell_odd_preprocess_forward_postprocess_fullgraph(env_ove
         output = impl.forward(query, key, value, None, meta)
         return impl.postprocess_output(output, meta)
 
+    from fastvideo.forward_context import set_forward_context
+
     torch._dynamo.reset()
     try:
-        with torch.inference_mode():
+        with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=meta):
             compiled = torch.compile(run, fullgraph=True)
             output = compiled(raw_qkv)
             torch.cuda.synchronize()
@@ -675,3 +679,125 @@ def test_real_ptx_blackwell_odd_preprocess_forward_postprocess_fullgraph(env_ove
 
     assert output.shape == (1, meta.total_seq_length, _HEADS, _DIM)
     assert torch.isfinite(output).all()
+
+
+def test_regional_compile_reuses_graphs_across_prompt_lengths(monkeypatch, env_overrides):
+    """Unseen prompt lengths reuse warm graphs and match the eager route exactly (#1940)."""
+    from torch._dynamo.testing import CompileCounter
+
+    from fastvideo.forward_context import set_forward_context
+
+    def mask_count(q, num):
+        return q + num.repeat_interleave(q.shape[2] // num.shape[-1], dim=-1).unsqueeze(-1).to(q.dtype)
+
+    fake_sm = _FakePtxBlackwell(supported=True)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell",
+                        lambda q, k, v, q2k_idx, q2k_num, vbs, need_lse=True: (mask_count(q, q2k_num), None))
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell_from_mask",
+                        lambda q, k, v, block_map, vbs: (mask_count(q, block_map.sum(-1)), None))
+    monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm)
+    monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index)
+    env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True))
+    monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+    assert impl.prepare_for_regional_compile(torch.device("cpu")) is None
+    builder = MiniMaxH3VSAMetadataBuilder()
+    counter = CompileCounter()
+
+    def run(raw, meta):
+        tiled = impl.preprocess_qkv(raw, meta)
+        query, key, value, gate = tiled.chunk(4, dim=0)
+        return impl.postprocess_output(impl.forward(query, key, value, gate, meta), meta)
+
+    def check(text_rows_list):
+        parities = set()
+        for text_rows in text_rows_list:
+            meta = _build_meta(prefix_segments=(text_rows, 30), sparsity=0.5, builder=builder)
+            parities.add(meta.variable_block_sizes.numel() % 2)
+            raw = torch.randn(4, meta.total_seq_length, _HEADS, _DIM, dtype=torch.bfloat16)
+            with set_forward_context(current_timestep=0, attn_metadata=meta):
+                actual = compiled(raw, meta)
+                expected = run(raw.clone(), meta)
+            torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        return parities
+
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(run, backend=counter, fullgraph=True)
+        with torch.no_grad():
+            # Dynamo may still specialize on bounded shape classes (tile-count
+            # parity); once both are warm, unseen lengths must reuse a graph.
+            assert check((17, 40, 64, 81, 100, 130)) == {0, 1}
+            warm_graphs = counter.frame_count
+            assert check((151, 200, 233, 300, 47, 113)) == {0, 1}
+        assert counter.frame_count == warm_graphs, (
+            f"unseen prompt lengths compiled {counter.frame_count - warm_graphs} new graphs")
+    finally:
+        torch._dynamo.reset()
+
+
+def test_regional_compile_reuses_graphs_across_sp_padding_lengths(monkeypatch, env_overrides):
+    """Sequence lengths that do and don't divide by the SP degree reuse warm graphs (#1940).
+
+    ``DistributedAttention_VSA`` slices the SP padding off before the backend and pads the output back, so the
+    padded length (a function of ``total_seq_length % sp``) must not leak into the compiled graph's guards.
+    """
+    from torch._dynamo.testing import CompileCounter
+
+    from fastvideo.attention import layer as attn_layer
+    from fastvideo.forward_context import set_forward_context
+
+    sp_degree = 4
+
+    def mask_count(q, num):
+        return q + num.repeat_interleave(q.shape[2] // num.shape[-1], dim=-1).unsqueeze(-1).to(q.dtype)
+
+    fake_sm = _FakePtxBlackwell(supported=True)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell",
+                        lambda q, k, v, q2k_idx, q2k_num, vbs, need_lse=True: (mask_count(q, q2k_num), None))
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_ptx_blackwell_from_mask",
+                        lambda q, k, v, block_map, vbs: (mask_count(q, block_map.sum(-1)), None))
+    monkeypatch.setattr(vsa_h3, "_ptx_blackwell", fake_sm)
+    monkeypatch.setattr(vsa_h3, "map_to_index", _fake_map_to_index)
+    monkeypatch.setattr(attn_layer, "sequence_model_parallel_all_to_all_4D", lambda x, **_: x)
+    env_overrides.enter_context(envs.FASTVIDEO_VSA_PTX_BLACKWELL.override(True))
+    monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None)
+    impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
+    assert impl.prepare_for_regional_compile(torch.device("cpu")) is None
+    layer = attn_layer.DistributedAttention_VSA.__new__(attn_layer.DistributedAttention_VSA)
+    torch.nn.Module.__init__(layer)
+    layer.attn_impl = impl
+    layer._compile_forward_enabled = True
+    builder = MiniMaxH3VSAMetadataBuilder()
+    counter = CompileCounter()
+
+    def run(q, k, v, total):
+        return layer(q, k, v, total)[0]
+
+    def check(text_rows_list):
+        pads = set()
+        for text_rows in text_rows_list:
+            meta = _build_meta(prefix_segments=(text_rows, 30), sparsity=0.5, builder=builder)
+            total = meta.total_seq_length
+            pad = -total % sp_degree
+            pads.add(pad)
+            q, k, v = (torch.randn(1, total + pad, _HEADS, _DIM, dtype=torch.bfloat16) for _ in range(3))
+            with set_forward_context(current_timestep=0, attn_metadata=meta):
+                actual = compiled(q, k, v, total)
+                expected = run(q, k, v, total)
+            assert actual.shape == (1, total + pad, _HEADS, _DIM)
+            torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        return pads
+
+    torch._dynamo.reset()
+    try:
+        compiled = torch.compile(run, backend=counter, fullgraph=True)
+        with torch.no_grad():
+            warm_pads = check((17, 40, 64, 81, 100, 130))
+            assert 0 in warm_pads and len(warm_pads) > 1, "need lengths that do and don't divide by the SP degree"
+            warm_graphs = counter.frame_count
+            check((151, 200, 233, 300, 47, 113, 18, 19, 20))
+        assert counter.frame_count == warm_graphs, (
+            f"unseen padded lengths compiled {counter.frame_count - warm_graphs} new graphs")
+    finally:
+        torch._dynamo.reset()