Perf/minimax h3 dit swiglu rope cache npu (#6410)

Signed-off-by: Wenjia Li <wjialish@gmail.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>

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()

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,
+    )

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)

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