From d9980c75e8fe68ba2e6aa84360f109182a75023b Mon Sep 17 00:00:00 2001 From: Wenjia Li Date: Sat, 29 Aug 2026 17:14:52 +0800 Subject: [PATCH] Perf/minimax h3 dit swiglu rope cache npu (#6410) Signed-off-by: Wenjia Li Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- .buildkite/npu/test-npu-ready.yml | 12 + .../cache/test_teacache_extractors.py | 242 +++++++++++++++++- tests/diffusion/layers/test_activation.py | 32 +++ .../layers/test_fused_qk_norm_rope_npu.py | 194 ++++++++++++++ .../minimax_h3/test_minimax_h3_parallel.py | 142 +++++++++- .../diffusion/cache/teacache/extractors.py | 101 ++++++-- vllm_omni/diffusion/cache/teacache/hook.py | 11 +- vllm_omni/diffusion/layers/activation.py | 5 +- .../diffusion/layers/fused_qk_norm_rope.py | 111 +++++++- .../models/minimax_h3/denoise_loop.py | 23 ++ .../minimax_h3/minimax_h3_transformer.py | 80 +++++- 11 files changed, 920 insertions(+), 33 deletions(-) create mode 100644 tests/diffusion/layers/test_activation.py create mode 100644 tests/diffusion/layers/test_fused_qk_norm_rope_npu.py diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 69b6f331b..9aa55e5d0 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -10,6 +10,18 @@ steps: -m 'core_model and npu and A3 and omni' --run-level 'core_model' + - label: "Diffusion NPU layer tests" + key: ready-diffusion-npu-layer-tests + depends_on: upload-ready-pipeline + mirror_hardwares: a3_npu_2 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + tests/diffusion/layers/test_fused_qk_norm_rope_npu.py + -m 'core_model and npu and A3' + --run-level 'core_model' + - group: ":card_index_dividers: Omni Model Test" key: ready-omni-group depends_on: upload-ready-pipeline diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index fdeca8e17..d821541d0 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for TeaCache extractor functions. @@ -18,6 +18,7 @@ Currently implemented: - TestMiniMaxH3Extractor: MiniMaxH3DiTModel extractor """ +import math from abc import ABC, abstractmethod from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -650,6 +651,29 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): torch.testing.assert_close(actual_video, expected_video) torch.testing.assert_close(actual_audio, expected_audio) + def test_extractor_uses_prepared_rope_table( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """TeaCache must consume the denoise-loop RoPE table without rebuilding it.""" + rope_table = minimax_h3_module.prepare_rope_table( + sample_inputs["img_position_ids"], + seq_len=sample_inputs["x"].shape[1], + ) + + def unexpected_rope_build(*args, **kwargs): + raise AssertionError("TeaCache extractor rebuilt the prepared RoPE table") + + monkeypatch.setattr(minimax_h3_module, "prepare_rope_table", unexpected_rope_build) + monkeypatch.setattr(minimax_h3_module.rope, "forward", unexpected_rope_build) + context = extract_minimax_h3_context( + minimax_h3_module, + **{**sample_inputs, "rope_table": rope_table}, + ) + context.run_transformer_blocks() + def test_run_transformer_blocks_forwards_packed_layout( self, minimax_h3_module, @@ -705,6 +729,222 @@ class TestMiniMaxH3Extractor(BaseExtractorTest): assert video_logits.shape[0] == sample_inputs["img_pos_info"]["position_ids"].shape[0] assert audio_logits.shape[0] == sample_inputs["audio_pos_info"]["position_ids"].shape[0] + @pytest.mark.sp + def test_teacache_strict_sp_keeps_cache_state_rank_local( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """Strict SP must not mix local TeaCache state with gathered block rows.""" + from vllm_omni.diffusion.attention.ops import minimax_h3_modulation + from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig + from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook + from vllm_omni.diffusion.distributed import parallel_state + + seq_len = sample_inputs["x"].shape[1] + local_len = seq_len // 2 + hidden_size = minimax_h3_module.hidden_size + video_width = minimax_h3_module.arch.latents_dim * math.prod(minimax_h3_module.arch.patch_size) + audio_width = minimax_h3_module.arch.audio_latents_dim + captured_modulation_indices: list[torch.Tensor] = [] + + class _StrictSPPrepare(nn.Module): + def __init__(self): + super().__init__() + self.global_indices: list[torch.Tensor] = [] + + def forward(self, hidden, rope_table, combined_indices): + assert hidden.shape[0] == local_len + assert rope_table.shape[0] == local_len + self.global_indices.append(combined_indices.clone()) + return hidden, rope_table, combined_indices.narrow(0, 0, local_len) + + class _StrictSPGather(nn.Module): + def __init__(self): + super().__init__() + self.inputs: list[torch.Tensor] = [] + self.outputs: list[torch.Tensor] = [] + + def forward(self, tensor): + self.inputs.append(tensor.clone()) + if tensor.shape[-1] == hidden_size: + raise AssertionError("TeaCache must not gather hidden states in strict SP") + assert tensor.shape == (local_len, video_width + audio_width) + # Model the all-gather of compact logits from every SP rank. + # postprocess() indexes the resulting global layout with + # infer_out_pos/audio_pos, including rows owned by another rank. + gathered = tensor.repeat(seq_len // local_len, 1) + assert gathered.shape == (seq_len, video_width + audio_width) + self.outputs.append(gathered) + return gathered + + class _StrictSPGroup: + def __init__(self): + super().__init__() + self.local_decisions: list[int] = [] + self.ops = [] + + def all_reduce(self, decision, op): + self.local_decisions.append(int(decision.item())) + self.ops.append(op) + # Simulate another rank requiring the second step to execute + # the collective block path even though this rank would cache. + if len(self.local_decisions) == 2: + decision.fill_(1) + return decision + + class _StrictSPBlock(nn.Module): + def __init__(self): + super().__init__() + self.block_indices: list[torch.Tensor] = [] + + def adaln_proj(self, t_emb): + del t_emb + zeros = torch.zeros(3, hidden_size, dtype=torch.float32) + return zeros, zeros + + def norm1(self, hidden): + return hidden.clone() + + def forward(self, hidden, *, combined_indices, rope_table, **kwargs): + del kwargs + assert hidden.shape[0] == local_len + assert rope_table.shape[0] == local_len + self.block_indices.append(combined_indices.clone()) + return hidden + 1 + + class _StrictSPFinalLayer(nn.Module): + def forward(self, hidden, *, t_emb, inverse_indices): + del t_emb, inverse_indices + assert hidden.shape == (local_len, hidden_size) + return hidden[:, :video_width].float(), hidden[:, :audio_width].float() + + local_prepare = _StrictSPPrepare() + gather = _StrictSPGather() + sp_group = _StrictSPGroup() + block = _StrictSPBlock() + minimax_h3_module.local_sp_prepare = local_prepare + minimax_h3_module.sp_gather = gather + minimax_h3_module.blocks = nn.ModuleList([block]) + minimax_h3_module.final_layer = _StrictSPFinalLayer() + monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, local_len)) + monkeypatch.setattr( + minimax_h3_module, + "_embed", + lambda **kwargs: ( + torch.arange(local_len * hidden_size, dtype=torch.float32).reshape(local_len, hidden_size), + torch.zeros(3, hidden_size, dtype=torch.float32), + ), + ) + monkeypatch.setattr( + minimax_h3_module, + "prepare_rope_table", + lambda *args, **kwargs: torch.zeros(local_len, 6, dtype=torch.float32), + ) + monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group) + + original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_ + + def capture_indexed_scale_shift(hidden, shift, scale, indices): + captured_modulation_indices.append(indices.clone()) + return original_indexed_scale_shift(hidden, shift, scale, indices) + + monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift) + + local_img_pos = torch.tensor([0, 1]) + local_audio_pos = torch.tensor([2, 3]) + token_tags = torch.full((seq_len,), -1, dtype=torch.long) + token_tags[local_img_pos] = 0 + token_tags[local_audio_pos] = 2 + inverse_indices = torch.zeros(seq_len, dtype=torch.long) + inputs = { + **sample_inputs, + "inverse_indices": inverse_indices, + "token_tags": token_tags, + "update_mask": torch.ones(local_img_pos.numel(), dtype=torch.bool), + "prompt_embeds": torch.empty(0, 6), + "img_pos_info": {"position_ids": local_img_pos}, + "audio_pos_info": {"position_ids": local_audio_pos}, + "text_pos_info": {"position_ids": torch.empty(0, dtype=torch.long)}, + "img_pos_for_infer_output_info": {"position_ids": local_img_pos}, + "refiner_packed_seq_params": { + "cu_seqlens_q": torch.tensor([0, 0], dtype=torch.int32), + "max_seqlen_q": 0, + }, + } + expected_local_indices = (inverse_indices * 3 + token_tags.clamp(min=0))[:local_len] + + hook = TeaCacheHook( + TeaCacheConfig( + transformer_type="MiniMaxH3DiTModel", + rel_l1_thresh=0.3, + ) + ) + hook.initialize_hook(minimax_h3_module) + + hook.new_forward(minimax_h3_module, **inputs) + hook.new_forward(minimax_h3_module, **inputs) + video_logits, audio_logits = hook.new_forward(minimax_h3_module, **inputs) + state = hook.state_manager.get_state() + + assert state.previous_residual is not None + assert state.previous_residual.shape == (local_len, hidden_size) + assert len(captured_modulation_indices) == 3 + assert all(torch.equal(indices, expected_local_indices) for indices in captured_modulation_indices) + assert sp_group.local_decisions == [1, 0, 0] + assert len(sp_group.ops) == 3 + assert all(op == torch.distributed.ReduceOp.MAX for op in sp_group.ops) + assert len(local_prepare.global_indices) == 2 + assert all(indices.shape == (seq_len,) for indices in local_prepare.global_indices) + assert len(block.block_indices) == 2 + assert all(torch.equal(indices, expected_local_indices) for indices in block.block_indices) + assert len(gather.inputs) == 3 + assert video_logits.shape == (local_img_pos.numel(), video_width) + assert audio_logits.shape == (local_audio_pos.numel(), audio_width) + + @pytest.mark.cpu + def test_teacache_non_strict_sp_gathers_hidden_once( + self, + minimax_h3_module, + sample_inputs, + monkeypatch, + ): + """The non-strict SP cache path must not gather global hidden rows twice.""" + from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig + from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook + + class _Prepare(nn.Module): + def forward(self, hidden, rope_table, combined_indices): + return hidden, rope_table, combined_indices + + class _Gather(nn.Module): + def __init__(self): + super().__init__() + self.inputs: list[torch.Tensor] = [] + + def forward(self, tensor): + self.inputs.append(tensor.clone()) + return tensor + + seq_len = sample_inputs["x"].shape[1] + gather = _Gather() + minimax_h3_module.sp_prepare = _Prepare() + minimax_h3_module.sp_gather = gather + monkeypatch.setattr(minimax_h3_module, "_rope_local_span", lambda _seq_len: (0, seq_len)) + + hook = TeaCacheHook( + TeaCacheConfig( + transformer_type="MiniMaxH3DiTModel", + rel_l1_thresh=0.3, + ) + ) + hook.initialize_hook(minimax_h3_module) + hook.new_forward(minimax_h3_module, **sample_inputs) + + assert len(gather.inputs) == 1 + assert gather.inputs[0].shape == (seq_len, minimax_h3_module.hidden_size) + def test_invalid_module_raises_error(self): """Test that invalid module without blocks raises ValueError.""" invalid_module = Mock() diff --git a/tests/diffusion/layers/test_activation.py b/tests/diffusion/layers/test_activation.py new file mode 100644 index 000000000..c3205502f --- /dev/null +++ b/tests/diffusion/layers/test_activation.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.layers.activation import SiluAndMul + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + + +@pytest.mark.cpu +def test_silu_and_mul_native_matches_packed_reference() -> None: + packed = torch.randn(257, 256, dtype=torch.float32) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up) + + +@hardware_test(res={"npu": "A3"}, num_cards=1) +def test_silu_and_mul_npu_matches_packed_reference() -> None: + packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16) + gate, up = packed.chunk(2, dim=-1) + + torch.testing.assert_close( + SiluAndMul()(packed), + F.silu(gate) * up, + atol=2e-2, + rtol=2e-2, + ) diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py new file mode 100644 index 000000000..35d990afa --- /dev/null +++ b/tests/diffusion/layers/test_fused_qk_norm_rope_npu.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +import types + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +_HEAD_DIM = 128 +_ROTARY_DIM = 96 +_EPS = 1e-5 + + +def _reference( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + half = _ROTARY_DIM // 2 + + def apply(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + x = F.rms_norm(x, (_HEAD_DIM,), weight, _EPS) + cos = rope_table[..., :half].unsqueeze(1) + sin = rope_table[..., half:].unsqueeze(1) + return torch.cat( + ( + x[..., :half] * cos - x[..., half:_ROTARY_DIM] * sin, + x[..., half:_ROTARY_DIM] * cos + x[..., :half] * sin, + x[..., _ROTARY_DIM:], + ), + dim=-1, + ) + + return apply(q, q_weight), apply(k, k_weight) + + +def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + torch.manual_seed(7) + tokens, heads = 11, 4 + q = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16) + k = torch.randn(tokens, heads, _HEAD_DIM, dtype=torch.bfloat16) + q_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16) + k_weight = torch.randn(_HEAD_DIM, dtype=torch.bfloat16) + freqs = torch.randn(tokens, _ROTARY_DIM // 2) + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16) + return q, k, q_weight, k_weight, rope_table + + +@pytest.mark.cpu +def test_npu_qk_norm_rope_uses_torch_npu_fused_primitives_without_mindiesd( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + calls: dict[str, int] = {"rms_norm": 0, "rotary_mul": 0} + + def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float): + calls["rms_norm"] += 1 + return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),) + + def npu_rotary_mul(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + calls["rotary_mul"] += 1 + assert x.shape[-1] == _ROTARY_DIM + assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM) + assert sin.shape == cos.shape + half = _ROTARY_DIM // 2 + return torch.cat( + ( + x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half], + x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:], + ), + dim=-1, + ) + + monkeypatch.setitem( + sys.modules, + "torch_npu", + types.SimpleNamespace(npu_rms_norm=npu_rms_norm, npu_rotary_mul=npu_rotary_mul), + ) + monkeypatch.setattr(fused, "find_spec", lambda _: None) + q, k, q_weight, k_weight, rope_table = _inputs() + + actual_q, actual_k = fused._npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + _ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == {"rms_norm": 2, "rotary_mul": 2} + torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + + +@pytest.mark.cpu +def test_npu_qk_norm_rope_uses_mindiesd_and_normalizes_packed_layout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + calls: dict[str, int] = {"rms_norm": 0, "mindiesd_rope": 0} + + def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float): + calls["rms_norm"] += 1 + return (F.rms_norm(x, (_HEAD_DIM,), gamma, epsilon),) + + def rotary_position_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs): + calls["mindiesd_rope"] += 1 + assert x.dim() == 4 + assert x.shape[0] == 1 + assert x.shape[-1] == _ROTARY_DIM + assert cos.shape == (1, x.shape[1], 1, _ROTARY_DIM) + assert sin.shape == cos.shape + assert kwargs == { + "rotated_mode": "rotated_half", + "head_first": False, + "fused": True, + } + half = _ROTARY_DIM // 2 + return torch.cat( + ( + x[..., :half] * cos[..., :half] - x[..., half:] * sin[..., :half], + x[..., half:] * cos[..., half:] + x[..., :half] * sin[..., half:], + ), + dim=-1, + ) + + monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm)) + monkeypatch.setitem( + sys.modules, + "mindiesd", + types.SimpleNamespace(rotary_position_embedding=rotary_position_embedding), + ) + monkeypatch.setattr(fused, "find_spec", lambda _: object()) + q, k, q_weight, k_weight, rope_table = _inputs() + + actual_q, actual_k = fused._npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + _ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == {"rms_norm": 2, "mindiesd_rope": 2} + torch.testing.assert_close(actual_q, expected_q, atol=0, rtol=0) + torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0) + + +@hardware_test(res={"npu": "A3"}, num_cards=1) +def test_fused_qk_norm_rope_npu_matches_reference(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the public Ascend dispatch with real torch_npu/MindIE operators.""" + from vllm_omni.diffusion.layers import fused_qk_norm_rope as fused + + q, k, q_weight, k_weight, rope_table = (tensor.to("npu") for tensor in _inputs()) + calls = 0 + original_npu_impl = fused._npu_qk_norm_rope + + def capture_npu_impl(*args, **kwargs): + nonlocal calls + calls += 1 + return original_npu_impl(*args, **kwargs) + + monkeypatch.setattr(fused, "_npu_qk_norm_rope", capture_npu_impl) + actual_q, actual_k = fused.fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + _EPS, + head_dim=_HEAD_DIM, + rotary_dim=_ROTARY_DIM, + ) + expected_q, expected_k = _reference(q, k, q_weight, k_weight, rope_table) + + assert calls == 1 + torch.testing.assert_close(actual_q, expected_q, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(actual_k, expected_k, atol=2e-2, rtol=2e-2) 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 18345f247..45f9229d6 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import pytest import torch @@ -113,6 +113,146 @@ def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout(): torch.testing.assert_close(actual, expected, atol=0, rtol=0) +def test_h3_prepared_rope_table_is_rank_local_and_validated(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3 + + class Rope(nn.Module): + def forward(self, position_ids): + freqs_half = position_ids[0].float() + return torch.cat((freqs_half, freqs_half), dim=-1) + + model = object.__new__(h3.MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.arch = h3.MiniMaxH3DiTArchConfig(rope_inv_freq_len=1) + model.rope = Rope() + model.local_sp_prepare = nn.Identity() + monkeypatch.setattr(h3, "_sequence_parallel_local_span", lambda *args, **kwargs: (1, 2)) + + position_ids = torch.arange(12, dtype=torch.long).reshape(1, 4, 3) + actual = model.prepare_rope_table(position_ids, seq_len=4) + expected = h3._build_rope_table(model.rope(position_ids[:, 1:3])) + + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + model._validate_prepared_rope_table(actual, local_len=2, device=torch.device("cpu")) + with pytest.raises(ValueError, match="local_seq_len"): + model._validate_prepared_rope_table(actual[:1], local_len=2, device=torch.device("cpu")) + + +def test_denoise_branch_prepares_rope_table_once_per_npu_run(monkeypatch: pytest.MonkeyPatch): + from vllm_omni.diffusion.models.minimax_h3 import denoise_loop + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + ) + + monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: True) + + packed = { + "seq_len": torch.tensor(3), + "img_pos": torch.tensor([1]), + "audio_pos": torch.tensor([2]), + "text_pos": torch.tensor([0]), + "update_mask": torch.tensor([True]), + "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32), + "img_position_ids": torch.zeros(3, 3, dtype=torch.long), + "latent_grid": torch.tensor([1, 1, 1]), + "video_row_start": torch.tensor(1), + } + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=torch.zeros(1, 2), + token_tags=torch.zeros(3, dtype=torch.long), + device=torch.device("cpu"), + ) + + class Model: + def __init__(self): + self.rope_table = torch.zeros(3, 6, dtype=torch.bfloat16) + self.prepare_calls = 0 + self.forward_calls = 0 + + def prepare_rope_table(self, img_position_ids, *, seq_len): + assert img_position_ids is branch.static_kwargs["img_position_ids"] + assert seq_len == 3 + self.prepare_calls += 1 + return self.rope_table + + def __call__(self, **kwargs): + assert kwargs["rope_table"] is self.rope_table + self.forward_calls += 1 + return torch.zeros(1, 96), torch.zeros(1, 32) + + model = Model() + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=torch.zeros(1, 96), + initial_audio_rows=torch.zeros(1, 32), + keyframe_cond_rows=None, + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.5, 0.0], + device=torch.device("cpu"), + ) + + assert model.prepare_calls == 1 + assert model.forward_calls == 2 + + +def test_denoise_branch_keeps_rope_construction_on_the_reference_path_off_npu( + monkeypatch: pytest.MonkeyPatch, +): + from vllm_omni.diffusion.models.minimax_h3 import denoise_loop + from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MiniMaxH3DenoiseBranch, + minimax_h3_denoise_loop, + ) + + monkeypatch.setattr(denoise_loop.current_omni_platform, "is_npu", lambda: False) + packed = { + "seq_len": torch.tensor(3), + "img_pos": torch.tensor([1]), + "audio_pos": torch.tensor([2]), + "text_pos": torch.tensor([0]), + "update_mask": torch.tensor([True]), + "cu_seqlens": torch.tensor([0, 3, 3], dtype=torch.int32), + "img_position_ids": torch.zeros(3, 3, dtype=torch.long), + "latent_grid": torch.tensor([1, 1, 1]), + "video_row_start": torch.tensor(1), + } + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=torch.zeros(1, 2), + token_tags=torch.zeros(3, dtype=torch.long), + device=torch.device("cpu"), + ) + + class Model: + def __init__(self): + self.forward_calls = 0 + + def prepare_rope_table(self, *_args, **_kwargs): + raise AssertionError("the reference path must not pre-build a rope_table") + + def __call__(self, **kwargs): + assert "rope_table" not in kwargs + self.forward_calls += 1 + return torch.zeros(1, 96), torch.zeros(1, 32) + + model = Model() + minimax_h3_denoise_loop( + model=model, + positive=branch, + initial_video_rows=torch.zeros(1, 96), + initial_audio_rows=torch.zeros(1, 32), + keyframe_cond_rows=None, + sigmas_video=[1.0, 0.5, 0.0], + sigmas_audio=[1.0, 0.5, 0.0], + device=torch.device("cpu"), + ) + + assert model.forward_calls == 2 + + 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 diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 5be8b53ab..7f0242d57 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Model-specific extractors for TeaCache. @@ -13,6 +13,7 @@ all model-specific information needed for generic caching, including preprocessi transformer execution, and postprocessing logic. """ +import math from collections.abc import Callable from dataclasses import dataclass from typing import Any @@ -1263,7 +1264,6 @@ def extract_minimax_h3_context( from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( _FORWARD_SUPPORTED_KWARGS, MINIMAX_H3_ADALN_MODALITY_NUM, - _build_rope_table, _required_kwarg, ) @@ -1320,8 +1320,20 @@ def extract_minimax_h3_context( if inverse_indices.shape[0] != seq_len: raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") device = x.device - - rope_table = _build_rope_table(module.rope(img_position_ids).to(device)) + local_span = module._rope_local_span(seq_len) + local_start, local_len = local_span + rope_table = kwargs.get("rope_table") + if rope_table is None: + rope_table = module.prepare_rope_table( + img_position_ids, + seq_len=seq_len, + ) + else: + module._validate_prepared_rope_table( + rope_table, + local_len=local_len, + device=device, + ) decoder_input, t_emb = module._embed( x=x, @@ -1336,28 +1348,59 @@ def extract_minimax_h3_context( num_requests=num_requests, seq_len=seq_len, device=device, - local_span=(0, seq_len), + local_span=local_span, ) combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) inverse_indices = inverse_indices.to(device) cu_seqlens = cu_seqlens.to(device) + # In strict SP, _embed already returns this rank's packed rows while + # combined_indices still describes the full packed sequence. TeaCache uses + # the first block's modulation as part of its rank-local cache state, so + # the modulation indices must be sliced to the same local layout. Keep the + # full indices below: local_sp_prepare owns the global-to-local conversion + # for transformer-block inputs. + state_combined_indices = combined_indices.narrow(0, local_start, local_len) + if local_len == seq_len: + state_combined_indices = combined_indices + shift_msa, scale_msa, *_ = module.blocks[0].adaln_proj(t_emb) modulated_input = module.blocks[0].norm1(decoder_input) modulated_input = indexed_scale_shift_( modulated_input, shift_msa, scale_msa, - combined_indices, + state_combined_indices, ) - def run_transformer_blocks() -> tuple[torch.Tensor, ...]: - hidden, block_rope, block_combined = module.sp_prepare( - decoder_input, - rope_table, - combined_indices, + def synchronize_cache_decision(local_should_compute: bool) -> bool: + if local_len == seq_len: + return local_should_compute + + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + decision = torch.tensor( + int(local_should_compute), + dtype=torch.int32, + device=device, ) + get_sp_group().all_reduce(decision, op=torch.distributed.ReduceOp.MAX) + return bool(decision.item()) + + def run_transformer_blocks() -> tuple[torch.Tensor, ...]: + if local_len == seq_len: + hidden, block_rope, block_combined = module.sp_prepare( + decoder_input, + rope_table, + combined_indices, + ) + else: + hidden, block_rope, block_combined = module.local_sp_prepare( + decoder_input, + rope_table, + combined_indices, + ) for block in module.blocks: hidden = block( hidden, @@ -1369,15 +1412,36 @@ def extract_minimax_h3_context( packed_total=seq_len, video_layout=video_layout, ) - hidden = module.sp_gather(hidden) + # TeaCache stores residuals as block_output - decoder_input. Strict SP + # must keep both tensors rank-local until postprocess projects local + # rows to compact logits and gathers them exactly once. + if local_len == seq_len: + hidden = module.sp_gather(hidden) return (hidden,) def postprocess(hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - video_logits, audio_logits = module.final_layer( - hidden, - t_emb=t_emb, - inverse_indices=inverse_indices, - ) + if local_len == seq_len: + video_logits, audio_logits = module.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 = module.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=local_inverse_indices, + ) + compact_logits = torch.cat((video_logits, audio_logits), dim=-1) + compact_logits = module.sp_gather(compact_logits) + video_width = module.arch.latents_dim * math.prod(module.arch.patch_size) + video_logits = compact_logits[..., :video_width] + audio_logits = compact_logits[..., video_width:] video_logits = video_logits.index_select(0, infer_out_pos.to(device)) audio_logits = audio_logits.index_select(0, audio_pos.to(device)) if not skip_mask_out_condition: @@ -1396,6 +1460,9 @@ def extract_minimax_h3_context( temb=t_emb, run_transformer_blocks=run_transformer_blocks, postprocess=postprocess, + extra_states={ + "synchronize_cache_decision": synchronize_cache_decision, + }, ) diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index 85a42adbb..b60e17836 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/vllm_omni/diffusion/cache/teacache/hook.py @@ -138,7 +138,16 @@ class TeaCacheHook(ModelHook): state = self.state_manager.get_state() # Decide whether to compute or cache based on modulated input similarity - should_compute = self._should_compute_full_transformer(state, ctx.modulated_input) + local_should_compute = self._should_compute_full_transformer(state, ctx.modulated_input) + sync_cache_decision = (ctx.extra_states or {}).get("synchronize_cache_decision") + if sync_cache_decision is not None: + should_compute = sync_cache_decision(local_should_compute) + # A rank that locally chose the cache path must reset its counter + # when another SP rank requires a full collective block execution. + if should_compute and not local_should_compute: + state.accumulated_rel_l1_distance = 0.0 + else: + should_compute = local_should_compute if not should_compute and state.previous_residual is not None: # ============================================================================ diff --git a/vllm_omni/diffusion/layers/activation.py b/vllm_omni/diffusion/layers/activation.py index 9d51b03d2..97a5a60fc 100644 --- a/vllm_omni/diffusion/layers/activation.py +++ b/vllm_omni/diffusion/layers/activation.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Custom activation functions for diffusion models.""" import torch @@ -31,7 +32,9 @@ class SiluAndMul(CustomOp): return out def forward_npu(self, x: torch.Tensor) -> torch.Tensor: - return self.forward_native(x) + import torch_npu + + return torch_npu.npu_swiglu(x, dim=-1) def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py index 64462d6c2..f7e27675c 100644 --- a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -11,11 +11,14 @@ The public contract is shared by diffusion attention implementations: ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K -or rotary-product intermediates. Unsupported inputs use the eager reference. +or rotary-product intermediates. Ascend composes its RMSNorm and rotary fused +primitives; unsupported inputs use the eager reference. """ from __future__ import annotations +from importlib.util import find_spec + import torch import torch.nn.functional as F from torch.library import Library @@ -23,6 +26,8 @@ from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.utils.torch_utils import direct_register_custom_op +from vllm_omni.platforms import current_omni_platform + _FUSED_HEAD_DIM = 128 _FUSED_ROTARY_DIM = 96 _HEADS_PER_PROGRAM = 8 @@ -130,6 +135,72 @@ def _eager_qk_norm_rope( ) +def _npu_apply_rope_table( + x: torch.Tensor, + rope_table: torch.Tensor, + rotary_dim: int, +) -> torch.Tensor: + """Apply H3's rotary dimensions through Ascend's fused rotary kernel. + + ``mindiesd.rotary_position_embedding`` takes a 4-D BSND tensor, whereas + MiniMax-H3 keeps packed activations as ``[tokens, heads, head_dim]``. + The shared MindIE-SD wrapper normalizes this 3-D layout to BSND and + restores it afterwards. It receives the half-width cos/sin values from + the packed table; the wrapper expands them to H3's non-interleaved 96-D + rotary layout. The 32 non-rotary head dimensions bypass the kernel. + + Some CANN environments package ``torch_npu`` without MindIE-SD. Keep + those deployments on the same Ascend fused rotary primitive rather than + silently falling back to eager elementwise RoPE. + """ + import torch_npu + + half = rotary_dim // 2 + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + cos = rope_table[..., :half] + sin = rope_table[..., half:] + + if find_spec("mindiesd") is not None: + from vllm_omni.diffusion.layers.rope import apply_rotary_emb_mindiesd + + x_rot = apply_rotary_emb_mindiesd( + x_rot, + cos, + sin, + interleaved=False, + half_head_dim=True, + ) + else: + # npu_rotary_mul uses BSND and full rotary-width cos/sin. H3 uses + # NeoX/rotated-half ordering, so duplicate rather than interleave the + # half-width table along the final dimension. + cos = cos.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + sin = sin.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) + x_rot = torch_npu.npu_rotary_mul(x_rot.unsqueeze(0), cos, sin).squeeze(0) + + return torch.cat((x_rot, x_pass), dim=-1) + + +def _npu_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + rotary_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Use Ascend RMSNorm and RoPE fused primitives for MiniMax-H3 DiT.""" + import torch_npu + + q_norm = torch_npu.npu_rms_norm(q, q_weight, epsilon=eps)[0] + k_norm = torch_npu.npu_rms_norm(k, k_weight, epsilon=eps)[0] + return ( + _npu_apply_rope_table(q_norm, rope_table, rotary_dim), + _npu_apply_rope_table(k_norm, rope_table, rotary_dim), + ) + + def _fused_cuda_supported( q: torch.Tensor, k: torch.Tensor, @@ -148,6 +219,24 @@ def _fused_cuda_supported( ) +def _fused_npu_supported( + q: torch.Tensor, + k: torch.Tensor, + head_dim: int, + rotary_dim: int, +) -> bool: + """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied.""" + return ( + current_omni_platform.is_npu() + and q.device.type == "npu" + and k.device.type == "npu" + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and head_dim == _FUSED_HEAD_DIM + and rotary_dim == _FUSED_ROTARY_DIM + ) + + def _launch_fused_rms_norm_rope( x: torch.Tensor, weight: torch.Tensor, @@ -192,6 +281,16 @@ def _fused_qk_norm_rope_impl( head_dim: int, rotary_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: + if _fused_npu_supported(q, k, head_dim, rotary_dim): + return _npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + rotary_dim, + ) if not _fused_cuda_supported(q, k, head_dim, rotary_dim): return _eager_qk_norm_rope( q, @@ -273,6 +372,16 @@ def fused_qk_norm_rope( q_weight = q_weight.contiguous() k_weight = k_weight.contiguous() rope_table = rope_table.contiguous() + if _fused_npu_supported(q, k, head_dim, rotary_dim): + return _npu_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + rotary_dim, + ) if not _fused_cuda_supported(q, k, head_dim, rotary_dim): return _fused_qk_norm_rope_impl( q, diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index af6feb963..f6b92aaf7 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 cfg-distilled full denoise loop. Per step, the positive presentation is forwarded exactly once. Video and audio @@ -19,6 +20,7 @@ from vllm_omni.diffusion.forward_context import ( set_forward_context_denoise_step_idx, set_forward_context_denoise_timestep, ) +from vllm_omni.platforms import current_omni_platform from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, @@ -137,6 +139,22 @@ class MiniMaxH3DenoiseBranch: latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])), ) + def prepare_rope_table(self, model: Any) -> None: + """Materialize the branch-local DiT RoPE table once per denoise run. + + ``img_position_ids`` is immutable for this branch while latents and + timesteps change every scheduler step. Keeping the table in + ``static_kwargs`` makes every model call reuse the exact same BF16 + tensor without extending its lifetime beyond this request branch. + """ + prepare = getattr(model, "prepare_rope_table", None) + if not callable(prepare): + return + self.static_kwargs["rope_table"] = prepare( + self.static_kwargs["img_position_ids"], + seq_len=self.seq_len, + ) + def forward_kwargs( self, *, @@ -253,6 +271,11 @@ def minimax_h3_denoise_loop( raise ValueError("video/audio sigma schedules must have equal length") if len(sigmas_video) < 2: raise ValueError("sigma schedules need at least 2 entries") + # Keep CUDA/CPU on the reference path: its accuracy CI compares the + # generated video against the official artifact. The table reuse is an + # Ascend-specific performance optimization and must not alter that path. + if current_omni_platform.is_npu(): + positive.prepare_rope_table(model) video_rows, audio_rows, cond_anchor, audio_anchor = minimax_h3_prepare_denoise_rows( positive=positive, initial_video_rows=initial_video_rows, 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 6bbc2a11e..5b5d9c213 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -49,6 +49,7 @@ from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -179,6 +180,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "packed_seq_params", "refiner_packed_seq_params", "video_token_layout", + "rope_table", } ) @@ -1139,6 +1141,54 @@ class MiniMaxH3DiTModel(nn.Module): for _, param in self.named_parameters(): param.missing_param_init = "error" + def _rope_local_span(self, seq_len: int) -> tuple[int, int]: + """Return the sequence-parallel rows owned by this DiT rank.""" + 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 + return _sequence_parallel_local_span( + seq_len, + hooks_applied=hooks_applied, + ) + + def prepare_rope_table( + self, + img_position_ids: torch.Tensor, + *, + seq_len: int, + ) -> torch.Tensor: + """Build the static local RoPE table once for one denoise branch. + + A MiniMax-H3 denoise branch reuses its packed position IDs at every + scheduler step. The returned table is local to the current sequence- + parallel rank, and therefore must be built by the model that will + consume it rather than cached globally across requests or ranks. + """ + local_start, local_len = self._rope_local_span(seq_len) + rope_position_ids = img_position_ids.narrow(1, local_start, local_len) + return _build_rope_table(self.rope(rope_position_ids).to(img_position_ids.device)) + + def _validate_prepared_rope_table( + self, + rope_table: torch.Tensor, + *, + local_len: int, + device: torch.device, + ) -> None: + expected_width = 6 * self.arch.rope_inv_freq_len + if rope_table.dim() != 2 or tuple(rope_table.shape) != (local_len, expected_width): + raise ValueError( + "rope_table must be [local_seq_len, rotary_dim] for the current " + f"sequence-parallel rank, got {list(rope_table.shape)}; expected " + f"[{local_len}, {expected_width}]." + ) + if rope_table.device != device: + raise ValueError(f"rope_table device {rope_table.device} must match x device {device}.") + if rope_table.dtype != _BF16_DTYPE: + raise ValueError(f"rope_table must be {_BF16_DTYPE}, got {rope_table.dtype}.") + def post_load_weights(self) -> None: for name, param in self.named_parameters(): if name in MINIMAX_H3_FP32_PARAM_NAMES and param.dtype != _FP32_DTYPE: @@ -1362,18 +1412,26 @@ 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 - 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_span = self._rope_local_span(seq_len) 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)) + rope_table = kwargs.get("rope_table") + if rope_table is None: + if current_omni_platform.is_npu(): + rope_table = self.prepare_rope_table( + img_position_ids, + seq_len=seq_len, + ) + else: + # Keep CUDA/CPU numerically and structurally identical to the + # main-branch reference path used by the H100 accuracy suite. + rope_position_ids = img_position_ids.narrow(1, local_start, local_len) + rope_table = _build_rope_table(self.rope(rope_position_ids).to(device)) + else: + self._validate_prepared_rope_table( + rope_table, + local_len=local_len, + device=device, + ) decoder_input, t_emb = self._embed( x=x,