nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf] MiniMax-H3: share compiled VSA graph across layers

Browse Source
This commit is contained in:
Will Lin
2026-08-22 06:35:55 +00:00
parent d20484ff97
commit 75c41e7e35
2 changed files with 87 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/attention/backends/video_sparse_attn_h3.py
+32 -4
View File
@@ -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'
fastvideo/tests/attention/test_vsa_h3_sm100a_route.py
+55 -2
View File
@@ -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)