[bugfix]: share compiled VSA graphs across H3 layers

Author: Satyam-53

fastvideo/attention/backends/video_sparse_attn_h3.py

@@ -442,13 +442,22 @@ def __init__(
         self.prefix = prefix
         self.layer_idx = layer_idx_from_prefix(prefix, default=-1)
         self.head_size = head_size
+        # Generic torch.compile must not specialize the shared VSA forward on
+        # the Python ``layer_idx`` value of each of H3's 50 blocks. This
+        # tensor is prepared after weights load and drives only the compiled
+        # dense-layer decision; it does not opt the module into sm_100a.
+        self._compile_layer_idx: torch.Tensor | None = None
         # None means the regional-compile preparation hook has not run.  The
         # eager path deliberately ignores this cache and preserves its
         # request-time env/probe/fallback behavior; only Dynamo capture reads
         # the prepared, static route.
         self._regional_compile_sm100a_enabled: bool | None = None
         self._regional_compile_layer_idx: torch.Tensor | None = None
 
+    def prepare_for_compile(self, device: torch.device) -> None:
+        """Tensorize per-layer state shared by every torch.compile route."""
+        self._compile_layer_idx = torch.tensor(self.layer_idx, device=device, dtype=torch.int64)
+
     def prepare_for_regional_compile(self, device: torch.device) -> str | None:
         """Resolve the inference-only sm_100a route before fullgraph capture.
 
@@ -459,6 +468,7 @@ def prepare_for_regional_compile(self, device: torch.device) -> str | None:
         the loaded model's device now, then let ``forward`` specialize on the
         resulting plain bool while Dynamo is compiling.
         """
+        self.prepare_for_compile(device)
         requested = os.environ.get(VSA_SM100A_ENV, "0") == "1"
         enabled = False
         reason = None if requested else f"{VSA_SM100A_ENV}=1 is required for compile-safe VSA-H3 attention"
@@ -603,13 +613,14 @@ def forward(  # type: ignore[override]
         logical_gate = gate_compress[:, :logical_seq_len] if gate_compress is not None else None
 
         # Probe-guided per-layer opt-out: diffuse layers run dense (all-True
-        # mask) while the rest keep the configured sparsity. During regional
-        # capture, keep the layer decision tensor-valued so the 50 block
-        # instances reuse one graph instead of specializing on layer_idx.
+        # mask) while the rest keep the configured sparsity. During any
+        # prepared capture, keep the layer decision tensor-valued so the 50
+        # block instances reuse one graph instead of specializing on the
+        # Python layer_idx attribute.
         force_dense = None
-        if regional_compiling:
-            assert self._regional_compile_layer_idx is not None
-            force_dense = (attn_metadata.dense_layers_tensor == self._regional_compile_layer_idx).any()
+        compile_layer_idx = self._compile_layer_idx if compiling else None
+        if compile_layer_idx is not None:
+            force_dense = (attn_metadata.dense_layers_tensor == compile_layer_idx).any()
             layer_sparsity = attn_metadata.VSA_sparsity
         else:
             layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity

fastvideo/models/dits/minimax_h3.py

@@ -733,31 +733,49 @@ def __init__(self, config: MiniMaxH3Config, hf_config: dict[str, Any]) -> None:
         )
         self.__post_init__()
 
+    @staticmethod
+    def _compile_setup_device(attention: MiniMaxH3Attention) -> torch.device:
+        """Return the loaded device even when FP8 replaced the query weight."""
+        query_state = next(attention.to_q.parameters(), None)
+        if query_state is None:
+            query_state = next(attention.to_q.buffers(), None)
+        if query_state is None:
+            raise RuntimeError("MiniMax H3 to_q has no materialized parameter or buffer for compile setup.")
+        return query_state.device
+
     def prepare_for_compile(self) -> None:
         """Pipeline hook, called once right before torch.compile wraps the blocks.
 
-        Resolve each loaded VSA compression gate eagerly. Generic and training
-        compile retain their established attention dispatch; only the
-        inference loader's separate ``prepare_for_regional_compile`` hook may
-        preselect the inference-only sm_100a path.
+        Resolve each loaded VSA compression gate eagerly and tensorize its
+        layer identity so repeated blocks share one Dynamo graph. Generic and
+        training compile retain their established attention dispatch; only
+        the inference loader's separate ``prepare_for_regional_compile`` hook
+        may preselect the inference-only sm_100a path.
 
         The inference-only Triton fusions expose fake-backed custom operators,
         so Dynamo can keep them active as opaque nodes inside each fullgraph
         block instead of tracing into their launcher implementation.
         """
         gate_states: list[bool] = []
+        prepared_vsa_impls = 0
         for block in self.transformer_blocks:
             attention = block.attn
             if attention.to_gate_compress is not None:
                 attention._resolve_gate_compress_for_compile()
                 assert attention._gate_compress_active is not None
                 gate_states.append(attention._gate_compress_active)
+            prepare_vsa = getattr(attention.distributed_attention.attn_impl, "prepare_for_compile", None)
+            if callable(prepare_vsa):
+                prepare_vsa(self._compile_setup_device(attention))
+                prepared_vsa_impls += 1
         if gate_states:
             logger.info(
                 "Resolved MiniMax H3 VSA compression gates before torch.compile: %d active, %d inactive",
                 sum(gate_states),
                 len(gate_states) - sum(gate_states),
             )
+        if prepared_vsa_impls:
+            logger.info("Prepared %d MiniMax H3 VSA layer indices for torch.compile", prepared_vsa_impls)
         if self.enabled_fusions:
             logger.info(
                 "MiniMax H3 inference fusions remain active under torch.compile through custom-op boundaries: %s",
@@ -774,14 +792,7 @@ def prepare_for_regional_compile(self) -> str | None:
             prepare_vsa = getattr(attention.distributed_attention.attn_impl, "prepare_for_regional_compile", None)
             if not callable(prepare_vsa):
                 continue
-            # Post-load FP8 conversion may replace to_q.weight with packed
-            # buffers. Either representation identifies the local device.
-            query_state = next(attention.to_q.parameters(), None)
-            if query_state is None:
-                query_state = next(attention.to_q.buffers(), None)
-            if query_state is None:
-                raise RuntimeError("MiniMax H3 to_q has no materialized parameter or buffer for compile setup.")
-            unsupported = prepare_vsa(query_state.device)
+            unsupported = prepare_vsa(self._compile_setup_device(attention))
             if unsupported:
                 unsupported_reasons.add(str(unsupported))
             prepared_vsa_impls += 1

fastvideo/tests/attention/test_vsa_h3_sm100a_route.py

@@ -308,6 +308,61 @@ def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes):
         torch._dynamo.reset()
 
 
+def test_generic_compile_reuses_graph_across_layer_indices_and_stays_on_triton(monkeypatch):
+    """Pipeline compile must share one graph without selecting sm_100a."""
+    fake_sm = _FakeSm100a(supported=True)
+    monkeypatch.setattr(vsa_h3, "_sm100a", fake_sm)
+    monkeypatch.setenv(VSA_SM100A_ENV, "1")
+    monkeypatch.setattr(vsa_h3, "probe_enabled", lambda: None)
+    meta = _build_meta(sparsity=0.5, dense_layers=(0, 17), prefix_segments=(64, 64))
+    q, k, v = _tiled_qkv(meta)
+
+    def fake_triton(q, k, v, block_map, variable_block_sizes):
+        del k, v, variable_block_sizes
+        return q + block_map.all().to(q.dtype), None
+
+    def fail_sm100a(*args, **kwargs):
+        raise AssertionError("generic torch.compile unexpectedly selected sm_100a")
+
+    monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fake_triton)
+    monkeypatch.setattr(fake_sm, "is_supported", fail_sm100a)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_sm100a)
+    monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", fail_sm100a)
+    monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_sm100a)
+
+    implementations = []
+    for layer_idx in range(20):
+        impl = MiniMaxH3VSAImpl(
+            num_heads=_HEADS,
+            head_size=_DIM,
+            causal=False,
+            softmax_scale=_DIM**-0.5,
+            prefix=f"transformer_blocks.{layer_idx}.attn",
+        )
+        impl.prepare_for_compile(torch.device("cpu"))
+        implementations.append(impl)
+
+    compiled_graphs = []
+
+    def recording_backend(graph_module, _example_inputs):
+        compiled_graphs.append(graph_module)
+        return graph_module.forward
+
+    torch._dynamo.reset()
+    try:
+        compiled = [torch.compile(impl.forward, backend=recording_backend, fullgraph=True)
+                    for impl in implementations]
+        with torch.inference_mode():
+            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
+                torch.testing.assert_close(actual, q + expected_delta, atol=0, rtol=0)
+    finally:
+        torch._dynamo.reset()
+
+    assert len(compiled_graphs) == 1
+
+
 def test_default_off_routes_triton(routed, monkeypatch):
     fake_sm, fake_triton, run, _ = routed
     monkeypatch.delenv(VSA_SM100A_ENV, raising=False)

