From 0033a9b11dc5eeddf9be54533483be37b7cfd603 Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:22:06 -0700 Subject: [PATCH] [bugfix] Fix FastH3 regional compile recompiling for every prompt length (#1942) Co-authored-by: Aryan Kumar Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../backends/video_sparse_attn_h3.py | 58 +++++--- fastvideo/models/dits/minimax_h3.py | 11 +- .../attention/test_vsa_h3_blackwell_route.py | 130 +++++++++++++++++- 3 files changed, 179 insertions(+), 20 deletions(-) diff --git a/fastvideo/attention/backends/video_sparse_attn_h3.py b/fastvideo/attention/backends/video_sparse_attn_h3.py index 5852f66ae..e369cf44f 100644 --- a/fastvideo/attention/backends/video_sparse_attn_h3.py +++ b/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, ) diff --git a/fastvideo/models/dits/minimax_h3.py b/fastvideo/models/dits/minimax_h3.py index 4eb579d9f..f29f6ce5f 100644 --- a/fastvideo/models/dits/minimax_h3.py +++ b/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) diff --git a/fastvideo/tests/attention/test_vsa_h3_blackwell_route.py b/fastvideo/tests/attention/test_vsa_h3_blackwell_route.py index cccfbcdd1..309149080 100644 --- a/fastvideo/tests/attention/test_vsa_h3_blackwell_route.py +++ b/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()