[Model][Performance] Optimize MiniMax-H3 strict Ulysses boundaries (#6173)

Signed-off-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: mokashliu <mokashliu@tencent.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -36,7 +36,13 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
     assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
     assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
     assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
-    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}
+    assert set(MiniMaxH3DiTModel._sp_plan) == {
+        "sp_prepare",
+        "local_sp_prepare",
+        "sp_gather",
+    }
+    assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
+    assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}
 
     model = object.__new__(MiniMaxH3DiTModel)
     nn.Module.__init__(model)
@@ -89,6 +95,279 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
     torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)
 
 
+def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        _build_rope_table,
+    )
+
+    local_freqs_half = torch.randn(2, 48)
+    local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1)
+
+    actual = _build_rope_table(local_freqs)
+    expected = torch.cat(
+        (torch.cos(local_freqs_half), torch.sin(local_freqs_half)),
+        dim=-1,
+    ).to(torch.bfloat16)
+
+    assert actual.shape == (2, 96)
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
+def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
+    from vllm_omni.diffusion import forward_context
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict")
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1)
+
+    assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2)
+
+
+@pytest.mark.parametrize(
+    ("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"),
+    [
+        (False, "strict", 16, 8, 1, 1),
+        (True, "advanced_uaa", 16, 8, 1, 1),
+        (True, "strict", 17, 8, 1, 1),
+        (True, "strict", 16, 4, 2, 1),
+        (True, "strict", 16, 1, 1, 8),
+    ],
+)
+def test_local_span_falls_back_when_local_embedding_is_unsafe(
+    monkeypatch,
+    hooks_applied,
+    mode,
+    seq_len,
+    ulysses,
+    ring,
+    allgather,
+):
+    from vllm_omni.diffusion import forward_context
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
+    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_ulysses_parallel_world_size",
+        lambda: ulysses,
+    )
+    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring)
+    monkeypatch.setattr(
+        parallel_state,
+        "get_allgather_parallel_world_size",
+        lambda: allgather,
+    )
+
+    assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == (
+        0,
+        seq_len,
+    )
+
+
+def test_local_embedding_spans_reassemble_multimodal_rows():
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    class _IdentityProjection(nn.Module):
+        def forward(self, x):
+            return x, None
+
+    class _IdentityRefiner(nn.Module):
+        def forward(self, x, **kwargs):
+            return x
+
+    model = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.hidden_size = 2
+    model.video_patch_proj = _IdentityProjection()
+    model.audio_patch_proj = _IdentityProjection()
+    model.condition_proj = _IdentityProjection()
+    model.token_refiner = _IdentityRefiner()
+    model.time_embedder = nn.Identity()
+
+    seq_len = 16
+    x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2)
+    audio_x = x + 100
+    text = torch.tensor(
+        [[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]],
+        dtype=torch.bfloat16,
+    )
+    common = {
+        "x": x,
+        "audio_x": audio_x,
+        "text_embeddings_selected": text,
+        "unique_timesteps": torch.tensor([0.2, 0.3]),
+        "img_pos": torch.tensor([1, 2, 6, 7, 11, 12]),
+        "audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]),
+        "text_pos": torch.tensor([0, 5, 10, 15]),
+        "refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32),
+        "refiner_max_seqlen": 4,
+        "seq_len": seq_len,
+        "device": torch.device("cpu"),
+        "local_span": (0, seq_len),
+    }
+
+    full, full_t_emb = model._embed(**common)
+    local_spans = []
+    for rank in range(8):
+        local_span = (rank * (seq_len // 8), seq_len // 8)
+        local, local_t_emb = model._embed(
+            **{
+                **common,
+                "local_span": local_span,
+            }
+        )
+        assert local.shape == (2, 2)
+        torch.testing.assert_close(local_t_emb, full_t_emb)
+        local_spans.append(local)
+
+    torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "local_projection",
+    [False, True],
+    ids=["fallback", "strict-local"],
+)
+def test_sp_output_boundary_preserves_logits_and_minimizes_gather(
+    monkeypatch,
+    local_projection,
+):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
+
+    class _StaticRope(nn.Module):
+        input_rows = None
+
+        def forward(self, position_ids):
+            self.input_rows = position_ids.shape[1]
+            return torch.zeros(position_ids.shape[1], 6)
+
+    class _Prepare(nn.Module):
+        def __init__(self, shard):
+            super().__init__()
+            self.shard = shard
+
+        def forward(self, hidden, rope, combined):
+            local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0]
+            return hidden[:local_len], rope[:local_len], combined[:local_len]
+
+    class _Gather(nn.Module):
+        def __init__(self, remote_rows, replace):
+            super().__init__()
+            self.remote_rows = remote_rows
+            self.replace = replace
+            self.input_width = None
+
+        def forward(self, tensor):
+            self.input_width = tensor.shape[-1]
+            if self.replace:
+                return self.remote_rows
+            return torch.cat((tensor, self.remote_rows), dim=0)
+
+    class _FinalLayer(nn.Module):
+        def forward(self, hidden, *, t_emb, inverse_indices):
+            del t_emb
+            inverse = inverse_indices.to(torch.float32).unsqueeze(-1)
+            return (
+                hidden[:, :2].to(torch.float32) + inverse,
+                hidden[:, 2:3].to(torch.float32) - inverse,
+            )
+
+    seq_len = 8
+    full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4)
+    inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0])
+    final_layer = _FinalLayer()
+    remote_video, remote_audio = final_layer(
+        full_hidden[seq_len // 2 :],
+        t_emb=torch.empty(0),
+        inverse_indices=inverse_indices[seq_len // 2 :],
+    )
+    if local_projection:
+        remote_rows = torch.cat((remote_video, remote_audio), dim=-1)
+    else:
+        remote_rows = full_hidden
+    gather = _Gather(remote_rows, replace=not local_projection)
+
+    model = object.__new__(h3.MiniMaxH3DiTModel)
+    nn.Module.__init__(model)
+    model.arch = h3.MiniMaxH3DiTArchConfig(
+        latents_dim=2,
+        audio_latents_dim=1,
+        patch_size=(1, 1, 1),
+    )
+    rope = _StaticRope()
+    model.rope = rope
+    model.sp_prepare = _Prepare(shard=True)
+    model.local_sp_prepare = _Prepare(shard=False)
+    model.sp_gather = gather
+    model.blocks = nn.ModuleList()
+    model.final_layer = final_layer
+    monkeypatch.setattr(
+        model,
+        "_embed",
+        lambda **kwargs: (
+            full_hidden[: seq_len // 2] if local_projection else full_hidden,
+            torch.tensor([[0.5]], dtype=torch.float32),
+        ),
+    )
+
+    def _local_span(*args, **kwargs):
+        del args, kwargs
+        return (0, seq_len // 2) if local_projection else (0, seq_len)
+
+    monkeypatch.setattr(
+        h3,
+        "_sequence_parallel_local_span",
+        _local_span,
+    )
+
+    img_pos = torch.tensor([0, 2, 4, 6])
+    audio_pos = torch.tensor([1, 3, 5, 7])
+    video, audio = model(
+        x=torch.zeros(1, seq_len, 2),
+        audio_x=torch.zeros(1, seq_len, 1),
+        img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long),
+        unique_timesteps=torch.tensor([0.5]),
+        inverse_indices=inverse_indices,
+        update_mask=torch.ones(img_pos.numel()),
+        token_tags=torch.zeros(seq_len, dtype=torch.long),
+        prompt_embeds=torch.empty(0, 2),
+        img_pos_info={"position_ids": img_pos},
+        audio_pos_info={"position_ids": audio_pos},
+        text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)},
+        img_pos_for_infer_output_info={"position_ids": img_pos},
+        packed_seq_params={
+            "cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32),
+            "max_seqlen_q": seq_len,
+        },
+        refiner_packed_seq_params={
+            "cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32),
+            "max_seqlen_q": 0,
+        },
+    )
+
+    expected_video, expected_audio = final_layer(
+        full_hidden,
+        t_emb=torch.empty(0),
+        inverse_indices=inverse_indices,
+    )
+    torch.testing.assert_close(video, expected_video.index_select(0, img_pos))
+    torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos))
+    expected_gather_width = 3 if local_projection else full_hidden.shape[-1]
+    expected_rope_rows = seq_len // 2 if local_projection else seq_len
+    assert gather.input_width == expected_gather_width
+    assert rope.input_rows == expected_rope_rows
+
+
 @pytest.mark.parametrize(
     ("tp_size", "message"),
     [

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1333,6 +1333,7 @@ def extract_minimax_h3_context(
         refiner_max_seqlen=refiner_max,
         seq_len=seq_len,
         device=device,
+        local_span=(0, seq_len),
     )
 
     combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -107,6 +107,7 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
 # video/text/audio tokens (padding is clamped to 0 before the embedding
 # lookup and masked out afterwards).
 MINIMAX_H3_ADALN_MODALITY_NUM = 3
+_LOCAL_SP_PREPARE_HOOK = "sp_input---local_sp_prepare"
 
 # Opt-in fp16-range protection for the NPU ascend_laser_attention kernel
 # (consumed only via the "laser_input_scale" extra key; other backends and
@@ -212,6 +213,51 @@ def _modulate_gate(
     return (x + gate.index_select(0, indices) * other).to(dtype)
 
 
+def _sequence_parallel_local_span(
+    seq_len: int,
+    *,
+    hooks_applied: bool,
+) -> tuple[int, int]:
+    """Return the packed-row span owned by this sequence-parallel rank."""
+    from vllm_omni.diffusion.forward_context import (
+        get_ulysses_mode,
+        is_forward_context_available,
+    )
+
+    if not hooks_applied or not is_forward_context_available():
+        return 0, seq_len
+    if get_ulysses_mode(default="strict") != "strict":
+        return 0, seq_len
+
+    try:
+        from vllm_omni.diffusion.distributed.parallel_state import (
+            get_allgather_parallel_world_size,
+            get_ring_parallel_world_size,
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+            get_ulysses_parallel_world_size,
+        )
+
+        world_size = int(get_sequence_parallel_world_size())
+        rank = int(get_sequence_parallel_rank())
+        ulysses_world_size = int(get_ulysses_parallel_world_size())
+        ring_world_size = int(get_ring_parallel_world_size())
+        allgather_world_size = int(get_allgather_parallel_world_size())
+    except AssertionError:
+        return 0, seq_len
+
+    if world_size <= 1 or ulysses_world_size != world_size:
+        return 0, seq_len
+    if ring_world_size != 1 or allgather_world_size != 1:
+        return 0, seq_len
+    if seq_len < world_size or seq_len % world_size:
+        return 0, seq_len
+
+    chunk_size = seq_len // world_size
+    start = rank * chunk_size
+    return start, chunk_size
+
+
 class MiniMaxH3Rope(nn.Module):
     """3D rope over (t, h, w); rotates 96 of 128 head dims (rotary_percent 0.75).
 
@@ -888,6 +934,13 @@ class MiniMaxH3DiTModel(nn.Module):
                 split_output=True,
             ),
         },
+        "local_sp_prepare": {
+            2: SequenceParallelInput(
+                split_dim=0,
+                expected_dims=1,
+                split_output=True,
+            ),
+        },
         "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2),
     }
     # The checkpoint already stores qkv and the MLP gate/up as single tensors
@@ -995,6 +1048,7 @@ class MiniMaxH3DiTModel(nn.Module):
             ]
         )
         self.sp_prepare = MiniMaxH3SPPrepare()
+        self.local_sp_prepare = MiniMaxH3SPPrepare()
         self.sp_gather = MiniMaxH3SPGather()
         self.final_layer = MiniMaxH3FinalLayer(
             arch,
@@ -1087,16 +1141,39 @@ class MiniMaxH3DiTModel(nn.Module):
         refiner_max_seqlen: int,
         seq_len: int,
         device: torch.device,
+        local_span: tuple[int, int],
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Build packed multimodal embeddings for the TP=1/SP=1 inference path.
+        """Build this rank's packed multimodal embedding rows.
 
-        Returns (decoder_input [S, H] bf16, t_emb [M, t_dim] fp32).
+        Returns (decoder_input [S_local, H] bf16, t_emb [M, t_dim] fp32).
         """
+        local_start, local_len = local_span
+        local_end = local_start + local_len
+        local_only = local_len != seq_len
+        if local_only:
+            img_mask = (img_pos >= local_start) & (img_pos < local_end)
+            audio_mask = (audio_pos >= local_start) & (audio_pos < local_end)
+            text_mask = (text_pos >= local_start) & (text_pos < local_end)
+            img_global_pos = img_pos[img_mask]
+            audio_global_pos = audio_pos[audio_mask]
+            img_local_pos = img_global_pos - local_start
+            audio_local_pos = audio_global_pos - local_start
+            text_local_pos = text_pos[text_mask] - local_start
+            text_local_indices = torch.nonzero(text_mask, as_tuple=False).view(-1)
+        else:
+            img_global_pos = img_pos
+            audio_global_pos = audio_pos
+            img_local_pos = img_pos
+            audio_local_pos = audio_pos
+            text_local_pos = text_pos
+            text_local_indices = None
+
         # Latent embedders stay fp32 in and out; their outputs are cast to the
         # bf16 sequence dtype only during indexed scattering.
-        x_rows = x.view(-1, x.shape[-1]).index_select(0, img_pos).to(_FP32_DTYPE)
+        x_rows = x.view(-1, x.shape[-1]).index_select(0, img_global_pos).to(_FP32_DTYPE)
         video_embed, _ = self.video_patch_proj(x_rows)
-        audio_rows = audio_x.view(-1, audio_x.shape[-1]).index_select(0, audio_pos).to(_FP32_DTYPE)
+        audio_rows = audio_x.view(-1, audio_x.shape[-1])
+        audio_rows = audio_rows.index_select(0, audio_global_pos).to(_FP32_DTYPE)
         audio_embed, _ = self.audio_patch_proj(audio_rows)
 
         text_rows = text_embeddings_selected.to(device=device, dtype=_BF16_DTYPE)
@@ -1106,11 +1183,29 @@ class MiniMaxH3DiTModel(nn.Module):
             cu_seqlens=refiner_cu_seqlens,
             max_seqlen=refiner_max_seqlen,
         )
+        if text_local_indices is not None:
+            text_embed = text_embed.index_select(0, text_local_indices)
 
-        embeddings = torch.zeros((seq_len, self.hidden_size), device=device, dtype=_BF16_DTYPE)
-        embeddings.index_add_(0, text_pos, text_embed.to(_BF16_DTYPE)[: text_pos.shape[0]])
-        embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]])
-        embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]])
+        embeddings = torch.zeros(
+            (local_len, self.hidden_size),
+            device=device,
+            dtype=_BF16_DTYPE,
+        )
+        embeddings.index_add_(
+            0,
+            text_local_pos,
+            text_embed.to(_BF16_DTYPE)[: text_local_pos.shape[0]],
+        )
+        embeddings.index_add_(
+            0,
+            img_local_pos,
+            video_embed.to(_BF16_DTYPE)[: img_local_pos.shape[0]],
+        )
+        embeddings.index_add_(
+            0,
+            audio_local_pos,
+            audio_embed.to(_BF16_DTYPE)[: audio_local_pos.shape[0]],
+        )
 
         t_emb = self.time_embedder(unique_timesteps)
         return embeddings, t_emb
@@ -1140,7 +1235,6 @@ class MiniMaxH3DiTModel(nn.Module):
         update_mask = _required_kwarg(kwargs, "update_mask")
         token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long)
         skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False))
-
         text_selected = _required_kwarg(kwargs, "prompt_embeds")
 
         img_pos = self._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info")
@@ -1170,8 +1264,18 @@ class MiniMaxH3DiTModel(nn.Module):
         if inverse_indices.shape[0] != seq_len:
             raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
         device = x.device
-        # Compute the packed RoPE table once for all transformer blocks.
-        rope_table = _build_rope_table(self.rope(img_position_ids).to(device))
+        local_sp_registry = getattr(self.local_sp_prepare, "_hook_registry", None)
+        hooks_applied = local_sp_registry is not None
+        if local_sp_registry is not None:
+            local_sp_hook = local_sp_registry.get_hook(_LOCAL_SP_PREPARE_HOOK)
+            hooks_applied = local_sp_hook is not None
+        local_span = _sequence_parallel_local_span(
+            seq_len,
+            hooks_applied=hooks_applied,
+        )
+        local_start, local_len = local_span
+        rope_position_ids = img_position_ids.narrow(1, local_start, local_len)
+        rope_table = _build_rope_table(self.rope(rope_position_ids).to(device))
 
         decoder_input, t_emb = self._embed(
             x=x,
@@ -1185,6 +1289,7 @@ class MiniMaxH3DiTModel(nn.Module):
             refiner_max_seqlen=refiner_max,
             seq_len=seq_len,
             device=device,
+            local_span=local_span,
         )
 
         combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device)
@@ -1195,11 +1300,18 @@ class MiniMaxH3DiTModel(nn.Module):
         block_rope = rope_table
         block_combined = combined_indices
 
-        hidden, block_rope, block_combined = self.sp_prepare(
-            hidden,
-            block_rope,
-            block_combined,
-        )
+        if local_len == seq_len:
+            hidden, block_rope, block_combined = self.sp_prepare(
+                hidden,
+                block_rope,
+                block_combined,
+            )
+        else:
+            hidden, block_rope, block_combined = self.local_sp_prepare(
+                hidden,
+                block_rope,
+                block_combined,
+            )
         for block in self.blocks:
             hidden = block(
                 hidden,
@@ -1211,13 +1323,29 @@ class MiniMaxH3DiTModel(nn.Module):
                 packed_total=seq_len,
                 video_layout=video_layout,
             )
-        hidden = self.sp_gather(hidden)
-
-        video_logits, audio_logits = self.final_layer(
-            hidden,
-            t_emb=t_emb,
-            inverse_indices=inverse_indices,
-        )
+        if local_len == seq_len:
+            hidden = self.sp_gather(hidden)
+            video_logits, audio_logits = self.final_layer(
+                hidden,
+                t_emb=t_emb,
+                inverse_indices=inverse_indices,
+            )
+        else:
+            local_inverse_indices = inverse_indices.narrow(
+                0,
+                local_start,
+                local_len,
+            )
+            video_logits, audio_logits = self.final_layer(
+                hidden,
+                t_emb=t_emb,
+                inverse_indices=local_inverse_indices,
+            )
+            compact_logits = torch.cat((video_logits, audio_logits), dim=-1)
+            compact_logits = self.sp_gather(compact_logits)
+            video_width = self.arch.latents_dim * math.prod(self.arch.patch_size)
+            video_logits = compact_logits[..., :video_width]
+            audio_logits = compact_logits[..., video_width:]
 
         # Select target and condition rows at inference-output positions, then
         # zero the condition rows.