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>

.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

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

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

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:
             # ============================================================================

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)

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,

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,

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,