diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index af6446255..f736ea68f 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/fastvideo/attention/backends/video_sparse_attn_h3.py @@ -233,6 +233,10 @@ class MiniMaxH3VSAMetadata(AttentionMetadata): exempt: bool variable_block_sizes: torch.Tensor untile_combined_index: torch.Tensor + # Device-side copy of ``dense_layers``. Regional fullgraph capture uses + # this tensor with each implementation's tensor-valued layer index so the + # shared block code does not specialize once per Python ``layer_idx``. + dense_layers_tensor: torch.Tensor # tokens per tile (256 or 64); selects the tile geometry AND the kernel # route in forward() (256 -> VSA-256 CuTe/Triton, 64 -> native Triton) tile_elems: int = _TILE_ELEMS @@ -278,6 +282,7 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): (_tile_partition_indices, variable_block_sizes, untile_combined_index, num_prefix_tiles, num_video_tiles) = _h3_tile_geometry(prefix_segments, dit_seq_shape, device, tile_shape) + dense_layers = tuple(int(layer) for layer in dense_layers) return MiniMaxH3VSAMetadata( current_timestep=current_timestep, VSA_sparsity=VSA_sparsity, @@ -288,7 +293,8 @@ class MiniMaxH3VSAMetadataBuilder(AttentionMetadataBuilder): variable_block_sizes=variable_block_sizes, untile_combined_index=untile_combined_index, tile_elems=int(tile_size), - dense_layers=tuple(int(layer) for layer in dense_layers), + dense_layers=dense_layers, + dense_layers_tensor=torch.tensor(dense_layers, device=device, dtype=torch.int64), tile_buf_holder=self._tile_buf_holder, ) @@ -372,6 +378,7 @@ class MiniMaxH3VSAImpl(AttentionImpl): # 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_regional_compile(self, device: torch.device) -> None: """Resolve the inference-only sm_100a route before fullgraph capture. @@ -383,6 +390,12 @@ class MiniMaxH3VSAImpl(AttentionImpl): the loaded model's device now, then let ``forward`` specialize on the resulting plain bool while Dynamo is compiling. """ + # Dynamo specializes Python integer attributes of nn.Module-owned + # objects. All 50 H3 blocks share one forward code object, so reading + # ``self.layer_idx`` in every regional graph exhausts the recompile + # cache. Tensor contents stay dynamic and let the blocks share a graph. + self._regional_compile_layer_idx = torch.tensor(self.layer_idx, device=device, dtype=torch.int64) + requested = os.environ.get(VSA_SM100A_ENV, "0") == "1" enabled = False reason = None @@ -458,9 +471,20 @@ class MiniMaxH3VSAImpl(AttentionImpl): # environment switch while Dynamo captures a regional full graph. compiling = torch.compiler.is_compiling() - # probe-guided per-layer opt-out: diffuse layers run dense (all-True - # mask) while the rest keep the configured sparsity - layer_sparsity = 0.0 if self.layer_idx in attn_metadata.dense_layers else attn_metadata.VSA_sparsity + # 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. + force_dense = None + if compiling: + if self._regional_compile_layer_idx is None: + raise RuntimeError( + "VSA-H3 layer routing was not resolved before torch.compile; " + "call prepare_for_regional_compile(device) on every MiniMaxH3VSAImpl after loading weights.") + force_dense = (attn_metadata.dense_layers_tensor == self._regional_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 probe_dir = None if compiling else probe_enabled() scores = None @@ -482,6 +506,10 @@ class MiniMaxH3VSAImpl(AttentionImpl): layer_sparsity, attn_metadata.exempt, ) + if force_dense is not None: + # A scalar bool tensor broadcasts over the block map. This exactly + # preserves the eager dense-layer contract without a Python branch. + mask = mask | force_dense if tile_elems == 64: # Native 64-token path: the block map is already at the kernels' diff --git a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py index ef8002e03..2d1af3879 100644 --- a/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py +++ b/fastvideo/tests/attention/test_vsa_h3_sm100a_route.py @@ -20,14 +20,15 @@ _SPEC = dict(raw_latent_shape=(4, 8, 16), patch_size=(1, 2, 2), prefix_segments= _HEADS, _DIM = 2, 128 -def _build_meta(): +def _build_meta(*, sparsity=0.0, dense_layers=()): return MiniMaxH3VSAMetadataBuilder().build( current_timestep=0, raw_latent_shape=_SPEC["raw_latent_shape"], patch_size=_SPEC["patch_size"], - VSA_sparsity=0.0, + VSA_sparsity=sparsity, prefix_segments=_SPEC["prefix_segments"], device=torch.device("cpu"), + dense_layers=dense_layers, tile_size=64, ) @@ -129,6 +130,8 @@ def test_prepare_for_regional_compile_resolves_supported_route(monkeypatch): assert probe_q.dtype == torch.bfloat16 assert probe_vbs.dtype == torch.int32 assert probe_vbs.tolist() == [64, 64] + assert impl._regional_compile_layer_idx is not None + assert impl._regional_compile_layer_idx.item() == -1 def test_prepare_for_regional_compile_env_off_skips_probe(monkeypatch): @@ -211,6 +214,56 @@ def test_prepared_route_fullgraph_avoids_eager_dispatch(monkeypatch): torch._dynamo.reset() +def test_prepared_route_reuses_graph_across_layer_indices_and_preserves_dense_overrides(monkeypatch): + """Fifty H3 blocks must not specialize the shared graph on layer_idx.""" + 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)) + assert meta.dense_layers_tensor.tolist() == [0, 17] + q, k, v = _tiled_qkv(meta) + + def fail_eager_dispatch(*args, **kwargs): + raise AssertionError("compiled route re-entered eager dispatch") + + def compile_safe_from_mask(q, k, v, block_map, variable_block_sizes): + del k, v, variable_block_sizes + # Encode the dense-layer decision in the result so this is a semantic + # test as well as a recompile-cache stress test. + return q + block_map.all().to(q.dtype), None + + 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_regional_compile(torch.device("cpu")) + implementations.append(impl) + + monkeypatch.setattr(fake_sm, "is_supported", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a", fail_eager_dispatch) + monkeypatch.setattr(fake_sm, "block_sparse_attn_sm100a_from_mask", compile_safe_from_mask) + monkeypatch.setattr(vsa_h3, "map_to_index", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "_sm100a_unavailable_reason", fail_eager_dispatch) + monkeypatch.setattr(vsa_h3, "block_sparse_attn_64_bhsd", fail_eager_dispatch) + + torch._dynamo.reset() + try: + compiled = [torch.compile(impl.forward, backend="eager", 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() + + def test_default_off_routes_triton(routed, monkeypatch): fake_sm, fake_triton, run, _ = routed monkeypatch.delenv(VSA_SM100A_ENV, raising=False)