fastvideo/tests/inference/test_inference_regional_compile.py

@@ -106,11 +106,15 @@ def test_h3_vsa_probe_degrades_regional_compile_to_eager(monkeypatch) -> None:
 class _RegionalPrepareProbe:
 
     def __init__(self, unsupported: str | None = None) -> None:
-        self.devices: list[torch.device] = []
+        self.compile_devices: list[torch.device] = []
+        self.regional_devices: list[torch.device] = []
         self.unsupported = unsupported
 
+    def prepare_for_compile(self, device: torch.device) -> None:
+        self.compile_devices.append(device)
+
     def prepare_for_regional_compile(self, device: torch.device) -> str | None:
-        self.devices.append(device)
+        self.regional_devices.append(device)
         return self.unsupported
 
 
@@ -157,7 +161,8 @@ def test_minimax_h3_prepare_for_compile_resolves_loaded_vsa_gates() -> None:
     assert [block.attn._gate_compress_active for block in model.transformer_blocks] == [False, True]
     for block in model.transformer_blocks:
         impl = block.attn.distributed_attention.attn_impl
-        assert impl.devices == []
+        assert impl.compile_devices == [next(block.parameters()).device]
+        assert impl.regional_devices == []
 
 
 def test_training_compile_prepare_does_not_probe_inference_kernel() -> None:
