[MiniMax-H3] Reuse Ulysses shard lengths across an inference forward

docs/source/en/training/distributed_inference.md

@@ -354,6 +354,20 @@ The default Ulysses Attention mechanism requires that the sequence length of hid
 pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True))
 ```
 
+MiniMax-H3 can reuse its packed-sequence shard lengths across its transformer blocks in one inference forward:
+
+```python
+transformer = pipeline.transformer
+with torch.no_grad():
+    output = transformer(**model_inputs, attention_kwargs={"reuse_ulysses_metadata": True})
+```
+
+Enable Ulysses Anything on the transformer first. Every rank must use the same opt-in setting for a forward and receive
+the same complete packed layout. The token refiner retains the default length collection. The model
+checks the global sequence length once on every forward, then passes an immutable per-call plan to its self-attention
+blocks. It does not cache plans between requests or infer a global layout from a rank-local length. Calls with gradients
+enabled are rejected for this option. Other models, normal Ulysses, Ring Attention, and the default H3 path are unchanged.
+
 > [!TIP] To avoid multiple forced CUDA sync caused by H2D and D2H transfers, please add the **gloo** backend in `init_process_group`. This will significantly reduce communication latency.
 
 We ran a benchmark for FLUX.1-dev with Ulysses, Ring, Unified Attention and Ulysses Anything Attention with [this script](https://github.com/huggingface/diffusers/pull/12996#issuecomment-3797695999) on a node of 4 L20 GPUs. The results are summarized as follows:

src/diffusers/models/_modeling_parallel.py

@@ -210,6 +210,8 @@ class ParallelConfig:
     context_parallel_config: ContextParallelConfig | None = None
     tensor_parallel_config: TensorParallelConfig | None = None
 
+    _ulysses_sequence_lengths: tuple[int, ...] | None = None
+
     _rank: int = None
     _world_size: int = None
     _device: torch.device = None

src/diffusers/models/attention_dispatch.py

@@ -2326,7 +2326,9 @@ def all_to_all_single_any_qkv_async(
     # S_LOCAL maybe not equal for all ranks in dynamic shape case,
     # since we don't know the actual shape before this timing, thus,
     # we have to use all gather to collect the S_LOCAL first.
-    output_split_sizes = gather_size_by_comm(S_LOCAL, group)
+    output_split_sizes = kwargs.get("seq_lengths")
+    if output_split_sizes is None:
+        output_split_sizes = gather_size_by_comm(S_LOCAL, group)
     x = x.flatten(0, 1)  # (world_size * S_LOCAL, B, H_LOCAL, D)
     x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group)
 
@@ -2365,7 +2367,9 @@ def all_to_all_single_any_o_async(x: torch.Tensor, group: dist.ProcessGroup, **k
     # b.tensor_split(4)[0].shape[1])
 
     S_LOCAL = kwargs.get("Q_S_LOCAL")
-    input_split_sizes = gather_size_by_comm(S_LOCAL, group)
+    input_split_sizes = kwargs.get("seq_lengths")
+    if input_split_sizes is None:
+        input_split_sizes = gather_size_by_comm(S_LOCAL, group)
     x = x.permute(1, 0, 2, 3).contiguous()  # (S_GLOBAL, B, H_LOCAL, D)
     output_split_sizes = [S_LOCAL] * world_size
     x = funcol.all_to_all_single(x, output_split_sizes, input_split_sizes, group)
@@ -2780,6 +2784,12 @@ def forward(
         _, S_KV_LOCAL, _, _ = key.shape
 
         metadata = ulysses_anything_metadata(query)
+        sequence_lengths = _parallel_config._ulysses_sequence_lengths
+        if sequence_lengths is not None:
+            local_length = sequence_lengths[dist.get_rank(group)]
+            if query.shape[1] != local_length or key.shape[1] != local_length:
+                raise ValueError("The forward sequence plan requires matching self-attention shard lengths.")
+            metadata = {**metadata, "seq_lengths": list(sequence_lengths)}
         query_wait = all_to_all_single_any_qkv_async(query, group, **metadata)
         key_wait = all_to_all_single_any_qkv_async(key, group, **metadata)
         value_wait = all_to_all_single_any_qkv_async(value, group, **metadata)

src/diffusers/models/transformers/transformer_minimax_h3.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 import torch
@@ -21,7 +21,7 @@
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import BaseOutput, apply_lora_scale, logging
-from .._modeling_parallel import ContextParallelInput, ContextParallelOutput
+from .._modeling_parallel import ContextParallelInput, ContextParallelOutput, ParallelConfig, gather_size_by_comm
 from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward
 from ..attention_dispatch import dispatch_attention_fn
 from ..cache_utils import CacheMixin
@@ -171,6 +171,7 @@ def __call__(
         hidden_states: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         attention_mask: torch.Tensor | None = None,
+        parallel_config: ParallelConfig | None = None,
     ) -> torch.Tensor:
         if attn.fused_projections:
             query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
@@ -201,7 +202,7 @@ def __call__(
             dropout_p=0.0,
             is_causal=False,
             backend=self._attention_backend,
-            parallel_config=self._parallel_config,
+            parallel_config=parallel_config if parallel_config is not None else self._parallel_config,
         )
         hidden_states = hidden_states.flatten(2, 3).type_as(query)
         hidden_states = attn.to_out[0](hidden_states)
@@ -243,8 +244,11 @@ def forward(
         hidden_states: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         attention_mask: torch.Tensor | None = None,
+        parallel_config: ParallelConfig | None = None,
     ) -> torch.Tensor:
-        return self.processor(self, hidden_states, rotary_emb, attention_mask)
+        if parallel_config is None:
+            return self.processor(self, hidden_states, rotary_emb, attention_mask)
+        return self.processor(self, hidden_states, rotary_emb, attention_mask, parallel_config=parallel_config)
 
 
 class MiniMaxH3TokenRefinerBlock(nn.Module):
@@ -351,6 +355,7 @@ def forward(
         adaln_indices: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None = None,
+        parallel_config: ParallelConfig | None = None,
     ) -> torch.Tensor:
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
 
@@ -359,7 +364,7 @@ def forward(
         norm_hidden_states = norm_hidden_states * (
             1.0 + scale_msa.index_select(0, adaln_indices)
         ) + shift_msa.index_select(0, adaln_indices)
-        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask)
+        attn_output = self.attn(norm_hidden_states, rotary_emb, attention_mask, parallel_config=parallel_config)
         hidden_states = residual + gate_msa.index_select(0, adaln_indices) * attn_output
 
         residual = hidden_states
@@ -601,7 +606,8 @@ def forward(
             text_indices (`torch.Tensor` of shape `(num_text_tokens,)`):
                 Positions of the text rows in the packed sequence.
             attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that, if specified, may carry a `scale` entry which is applied to the LoRA layers.
+                May carry a LoRA `scale` and `reuse_ulysses_metadata=True` to reuse packed-sequence shard lengths
+                within this inference forward when using Ulysses Anything context parallelism.
             return_dict (`bool`, defaults to `True`):
                 Whether to return a [`MiniMaxH3TransformerOutput`] instead of a plain tuple.
 
@@ -646,13 +652,29 @@ def forward(
         # 3. Row -> AdaLN table row.
         adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags
 
+        parallel_config = None
+        if attention_kwargs is not None and attention_kwargs.get("reuse_ulysses_metadata", False):
+            parallel_config = self._parallel_config
+            cp_config = parallel_config.context_parallel_config if parallel_config is not None else None
+            if torch.is_grad_enabled() or cp_config is None or not cp_config.ulysses_anything:
+                raise ValueError("Sequence metadata reuse requires inference with Ulysses Anything enabled.")
+            group = cp_config._ulysses_mesh.get_group()
+            global_lengths = gather_size_by_comm(sequence_length, group)
+            if any(length != sequence_length for length in global_lengths):
+                raise ValueError("All Ulysses ranks must agree on the packed sequence length for this forward.")
+            degree = len(global_lengths)
+            shard_lengths = tuple(
+                sequence_length // degree + int(rank < sequence_length % degree) for rank in range(degree)
+            )
+            parallel_config = replace(parallel_config, _ulysses_sequence_lengths=shard_lengths)
+
         for block in self.transformer_blocks:
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, temb, adaln_indices, rotary_emb
                 )
             else:
-                hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb)
+                hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb, parallel_config=parallel_config)
 
         # 5. Both heads run over every row, then the rows of each modality are selected. The heads are listed in
         # `_keep_in_fp32_modules`, so they stay float32 while the block stack runs in the requested `torch_dtype`;

tests/models/transformers/test_models_transformer_minimax_h3.py

@@ -13,9 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import replace
+from datetime import timedelta
+
+import pytest
 import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
 
 from diffusers import MiniMaxH3Transformer3DModel
+from diffusers.models._modeling_parallel import ContextParallelConfig, gather_size_by_comm
+from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput
 from diffusers.utils.torch_utils import randn_tensor
 
@@ -36,6 +44,77 @@
 enable_full_determinism()
 
 
+def _sequence_plan_worker(rank, port):
+    torch.set_num_threads(1)
+    dist.init_process_group(
+        "cpu:gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=2, timeout=timedelta(seconds=60)
+    )
+    try:
+        fixture = MiniMaxH3TransformerTesterConfig()
+        torch.manual_seed(7)
+        model = MiniMaxH3Transformer3DModel(**fixture.get_init_dict()).eval()
+        model.set_attention_backend("native")
+        mesh = dist.device_mesh.init_device_mesh("cpu", (1, 2), mesh_dim_names=("ring", "ulysses"))
+        model.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2, ulysses_anything=True, mesh=mesh))
+        original_config = model._parallel_config
+        with torch.no_grad():
+            for video_tokens in (8, 10, 8):
+                inputs = fixture.get_dummy_inputs(num_video_tokens=video_tokens, batch_size=1)
+                with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as baseline_profile:
+                    expected = model(**inputs, return_dict=False)
+                with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as planned_profile:
+                    actual = model(**inputs, attention_kwargs={"reuse_ulysses_metadata": True}, return_dict=False)
+                for output, reference in zip(actual, expected):
+                    torch.testing.assert_close(output, reference, rtol=0, atol=0)
+                counts = [
+                    sum(event.count for event in profile.key_averages() if "allgather" in event.key.lower())
+                    for profile in (baseline_profile, planned_profile)
+                ]
+                assert counts[1] < counts[0], counts
+                assert model._parallel_config is original_config
+                assert original_config._ulysses_sequence_lengths is None
+
+            # Rank zero keeps five rows across 9 -> 10 -> 9, while rank one's length changes.
+            for sequence_length in (9, 10, 9):
+                lengths = gather_size_by_comm(sequence_length, mesh["ulysses"].get_group())
+                assert lengths == [sequence_length, sequence_length]
+                plan = replace(
+                    original_config,
+                    _ulysses_sequence_lengths=(sequence_length // 2 + sequence_length % 2, sequence_length // 2),
+                )
+                torch.manual_seed(5)
+                q, k, v = (torch.randn(1, sequence_length, 2, 16) for _ in range(3))
+                q, k, v = (x.tensor_split(2, dim=1)[rank].contiguous() for x in (q, k, v))
+                expected = dispatch_attention_fn(q, k, v, backend="native", parallel_config=original_config)
+                actual = dispatch_attention_fn(q, k, v, backend="native", parallel_config=plan)
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+            with pytest.raises(ValueError, match="agree on the packed sequence length"):
+                model(
+                    **fixture.get_dummy_inputs(num_video_tokens=8 + 2 * rank, batch_size=1),
+                    attention_kwargs={"reuse_ulysses_metadata": True},
+                )
+            assert model._parallel_config is original_config
+            assert original_config._ulysses_sequence_lengths is None
+    finally:
+        dist.destroy_process_group()
+
+
+class TestMiniMaxH3SequencePlan:
+    def test_forward_plan_reuses_collectives_without_leaking_between_requests(self):
+        if torch_device != "cpu" or not dist.is_gloo_available():
+            pytest.skip("This focused distributed test uses CPU Gloo.")
+        from ..testing_utils.parallelism import _find_free_port
+
+        mp.spawn(_sequence_plan_worker, args=(_find_free_port(),), nprocs=2, join=True)
+
+    def test_metadata_reuse_requires_inference_context_parallelism(self):
+        fixture = MiniMaxH3TransformerTesterConfig()
+        model = MiniMaxH3Transformer3DModel(**fixture.get_init_dict())
+        with pytest.raises(ValueError, match="inference with Ulysses Anything"):
+            model(**fixture.get_dummy_inputs(), attention_kwargs={"reuse_ulysses_metadata": True})
+
+
 # The packed layout the dummy inputs describe: a text block, then the audio rows, then the target video rows, which
 # is the layout `MiniMaxH3Blocks` builds.
 NUM_TEXT_TOKENS = 4