From 7b76b6446d29adc23f563b75f604159ab11ecac8 Mon Sep 17 00:00:00 2001 From: mokeke <103873176+mo-ke-ke@users.noreply.github.com> Date: Sat, 15 Aug 2026 16:27:52 +0800 Subject: [PATCH] [Model][Performance] Optimize MiniMax-H3 strict Ulysses boundaries (#6173) Signed-off-by: mokashliu Co-authored-by: mokashliu Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- .../minimax_h3/test_minimax_h3_parallel.py | 281 +++++++++++++++++- .../diffusion/cache/teacache/extractors.py | 1 + .../minimax_h3/minimax_h3_transformer.py | 174 +++++++++-- 3 files changed, 432 insertions(+), 24 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py index 2b168499e..18345f247 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/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"), [ diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 8689c2c54..93e306ebb 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 659b64bc4..52f9ab589 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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.