@@ -168,7 +173,9 @@ def test_training_compile_prepare_does_not_probe_inference_kernel() -> None:
     assert reason is None
     attention = model.transformer_blocks[0].attn
     assert attention._gate_compress_active is True
-    assert attention.distributed_attention.attn_impl.devices == []
+    impl = attention.distributed_attention.attn_impl
+    assert impl.compile_devices == [next(model.parameters()).device]
+    assert impl.regional_devices == []
 
 
 def test_regional_compile_prepare_prefers_specialized_hook() -> None:
@@ -179,7 +186,8 @@ def test_regional_compile_prepare_prefers_specialized_hook() -> None:
 
     assert reason is None
     impl = model.transformer_blocks[0].attn.distributed_attention.attn_impl
-    assert impl.devices == [expected_device]
+    assert impl.compile_devices == [expected_device]
+    assert impl.regional_devices == [expected_device]
 
 
 def test_minimax_h3_prepare_for_regional_compile_does_not_require_quantized_q_weight() -> None:
@@ -193,7 +201,8 @@ def test_minimax_h3_prepare_for_regional_compile_does_not_require_quantized_q_we
 
     assert reason is None
     impl = attention.distributed_attention.attn_impl
-    assert impl.devices == [expected_device]
+    assert impl.compile_devices == [expected_device]
+    assert impl.regional_devices == [expected_device]
 
 
 def test_minimax_h3_prepare_for_regional_compile_propagates_backend_rejection() -> None: