From a3d9725d1552df94174e9599383bab2b97087a44 Mon Sep 17 00:00:00 2001 From: Satyam Srivastava Date: Sat, 29 Aug 2026 16:37:51 -0700 Subject: [PATCH] [bugfix]: share compiled VSA graphs across H3 layers --- .../backends/video_sparse_attn_h3.py | 23 ++++++-- fastvideo/models/dits/minimax_h3.py | 35 ++++++++---- .../attention/test_vsa_h3_sm100a_route.py | 55 +++++++++++++++++++ .../test_inference_regional_compile.py | 21 +++++-- 4 files changed, 110 insertions(+), 24 deletions(-) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 38f0daa04..c5b74fcf5 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -442,6 +442,11 @@ class MiniMaxH3VSAImpl(AttentionImpl): 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 @@ -449,6 +454,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): 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 @@ class MiniMaxH3VSAImpl(AttentionImpl): 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 @@ class MiniMaxH3VSAImpl(AttentionImpl): 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 diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index ec168cec1..3dff637fe 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/fastvideo/models/dits/minimax_h3.py @@ -733,31 +733,49 @@ class MiniMaxH3Transformer3DModel(BaseDiT): ) 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 @@ class MiniMaxH3Transformer3DModel(BaseDiT): 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 diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index 904c3ca3e..df8b2a366 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -308,6 +308,61 @@ def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_ov 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) diff --git a/fastvideo/tests/inference/test_inference_regional_compile.py b/fastvideo/tests/inference/test_inference_regional_compile.py index ffb5d0b4b..e3aa224e0 100644 --- a/fastvideo/tests/inference/test_inference_regional_compile.py +++ b/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: