[Model] Fuse MammothModa2 QK norm and RoPE (#7969)

Signed-off-by: levius <2114377220@qq.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>

benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Benchmark MammothModa2's Q/K RMSNorm + adjacent-pair real RoPE.
+
+Example:
+
+    CUDA_VISIBLE_DEVICES=0 python \
+      benchmarks/diffusion/benchmark_mammoth_moda2_qk_norm_rope.py \
+      --tokens 512 4173 8346 --warmup 50 --iters 200
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import os
+import statistics
+from collections.abc import Callable
+
+import torch
+
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+    _fused_cuda_supported,
+    fused_qk_norm_rope,
+)
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.mammoth_moda2.rope_real import (
+    apply_real_rotary_emb,
+)
+
+_HEAD_DIM = 120
+_Q_HEADS = 21
+_KV_HEADS = 7
+_EPS = 1e-5
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--tokens", nargs="+", type=int, default=[512, 4173, 8346])
+    parser.add_argument("--warmup", type=int, default=50)
+    parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--rounds", type=int, default=3)
+    parser.add_argument("--attention-sequences", nargs="*", type=int, default=[])
+    return parser.parse_args()
+
+
+def _version(distribution: str) -> str | None:
+    try:
+        return importlib.metadata.version(distribution)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def _measure(
+    fn: Callable[[], object],
+    warmup: int,
+    iters: int,
+) -> dict[str, float]:
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+
+    torch.accelerator.reset_peak_memory_stats()
+    baseline_bytes = torch.accelerator.memory_allocated()
+    samples = []
+    for _ in range(iters):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(float(start.elapsed_time(end)))
+
+    ordered = sorted(samples)
+    return {
+        "median_ms": statistics.median(ordered),
+        "mean_ms": statistics.mean(ordered),
+        "stddev_ms": statistics.pstdev(ordered),
+        "min_ms": ordered[0],
+        "max_ms": ordered[-1],
+        "p90_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.9))],
+        "extra_peak_mib": (torch.accelerator.max_memory_allocated() - baseline_bytes) / 2**20,
+    }
+
+
+@torch.inference_mode()
+def _run_shape(tokens: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
+    dtype = torch.bfloat16
+    device = torch.device("cuda")
+    torch.manual_seed(42)
+    query = torch.randn(tokens, _Q_HEADS, _HEAD_DIM, device=device, dtype=dtype)
+    key = torch.randn(tokens, _KV_HEADS, _HEAD_DIM, device=device, dtype=dtype)
+    norm_q = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
+    norm_k = RMSNorm(_HEAD_DIM, eps=_EPS).to(device=device, dtype=dtype)
+    norm_q.weight.normal_(mean=1.0, std=0.2)
+    norm_k.weight.normal_(mean=1.0, std=0.2)
+    angles = torch.rand(tokens, _HEAD_DIM // 2, device=device) * (2 * torch.pi)
+    cos = angles.cos().repeat_interleave(2, dim=-1).to(dtype)
+    sin = angles.sin().repeat_interleave(2, dim=-1).to(dtype)
+
+    def native() -> tuple[torch.Tensor, torch.Tensor]:
+        q_out = apply_real_rotary_emb(norm_q(query).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
+        k_out = apply_real_rotary_emb(norm_k(key).unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0)).squeeze(0)
+        return q_out, k_out
+
+    def fused() -> tuple[torch.Tensor, torch.Tensor]:
+        rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1)
+        return fused_qk_norm_rope(
+            query,
+            key,
+            norm_q.weight,
+            norm_k.weight,
+            rope_table,
+            _EPS,
+            head_dim=_HEAD_DIM,
+            rotary_dim=_HEAD_DIM,
+            interleaved=True,
+        )
+
+    expected_q, expected_k = native()
+    actual_q, actual_k = fused()
+    errors = {
+        "q_max_abs": (actual_q.float() - expected_q.float()).abs().max().item(),
+        "q_mean_abs": (actual_q.float() - expected_q.float()).abs().mean().item(),
+        "k_max_abs": (actual_k.float() - expected_k.float()).abs().max().item(),
+        "k_mean_abs": (actual_k.float() - expected_k.float()).abs().mean().item(),
+    }
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)
+
+    functions = {"native": native, "fused": fused}
+    order = ("fused", "native") if fused_first else ("native", "fused")
+    measured = {name: _measure(functions[name], warmup, iters) for name in order}
+    native_stats, fused_stats = measured["native"], measured["fused"]
+    return {
+        "tokens": tokens,
+        "native": native_stats,
+        "fused_including_rope_pack": fused_stats,
+        "median_speedup": native_stats["median_ms"] / fused_stats["median_ms"],
+        "errors": errors,
+    }
+
+
+@torch.inference_mode()
+def _run_attention(seq: int, warmup: int, iters: int, *, fused_first: bool = False) -> dict[str, object]:
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
+
+    torch.manual_seed(42)
+    block = (
+        TransformerBlock(
+            2520,
+            _Q_HEADS,
+            _KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=_EPS,
+            rope_repeats_pairs=True,
+        )
+        .to(device="cuda", dtype=torch.bfloat16)
+        .eval()
+    )
+    block.attn.norm_q.weight.normal_(1.0, 0.2)
+    block.attn.norm_k.weight.normal_(1.0, 0.2)
+    hidden = torch.randn(2, seq, 2520, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, device="cuda", dtype=torch.bool)
+    mask[0, seq - min(300, seq // 2) :] = False
+    mask[1, seq - min(17, seq // 2) :] = False
+    angles = torch.rand(2, seq, _HEAD_DIM // 2, device="cuda") * (2 * torch.pi)
+    # The production producer supplies FP32 tables; each example has its own positions.
+    rotary = angles.cos().repeat_interleave(2, dim=-1), angles.sin().repeat_interleave(2, dim=-1)
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    previous = os.environ.get(env)
+
+    def run():
+        return block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+
+    try:
+        outputs = {}
+        timings = {}
+        order = ("fused", "native") if fused_first else ("native", "fused")
+        for name in order:
+            os.environ[env] = "0" if name == "fused" else str(2 * seq + 1)
+            outputs[name] = run()
+            timings[name] = _measure(run, warmup, iters)
+        diff = (outputs["native"].float() - outputs["fused"].float()).abs()[mask]
+        max_diff, mean_diff = diff.max().item(), diff.mean().item()
+        assert max_diff < 2e-2 and mean_diff < 1e-3, (max_diff, mean_diff)
+        assert torch.count_nonzero(outputs["fused"][~mask]) == 0
+        return {
+            "sequence": seq,
+            "batch": 2,
+            "backend": block.attn.omni_attn.attn_backend.__name__,
+            "native": timings["native"],
+            "fused": timings["fused"],
+            "median_speedup": timings["native"]["median_ms"] / timings["fused"]["median_ms"],
+            "max_abs_diff": max_diff,
+            "mean_abs_diff": mean_diff,
+        }
+    finally:
+        if previous is None:
+            os.environ.pop(env, None)
+        else:
+            os.environ[env] = previous
+
+
+def main() -> None:
+    args = _parse_args()
+    if (
+        args.rounds < 1
+        or args.iters < 1
+        or args.warmup < 0
+        or any(value < 1 for value in [*args.tokens, *args.attention_sequences])
+    ):
+        raise ValueError("rounds, iters and sequence lengths must be positive; warmup must be nonnegative")
+    if not torch.cuda.is_available():
+        raise RuntimeError("This benchmark requires CUDA")
+
+    probe_q = torch.empty(1, _Q_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    probe_k = torch.empty(1, _KV_HEADS, _HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    if not _fused_cuda_supported(probe_q, probe_k, _HEAD_DIM, _HEAD_DIM, interleaved=True):
+        raise RuntimeError("The fused CUDA QK norm/RoPE path is unavailable")
+
+    result = {
+        "environment": {
+            "gpu": torch.cuda.get_device_name(),
+            "torch": torch.__version__,
+            "cuda": torch.version.cuda,
+            "triton": _version("triton"),
+            "vllm": _version("vllm"),
+            "vllm_omni": _version("vllm-omni"),
+        },
+        "config": {
+            "dtype": "bfloat16",
+            "q_heads": _Q_HEADS,
+            "kv_heads": _KV_HEADS,
+            "head_dim": _HEAD_DIM,
+            "warmup": args.warmup,
+            "iters": args.iters,
+            "rounds": args.rounds,
+            "inference_mode": True,
+            "alternating_order": True,
+            "norm_weights": "normal(mean=1.0, std=0.2)",
+        },
+        "results": [
+            {"round": round_id, **_run_shape(tokens, args.warmup, args.iters, fused_first=bool(round_id % 2))}
+            for round_id in range(args.rounds)
+            for tokens in args.tokens
+        ],
+        "attention_results": [
+            {"round": round_id, **_run_attention(seq, args.warmup, args.iters, fused_first=bool(round_id % 2))}
+            for round_id in range(args.rounds)
+            for seq in args.attention_sequences
+        ],
+    }
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

docs/configuration/environment_variables.md

@@ -130,7 +130,7 @@ settings.
 
 | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle |
 | --- | --- | --- | --- | --- |
-| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control |
+| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`; MammothModa2: `0`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control |
 | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary |
 | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control |
 

tests/diffusion/models/mammoth_moda2/test_dit_attention.py

@@ -187,3 +187,69 @@ def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask):
         image_rotary_emb=(angles.cos(), angles.sin()),
     )
     assert out.shape == (BATCH, 0, DIM)
+
+
+def test_model_owns_repeated_pair_rope_contract():
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
+    from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
+
+    with set_current_diffusion_config(_SDPA_CONFIG):
+        model = Transformer2DModel(
+            hidden_size=48,
+            num_attention_heads=6,
+            num_kv_heads=2,
+            num_layers=1,
+            num_refiner_layers=1,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            axes_dim_rope=(2, 2, 4),
+            axes_lens=(8, 8, 8),
+            text_feat_dim=32,
+        )
+    for blocks in (model.layers, model.noise_refiner, model.ref_image_refiner, model.context_refiner):
+        assert all(block.attn.processor.rope_repeats_pairs for block in blocks)
+    assert not _block(2).attn.processor.rope_repeats_pairs
+    freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000)
+    ids = torch.tensor([[[0, 1, 2], [3, 4, 5]], [[1, 0, 3], [2, 5, 4]]])
+    cos, sin = model.rope_embedder._get_freqs_real(freqs, ids)
+    for table in (cos, sin):
+        torch.testing.assert_close(table[..., 0::2], table[..., 1::2], atol=0, rtol=0)
+    assert not torch.equal(cos[0], cos[1])
+
+
+def test_embedding_preparation_preserves_stubbed_rope_tables(monkeypatch):
+    from types import SimpleNamespace
+
+    from torch import nn
+
+    import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit
+
+    # Embedding-only tests can stub the RoPE producer without constructing
+    # attention layers or supplying full transformer configuration.
+    model = object.__new__(mammoth_dit.Transformer2DModel)
+    nn.Module.__init__(model)
+    object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2))
+
+    class CaptionEmbedding(nn.Module):
+        def forward(self, timestep, text, dtype):
+            return torch.zeros(2, 8), text
+
+    model.time_caption_embed = CaptionEmbedding()
+    model.x_embedder = nn.Identity()
+    model.rope_embedder = lambda *args: (None, None, None, None, [4, 4], [8, 8])
+
+    def forbid_gate(*args):
+        raise AssertionError("Embedding preparation must not resolve the fusion token gate")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate)
+    prepared = model._prepare_embeddings(
+        torch.zeros(2, 4, 4, 4),
+        torch.ones(2),
+        torch.zeros(2, 4, 8),
+        torch.ones(2, 4, dtype=torch.bool),
+        torch.zeros(1),
+        2,
+        4,
+        4,
+    )
+    assert prepared[6:9] == (None, None, None)

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared
-backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic
-in bf16, with the padding mask the joint text+image sequence carries."""
+"""Real-shape checks for MammothModa2's fused QK norm/RoPE and shared CUDA
+attention backend against the native shared RMSNorm path."""
 
 import pytest
 import torch
+from diffusers.models.attention_processor import Attention
 
 import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend
+import vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model as mammoth_dit
 from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock
 
 from .test_dit_attention import _reference_attention
@@ -24,11 +26,127 @@ pytestmark = [
 DIM, HEADS, KV_HEADS = 2520, 21, 7
 
 
+def _norm_attention(head_dim, dtype):
+    attn = Attention(query_dim=head_dim * HEADS, heads=HEADS, kv_heads=KV_HEADS, dim_head=head_dim)
+    attn.norm_q = RMSNorm(head_dim)
+    attn.norm_k = RMSNorm(head_dim)
+    return attn.to(device="cuda", dtype=dtype)
+
+
+def test_unpaired_rope_preserves_native_arithmetic(monkeypatch):
+    """A generic block accepts independently valued even and odd RoPE lanes."""
+    torch.manual_seed(42)
+    q = torch.randn(2, 17, HEADS, 120, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(2, 17, KV_HEADS, 120, device="cuda", dtype=torch.bfloat16)
+    attn = _norm_attention(120, torch.bfloat16)
+    angles = torch.rand(2, 17, 120, device="cuda") * 6.283
+    rotary = angles.cos().to(q.dtype), angles.sin().to(q.dtype)
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary)
+        expected = (
+            mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary),
+            mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary),
+        )
+    for actual, want in zip(got, expected):
+        torch.testing.assert_close(actual, want, atol=0, rtol=0)
+
+
+def test_qk_norm_rope_benchmark_measures_inference(monkeypatch):
+    from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark
+
+    calls = []
+
+    def measure(fn, warmup, iters):
+        outputs = fn()
+        assert not torch.is_grad_enabled()
+        assert all(not output.requires_grad for output in outputs)
+        calls.append(fn.__name__)
+        return {"median_ms": 1.0}
+
+    monkeypatch.setattr(benchmark, "_measure", measure)
+    benchmark._run_shape(17, warmup=0, iters=1)
+    assert calls == ["native", "fused"]
+
+
+def _production_rotary(batch, seq, head_dim, dtype):
+    """Mammoth's real RoPE repeats one angle across each adjacent pair."""
+    angles = torch.rand(batch, seq, head_dim // 2, device="cuda") * 6.283
+    return (
+        angles.cos().repeat_interleave(2, dim=-1).to(dtype),
+        angles.sin().repeat_interleave(2, dim=-1).to(dtype),
+    )
+
+
+@pytest.mark.parametrize("force_native", [False, True])
+def test_production_rope_uses_fused_qk_norm_rope(monkeypatch, force_native):
+    torch.manual_seed(0)
+    seq = 512
+    block = (
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
+    mask[0, -37:] = False
+    rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16)
+    calls = []
+    fused_qk_norm_rope = mammoth_dit.fused_qk_norm_rope
+
+    def record_fused(q, k, q_weight, k_weight, rope_table, eps, **kwargs):
+        calls.append((q.shape, k.shape, rope_table.shape, kwargs))
+        return fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, **kwargs)
+
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record_fused)
+    with torch.no_grad():
+        got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary)
+        want = _reference_attention(block.attn, hidden, mask, rotary)
+
+    supported = mammoth_dit._fused_cuda_supported(hidden, hidden, block.head_dim, block.head_dim, interleaved=True)
+    expected_calls = (
+        [
+            (
+                torch.Size((2 * seq, HEADS, block.head_dim)),
+                torch.Size((2 * seq, KV_HEADS, block.head_dim)),
+                torch.Size((2 * seq, block.head_dim)),
+                {"head_dim": block.head_dim, "rotary_dim": block.head_dim, "interleaved": True},
+            )
+        ]
+        if supported
+        else []
+    )
+    assert calls == expected_calls
+    diff = (got.float() - want.float()).abs()[mask]
+    assert diff.max().item() < 2e-2, diff.max().item()
+    assert diff.mean().item() < 1e-3, diff.mean().item()
+
+
 @pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"])
 def test_real_shape_matches_previous_arithmetic_bf16(seq):
     torch.manual_seed(0)
     block = (
-        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
         .cuda()
         .to(torch.bfloat16)
         .eval()
@@ -37,8 +155,7 @@ def test_real_shape_matches_previous_arithmetic_bf16(seq):
     mask = torch.ones(2, seq, dtype=torch.bool, device="cuda")
     mask[0, seq - 300 :] = False
     mask[1, seq - 17 :] = False
-    angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283
-    rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16))
+    rotary = _production_rotary(1, seq, block.head_dim, torch.bfloat16)
     with torch.no_grad():
         got = block.attn(
             hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary
@@ -88,7 +205,15 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
     torch.manual_seed(0)
     dt = torch.float32
     block = (
-        TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5)
+        TransformerBlock(
+            DIM,
+            HEADS,
+            KV_HEADS,
+            multiple_of=256,
+            ffn_dim_multiplier=1.0,
+            norm_eps=1e-5,
+            rope_repeats_pairs=True,
+        )
         .cuda()
         .to(dt)
         .eval()
@@ -122,7 +247,11 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
         calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs))
         return sdpa(query, key, value, **kwargs)
 
+    def reject_fused(*args, **kwargs):
+        raise AssertionError("FP32 must keep Mammoth's native QK norm/RoPE path")
+
     with torch.no_grad(), monkeypatch.context() as patch:
+        patch.setattr(mammoth_dit, "fused_qk_norm_rope", reject_fused)
         patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa)
         patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa)
         patch.setattr(strategy, "pre_attention", prepare)
@@ -147,3 +276,206 @@ def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_g
     assert torch.count_nonzero(got[~mask]) == 0
     diff = (got.float() - want.float()).abs()[mask]
     assert diff.max().item() < 1e-3, diff.max().item()
+
+
+@pytest.mark.parametrize("layout", ["unbatched", "broadcast", "per_example", "strided"])
+@pytest.mark.parametrize("table_dtype", [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("force_native", [False, True])
+def test_paired_rope_with_learned_norm_weights(monkeypatch, layout, table_dtype, force_native):
+    torch.manual_seed(42)
+    batch, seq, head_dim = 2, 19, 120
+    attn = _norm_attention(head_dim, torch.bfloat16)
+    with torch.no_grad():
+        attn.norm_q.weight.normal_(1.0, 0.2)
+        attn.norm_k.weight.normal_(1.0, 0.2)
+    # Mimic views into packed QKV projections, including their nontrivial stride.
+    packed = torch.randn(batch, seq, HEADS + 2 * KV_HEADS, head_dim, device="cuda", dtype=torch.bfloat16)
+    q, k, _ = packed.split((HEADS, KV_HEADS, KV_HEADS), dim=2)
+    rotary = _production_rotary(batch if layout in ("per_example", "strided") else 1, seq, head_dim, table_dtype)
+    if layout == "unbatched":
+        rotary = tuple(t.squeeze(0) for t in rotary)
+    elif layout == "strided":
+        rotary = tuple(t.transpose(0, 1).contiguous().transpose(0, 1) for t in rotary)
+        assert not rotary[0].is_contiguous()
+    calls = []
+    fused = mammoth_dit.fused_qk_norm_rope
+
+    def record(*args, **kwargs):
+        calls.append(args[4].shape)
+        return fused(*args, **kwargs)
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record)
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "0")
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+        expected = (
+            mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), *rotary),
+            mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), *rotary),
+        )
+    supported = mammoth_dit._fused_cuda_supported(q, k, head_dim, head_dim, interleaved=True)
+    assert calls == ([torch.Size((batch * seq, head_dim))] if supported else [])
+    for actual, want in zip(got, expected):
+        torch.testing.assert_close(actual, want.to(actual.dtype), atol=0.0625, rtol=0.02)
+
+
+@pytest.mark.parametrize(
+    "reason", ["support_query", "token_gate", "no_q_norm", "no_k_norm", "epsilon", "fp16", "head_dim", "no_rope"]
+)
+def test_qk_norm_rope_unsupported_inputs_keep_native_chain(monkeypatch, reason):
+    torch.manual_seed(42)
+    dtype = torch.float16 if reason == "fp16" else torch.bfloat16
+    dim = 260 if reason == "head_dim" else 120
+    attn = _norm_attention(dim, dtype)
+    q = torch.randn(2, 19, HEADS, dim, device="cuda", dtype=dtype)
+    k = torch.randn(2, 19, KV_HEADS, dim, device="cuda", dtype=dtype)
+    rotary = _production_rotary(2, 19, dim, dtype) if reason != "no_rope" else None
+    if reason == "support_query":
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    elif reason == "no_q_norm":
+        attn.norm_q = None
+    elif reason == "no_k_norm":
+        attn.norm_k = None
+    elif reason == "epsilon":
+        attn.norm_k.variance_epsilon = 1e-3
+    monkeypatch.setenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "39" if reason == "token_gate" else "0")
+
+    def reject(*args, **kwargs):
+        raise AssertionError(f"Unsupported {reason} must keep the native chain")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", reject)
+    with torch.no_grad():
+        got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+        expected_q = attn.norm_q(q) if attn.norm_q is not None else q
+        expected_k = attn.norm_k(k) if attn.norm_k is not None else k
+        if rotary is not None:
+            expected_q = mammoth_dit.apply_real_rotary_emb(expected_q, *rotary)
+            expected_k = mammoth_dit.apply_real_rotary_emb(expected_k, *rotary)
+    torch.testing.assert_close(got[0], expected_q, atol=0, rtol=0)
+    torch.testing.assert_close(got[1], expected_k, atol=0, rtol=0)
+
+
+def test_attention_benchmark_single_token_and_restores_gate(monkeypatch):
+    from benchmarks.diffusion import benchmark_mammoth_moda2_qk_norm_rope as benchmark
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    monkeypatch.setenv(env, "123")
+    result = benchmark._run_attention(1, warmup=0, iters=1)
+    assert result["sequence"] == 1
+    assert benchmark.os.environ[env] == "123"
+
+
+@pytest.mark.parametrize("packed", [False, True], ids=["gated_native", "reused_packed"])
+def test_prepared_rope_table_skips_per_layer_gate_and_packing(monkeypatch, packed):
+    torch.manual_seed(7)
+    batch, seq, dim = 2, 19, 120
+    attn = _norm_attention(dim, torch.bfloat16)
+    q = torch.randn(batch, seq, HEADS, dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(batch, seq, KV_HEADS, dim, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _production_rotary(batch, seq, dim, torch.float32)
+    table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1).reshape(batch * seq, dim) if packed else None
+    rotary = cos, sin, table
+    calls = []
+    fused = mammoth_dit.fused_qk_norm_rope
+
+    def record(*args, **kwargs):
+        calls.append(args[4])
+        return fused(*args, **kwargs)
+
+    def forbid_gate(*args, **kwargs):
+        raise AssertionError("Prepared RoPE must not resolve the token gate per attention layer")
+
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope", record)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", forbid_gate)
+    with torch.no_grad():
+        for _ in range(2):
+            got = mammoth_dit._apply_qk_norm_rope(attn, q, k, rotary, rope_repeats_pairs=True)
+            expected = (
+                mammoth_dit.apply_real_rotary_emb(attn.norm_q(q), cos, sin),
+                mammoth_dit.apply_real_rotary_emb(attn.norm_k(k), cos, sin),
+            )
+            for actual, want in zip(got, expected):
+                torch.testing.assert_close(
+                    actual, want.to(actual.dtype), atol=0.0625 if packed else 0, rtol=0.02 if packed else 0
+                )
+    supported = packed and mammoth_dit._fused_cuda_supported(q, k, dim, dim, interleaved=True)
+    assert calls == ([table, table] if supported else [])
+
+
+@pytest.mark.parametrize("force_native", [False, True])
+def test_model_prepares_three_rope_tables_with_one_threshold_lookup(monkeypatch, force_native):
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel
+    from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
+
+    torch.manual_seed(7)
+    model = (
+        Transformer2DModel(
+            patch_size=2,
+            in_channels=4,
+            hidden_size=48,
+            num_layers=1,
+            num_refiner_layers=1,
+            num_attention_heads=6,
+            num_kv_heads=2,
+            multiple_of=8,
+            ffn_dim_multiplier=1.0,
+            axes_dim_rope=(2, 2, 4),
+            axes_lens=(8, 8, 8),
+            text_feat_dim=32,
+        )
+        .cuda()
+        .to(torch.bfloat16)
+        .eval()
+    )
+    hidden = torch.randn(2, 4, 4, 4, device="cuda", dtype=torch.bfloat16)
+    text = torch.randn(2, 4, 32, device="cuda", dtype=torch.bfloat16)
+    mask = torch.ones(2, 4, device="cuda", dtype=torch.bool)
+    freqs = RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (8, 8, 8), 10000)
+    calls = []
+
+    def threshold(default):
+        calls.append(default)
+        return 9
+
+    if force_native:
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+    supported = mammoth_dit._fused_cuda_supported(hidden, hidden, 8, 8, interleaved=True)
+    tables: list[mammoth_dit.RotaryEmbedding] = []
+    refiners = model._apply_refiners
+    layers = model._apply_transformer_layers
+
+    def record_refiners(text, text_mask, context, img, img_mask, noise, temb):
+        tables.extend((context, noise))
+        return refiners(text, text_mask, context, img, img_mask, noise, temb)
+
+    def record_layers(hidden, mask, joint, temb):
+        tables.append(joint)
+        return layers(hidden, mask, joint, temb)
+
+    monkeypatch.setattr(model, "_apply_refiners", record_refiners)
+    monkeypatch.setattr(model, "_apply_transformer_layers", record_layers)
+    monkeypatch.setattr(mammoth_dit, "fused_qk_norm_rope_min_tokens", threshold)
+    args = (hidden, torch.ones(2, device="cuda"), text, freqs, mask)
+    with torch.no_grad():
+        fused_output = model(*args)
+    context, noise, joint = tables
+    assert calls == ([0] if supported else [])
+    assert all(len(rotary) == 3 for rotary in (context, noise, joint))
+    assert context[2] is None and noise[2] is None
+    if supported:
+        assert joint[2].shape == (2 * 8, 8)
+    else:
+        assert joint[2] is None
+    for rotary in (context, noise, joint):
+        if rotary[2] is not None:
+            expected = torch.cat((rotary[0][..., 0::2], rotary[1][..., 0::2]), dim=-1).reshape(-1, 8).float()
+            torch.testing.assert_close(rotary[2], expected, atol=0, rtol=0)
+
+    # Exercise the prepared tuple through every real refiner and joint block,
+    # then compare the same model forward with the unsupported-device fallback.
+    with torch.no_grad():
+        monkeypatch.setattr(mammoth_dit, "_fused_cuda_supported", lambda *args, **kwargs: False)
+        native_output = model(*args)
+    assert fused_output.shape == hidden.shape
+    torch.testing.assert_close(fused_output, native_output, atol=0.0625, rtol=0.02)

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -12,10 +12,112 @@ from torch import nn
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+    _fused_cuda_supported,
+    fused_qk_norm_rope,
+    fused_qk_norm_rope_min_tokens,
+)
 from vllm_omni.diffusion.layers.norm import RMSNorm
 
 from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb
 
+# On the measured RTX 4090, fusion won even for the shortest tested tables;
+# deployments can tune the crossover with VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS.
+_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS = 0
+
+RotaryEmbedding = tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]
+
+
+def _with_packed_rope_table(rotary_emb: tuple[torch.Tensor, torch.Tensor], min_tokens: int | None) -> RotaryEmbedding:
+    """Pack a model-produced adjacent-pair RoPE table once for all layers.
+
+    A third element of ``None`` records that the support or token gate chose
+    the eager path, so attention layers need not resolve that gate again.
+    """
+    cos, sin = rotary_emb
+    if min_tokens is None or cos.shape[0] * cos.shape[1] < min_tokens:
+        return cos, sin, None
+    packed = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1)
+    return cos, sin, packed.reshape(-1, cos.shape[-1]).float()
+
+
+def _apply_qk_norm_rope(
+    attn: Attention,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    image_rotary_emb: RotaryEmbedding | None,
+    *,
+    rope_repeats_pairs: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply Mammoth's per-head Q/K norm and adjacent-pair real RoPE."""
+    batch_size, sequence_length, _, head_dim = query.shape
+    cos, sin = (image_rotary_emb[0], image_rotary_emb[1]) if image_rotary_emb is not None else (None, None)
+    prepared = image_rotary_emb is not None and len(image_rotary_emb) == 3
+    packed = image_rotary_emb[2] if prepared else None
+    use_fused = (
+        rope_repeats_pairs
+        and cos is not None
+        and sin is not None
+        and attn.norm_q is not None
+        and attn.norm_k is not None
+        and attn.norm_q.variance_epsilon == attn.norm_k.variance_epsilon
+        and key.shape[0] == batch_size
+        and key.shape[1] == sequence_length
+        and key.shape[-1] == head_dim
+        and _fused_cuda_supported(query, key, head_dim, head_dim, interleaved=True)
+        and cos.shape[-2:] == (sequence_length, head_dim)
+        and sin.shape == cos.shape
+        and cos.ndim in (2, 3)
+        and (cos.ndim == 2 or cos.shape[0] in (1, batch_size))
+        and cos.device == query.device
+        and cos.dtype in (query.dtype, torch.float32)
+        and sin.device == query.device
+        and sin.dtype == cos.dtype
+        and (
+            packed is not None
+            if prepared
+            else batch_size * sequence_length >= fused_qk_norm_rope_min_tokens(_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS)
+        )
+    )
+    if use_fused:
+        assert cos is not None and sin is not None
+        if prepared:
+            assert packed is not None
+            rope_table = packed
+        else:
+            if cos.ndim == 2:
+                cos, sin = cos.unsqueeze(0), sin.unsqueeze(0)
+            if cos.shape[0] == 1 and batch_size > 1:
+                cos = cos.expand(batch_size, -1, -1)
+                sin = sin.expand(batch_size, -1, -1)
+            rope_table = torch.cat((cos[..., 0::2], sin[..., 0::2]), dim=-1).reshape(
+                batch_size * sequence_length, head_dim
+            )
+        query, key = fused_qk_norm_rope(
+            query.reshape(batch_size * sequence_length, query.shape[2], head_dim),
+            key.reshape(batch_size * sequence_length, key.shape[2], head_dim),
+            attn.norm_q.weight,
+            attn.norm_k.weight,
+            rope_table,
+            attn.norm_q.variance_epsilon,
+            head_dim=head_dim,
+            rotary_dim=head_dim,
+            interleaved=True,
+        )
+        return (
+            query.reshape(batch_size, sequence_length, query.shape[1], head_dim),
+            key.reshape(batch_size, sequence_length, key.shape[1], head_dim),
+        )
+
+    if attn.norm_q is not None:
+        query = attn.norm_q(query)
+    if attn.norm_k is not None:
+        key = attn.norm_k(key)
+    if image_rotary_emb is not None:
+        query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1])
+        key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1])
+    return query, key
+
 
 class LuminaRMSNormZero(nn.Module):
     """
@@ -275,13 +377,18 @@ class AttnProcessor:
     head count first.
     """
 
+    def __init__(self, *, rope_repeats_pairs: bool = False) -> None:
+        # Only the model's RotaryPosEmbedReal producer guarantees equal adjacent
+        # lanes. Generic processors retain the native four-lane arithmetic.
+        self.rope_repeats_pairs = rope_repeats_pairs
+
     def __call__(
         self,
         attn: Attention,
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         attention_mask: torch.Tensor | None = None,
-        image_rotary_emb: torch.Tensor | None = None,
+        image_rotary_emb: RotaryEmbedding | None = None,
     ) -> torch.Tensor:
         batch_size, sequence_length, _ = hidden_states.shape
 
@@ -307,14 +414,7 @@ class AttnProcessor:
         key = key.view(batch_size, -1, kv_heads, head_dim)
         value = value.view(batch_size, -1, kv_heads, head_dim)
 
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-
-        if image_rotary_emb is not None:
-            query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1])
-            key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1])
+        query, key = _apply_qk_norm_rope(attn, query, key, image_rotary_emb, rope_repeats_pairs=self.rope_repeats_pairs)
 
         query, key = query.to(dtype), key.to(dtype)
 
@@ -347,13 +447,15 @@ class TransformerBlock(nn.Module):
         ffn_dim_multiplier: float,
         norm_eps: float,
         modulation: bool = True,
+        *,
+        rope_repeats_pairs: bool = False,
     ) -> None:
         """Initialize the transformer block."""
         super().__init__()
         self.head_dim = dim // num_attention_heads
         self.modulation = modulation
 
-        processor = AttnProcessor()
+        processor = AttnProcessor(rope_repeats_pairs=rope_repeats_pairs)
 
         # Initialize attention layer
         self.attn = Attention(
@@ -403,7 +505,7 @@ class TransformerBlock(nn.Module):
         self,
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
-        image_rotary_emb: torch.Tensor,
+        image_rotary_emb: RotaryEmbedding,
         temb: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if self.modulation:
@@ -505,6 +607,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                     ffn_dim_multiplier,
                     norm_eps,
                     modulation=True,
+                    rope_repeats_pairs=True,
                 )
                 for _ in range(num_refiner_layers)
             ]
@@ -520,6 +623,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                     ffn_dim_multiplier,
                     norm_eps,
                     modulation=True,
+                    rope_repeats_pairs=True,
                 )
                 for _ in range(num_refiner_layers)
             ]
@@ -535,6 +639,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                     ffn_dim_multiplier,
                     norm_eps,
                     modulation=False,
+                    rope_repeats_pairs=True,
                 )
                 for _ in range(num_refiner_layers)
             ]
@@ -551,6 +656,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                     ffn_dim_multiplier,
                     norm_eps,
                     modulation=True,
+                    rope_repeats_pairs=True,
                 )
                 for _ in range(num_layers)
             ]
@@ -705,10 +811,10 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         self,
         text_hidden_states: torch.Tensor,
         text_attention_mask: torch.Tensor,
-        context_rotary_emb: torch.Tensor,
+        context_rotary_emb: RotaryEmbedding,
         img_tokens: torch.Tensor,
         img_mask: torch.Tensor,
-        noise_rotary_emb: torch.Tensor,
+        noise_rotary_emb: RotaryEmbedding,
         temb: torch.Tensor,
     ):
         for layer in self.context_refiner:
@@ -777,6 +883,20 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
             ar_image_attention_mask,
         )
 
+        # The producer repeats each frequency across adjacent lanes. Resolve
+        # support and the token gate once, then share each of the three packed
+        # tables across its refiner or transformer layers. Keep preparation in
+        # forward so embedding-only callers retain the producer's return values.
+        head_dim = rotary_emb[0].shape[-1]
+        min_tokens = (
+            fused_qk_norm_rope_min_tokens(_MAMMOTH_FUSED_QK_NORM_ROPE_MIN_TOKENS)
+            if _fused_cuda_supported(hidden_states, hidden_states, head_dim, head_dim, interleaved=True)
+            else None
+        )
+        context_rotary_emb = _with_packed_rope_table(context_rotary_emb, min_tokens)
+        noise_rotary_emb = _with_packed_rope_table(noise_rotary_emb, min_tokens)
+        rotary_emb = _with_packed_rope_table(rotary_emb, min_tokens)
+
         text_hidden_states, img_tokens = self._apply_refiners(
             text_hidden_states,
             text_attention_mask,