[Kernel][Boogu-Image] Fuse Q/K RMSNorm + interleaved RoPE via fused_qk_norm_rope (#6982)

Signed-off-by: Qihan Kang <rollykanggg@gmail.com>

docs/configuration/environment_variables.md

@@ -112,6 +112,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_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/config/test_environment_variables.py

@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     """Make an inventory expansion an explicit review decision."""
     category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
     assert category_counts == {
-        EnvironmentVariableCategory.PUBLIC_OMNI: 24,
+        EnvironmentVariableCategory.PUBLIC_OMNI: 25,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 63,

tests/diffusion/layers/test_fused_qk_norm_rope.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
@@ -80,3 +80,96 @@ def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
 
     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)
+
+
+# ---------------------------------------------------------------------------
+# General geometry (any even head_dim, here Boogu-Image's 120) in both
+# pairing modes, against the module's own eager reference at the same
+# tolerance as the MiniMax-H3 test above.
+# ---------------------------------------------------------------------------
+
+_BOOGU_HEAD_DIM = 120
+
+
+def _boogu_inputs(strided: bool):
+    torch.manual_seed(11)
+    if strided:
+        # Slices of a wider head axis: the op must honour q/k strides
+        # (merged-QKV projections hand the op such views).
+        q = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :28]
+        k = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :7]
+    else:
+        q = torch.randn(4139, 28, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn(4139, 7, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    q_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    freqs = torch.randn(4139, _BOOGU_HEAD_DIM // 2, device="cuda", dtype=torch.float32)
+    rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1)
+    return q, k, q_weight, k_weight, rope_table
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("strided", [False, True])
+def test_fused_qk_norm_rope_interleaved(strided):
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        _eager_qk_norm_rope,
+        _launch_fused_qk_norm_rope,
+        fused_qk_norm_rope,
+    )
+
+    q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided)
+    expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM, True)
+    # Check the public op AND the launcher directly: the latter cannot fall
+    # back to eager, so a silent dispatch regression cannot go green here.
+    for actual in (
+        fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True),
+        _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True),
+    ):
+        torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02)
+        torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+def test_fused_qk_norm_rope_half_split_general_dim():
+    """head_dim=120 off the MiniMax-H3 128 pin, half-split pairing.
+
+    Exercises the generalised combined kernel's half-split mode via the
+    launcher. It is not routed in production (half-split traffic keeps the
+    untouched pre-existing per-tensor kernel and its 128/96 contract, so the
+    public op falls back to eager at this geometry), pending the
+    maintainers' call.
+    """
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        _eager_qk_norm_rope,
+        _launch_fused_qk_norm_rope,
+    )
+
+    q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided=False)
+    expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM)
+    actual = _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=False)
+    torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02)
+
+
+def test_fused_qk_norm_rope_min_tokens_resolution(monkeypatch):
+    """The op-level token-gate override: unset or blank -> caller's default; a
+    non-negative integer string overrides it (``0`` = always fuse); anything
+    else is rejected naming the variable."""
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope_min_tokens
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    monkeypatch.delenv(env, raising=False)
+    assert fused_qk_norm_rope_min_tokens(2048) == 2048
+    for blank in ("", "  "):
+        monkeypatch.setenv(env, blank)
+        assert fused_qk_norm_rope_min_tokens(2048) == 2048
+    monkeypatch.setenv(env, "0")
+    assert fused_qk_norm_rope_min_tokens(2048) == 0
+    monkeypatch.setenv(env, "4096")
+    assert fused_qk_norm_rope_min_tokens(2048) == 4096
+    for bad in ("-1", "abc"):
+        monkeypatch.setenv(env, bad)
+        with pytest.raises(ValueError, match=env):
+            fused_qk_norm_rope_min_tokens(2048)

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -26,7 +26,9 @@ NORM_EPS = 1e-5
 def _init_distributed(request):
     """Initialize the minimal single-rank distributed environment required by
     the vLLM parallel linear layers (tensor-parallel group must exist)."""
-    if request.node.name.startswith("test_real_rotary_"):
+    if request.node.name.startswith(
+        ("test_real_rotary_", "test_packed_rope_table", "test_fused_qk_norm_rope", "test_fallback_with_no")
+    ):
         yield
         return
 
@@ -600,3 +602,152 @@ def test_transformer_forward_ti2i_shape():
 
     assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
     assert torch.isfinite(out).all()
+
+
+# ---------------------------------------------------------------------------
+# The three tests below are for the fused qk-norm+RoPE kernel's wiring into
+# this transformer: the packed-table layout conversion, the fused path against
+# the previous eager chain, and the fallback when no packed table is present.
+# They need CUDA (the module-level marks are CPU; these add the cuda mark and
+# skip on CPU-only runners) and real Boogu geometry rather than the mock dims
+# above.
+# ---------------------------------------------------------------------------
+
+_FUSED_HEAD_DIM = 120
+_FUSED_Q_HEADS = 28
+_FUSED_KV_HEADS = 7
+_FUSED_EPS = 1e-5
+
+
+def _fused_rotary_pair(tokens: int):
+    theta = torch.randn(1, tokens, _FUSED_HEAD_DIM // 2, device="cuda", dtype=torch.float32)
+    return (
+        torch.cos(theta).repeat_interleave(2, dim=-1),
+        torch.sin(theta).repeat_interleave(2, dim=-1),
+    )
+
+
+def _fused_norms():
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.layers.layernorm import RMSNorm
+
+    with set_current_vllm_config(VllmConfig()):
+        norm_q = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16)
+        norm_k = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16)
+    with torch.no_grad():
+        norm_q.weight.copy_(torch.randn(_FUSED_HEAD_DIM))
+        norm_k.weight.copy_(torch.randn(_FUSED_HEAD_DIM))
+    return norm_q, norm_k
+
+
+def _operand_ulp_bound(x, weight, cos, sin):
+    """One bf16 rounding of a normalized operand, propagated through RoPE."""
+    import torch.nn.functional as F
+
+    n = F.rms_norm(x, (_FUSED_HEAD_DIM,), weight, _FUSED_EPS).float().abs()
+    c = cos[..., ::2].float().abs().unsqueeze(2)
+    s = sin[..., ::2].float().abs().unsqueeze(2)
+    even_mag = n[..., ::2] * c + n[..., 1::2] * s
+    odd_mag = n[..., ::2] * s + n[..., 1::2] * c
+    return 2.0**-6 * torch.stack((even_mag, odd_mag), dim=-1).flatten(-2) + 1e-6
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_packed_rope_table_layout(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _with_packed_rope_table,
+    )
+
+    monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False)
+    torch.manual_seed(0)
+    tokens = 2311  # above _FUSED_MIN_TOKENS
+    cos, sin = _fused_rotary_pair(tokens)
+    out = _with_packed_rope_table((cos, sin))
+    assert len(out) == 3
+    assert out[0] is cos and out[1] is sin
+    packed = out[2]
+    assert packed.shape == (tokens, _FUSED_HEAD_DIM)
+    assert packed.dtype == torch.float32 and packed.is_contiguous()
+    half = _FUSED_HEAD_DIM // 2
+    assert torch.equal(packed[:, :half], cos[0, :, ::2])
+    assert torch.equal(packed[:, half:], sin[0, :, ::2])
+    # Below _FUSED_MIN_TOKENS the tuple passes through unchanged (host-bound
+    # regime: the fused path would cost more than it saves).
+    short = _with_packed_rope_table(_fused_rotary_pair(17))
+    assert len(short) == 2
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fused_qk_norm_rope_matches_eager_chain():
+    """The fused path vs the previous eager chain at real Boogu shapes.
+
+    Every element must sit within one bf16 ulp of a normalized operand
+    propagated through the rotation (the two paths legitimately differ by
+    single- vs double-rounding of the norm's weight multiply).
+    """
+    from vllm.triton_utils import HAS_TRITON
+
+    if not HAS_TRITON:
+        pytest.skip("Triton required")
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _qk_norm_rope,
+        _with_packed_rope_table,
+    )
+
+    torch.manual_seed(0)
+    norm_q, norm_k = _fused_norms()
+    q = torch.randn(1, 4139, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 4139, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _fused_rotary_pair(4139)
+    with_table = _with_packed_rope_table((cos, sin))
+
+    fused = _qk_norm_rope(q, k, norm_q, norm_k, with_table, torch.bfloat16)
+    eager = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16)
+    for fused_t, eager_t, x, w in zip(fused, eager, (q, k), (norm_q.weight, norm_k.weight)):
+        bound = _operand_ulp_bound(x, w, cos, sin)
+        diff = (fused_t.float() - eager_t.float()).abs()
+        assert (diff <= bound).all(), "fused path beyond one operand ulp of the eager chain"
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_fallback_with_no_packed_table():
+    """Without a packed table the helper is bit-exact to the old chain."""
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _qk_norm_rope,
+        apply_rotary_emb,
+    )
+
+    torch.manual_seed(0)
+    norm_q, norm_k = _fused_norms()
+    q = torch.randn(1, 64, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 64, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    cos, sin = _fused_rotary_pair(64)
+
+    out_q, out_k = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16)
+    ref_q = apply_rotary_emb(norm_q(q), (cos, sin)).to(torch.bfloat16)
+    ref_k = apply_rotary_emb(norm_k(k), (cos, sin)).to(torch.bfloat16)
+    assert torch.equal(out_q, ref_q)
+    assert torch.equal(out_k, ref_k)
+
+
+def test_packed_rope_table_env_override(monkeypatch):
+    """VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS steers the packing gate: 0 packs
+    even a 43-token table, a huge value leaves a 4139-token table untouched,
+    unset keeps the 2048 default routing."""
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        _with_packed_rope_table,
+    )
+
+    env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+    short, long = _identity_rotary_emb(1, 43), _identity_rotary_emb(1, 4139)
+
+    monkeypatch.setenv(env, "0")
+    assert len(_with_packed_rope_table(short)) == 3
+    monkeypatch.setenv(env, str(10**9))
+    assert _with_packed_rope_table(long) is long
+    monkeypatch.delenv(env)
+    assert len(_with_packed_rope_table(long)) == 3
+    assert _with_packed_rope_table(short) is short

vllm_omni/config/environment_variable_inventory.py

@@ -65,6 +65,7 @@ _PUBLIC_OMNI = (
     "SPEAKER_SAMPLES_DIR",
     "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT",
     "VLLM_OMNI_EVENT_DRIVEN_ORCH",
+    "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS",
     "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S",
     "VLLM_OMNI_ORCH_MONITOR_PATH",
     "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY",

vllm_omni/diffusion/envs.py

@@ -15,6 +15,7 @@ if TYPE_CHECKING:
     MASTER_PORT: int | None = None
     CUDA_HOME: str | None = None
     LOCAL_RANK: int = 0
+    VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS: str | None = None
 
 environment_variables: dict[str, Callable[[], Any]] = {
     # ================== Runtime Env Vars ==================
@@ -28,6 +29,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # local rank of the process in the distributed setting, used to determine
     # the GPU device id
     "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
+    # Minimum rotary-table span (B*S positions) at which consumers of the shared
+    # fused_qk_norm_rope op take the fused path instead of their eager chain;
+    # "0" = always fuse; unset = each consumer's own measured default. Raw
+    # string or None; validated by fused_qk_norm_rope_min_tokens().
+    "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS": lambda: os.environ.get("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", None),
 }
 
 

vllm_omni/diffusion/layers/fused_qk_norm_rope.py

@@ -1,18 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Q/K RMSNorm followed by packed non-interleaved RoPE.
+"""Q/K RMSNorm followed by packed RoPE (non-interleaved or interleaved).
 
 The public contract is shared by diffusion attention implementations:
 
 * ``q`` and ``k`` are ``[tokens, heads, head_dim]``;
 * norm weights are one-dimensional ``[head_dim]`` tensors;
 * ``rope_table`` is ``[tokens, rotary_dim]`` and stores
-  ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``.
+  ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``;
+  its dtype is either the activation dtype or float32;
+* ``interleaved=False`` rotates half-split pairs ``(d, d + rotary_dim/2)``
+  (MiniMax-H3); ``interleaved=True`` rotates adjacent pairs ``(2i, 2i + 1)``
+  with ``theta_i`` (Boogu-Image, matching its ``apply_rotary_emb``).
 
 The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K
-or rotary-product intermediates. Ascend composes its RMSNorm and rotary fused
-primitives; unsupported inputs use the eager reference.
+or rotary-product intermediates. The interleaved mode supports any even
+``rotary_dim <= head_dim <= 256`` (``tl.arange`` padding to the next power of
+two); the half-split mode keeps the pre-existing kernel and its MiniMax-H3
+geometry contract (``head_dim == 128``, ``rotary_dim == 96``). Ascend
+composes its RMSNorm and rotary fused primitives on the MiniMax-H3 geometry;
+unsupported inputs use the eager reference.
 """
 
 from __future__ import annotations
@@ -26,31 +34,73 @@ 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.diffusion import envs
 from vllm_omni.platforms import current_omni_platform
 
 _FUSED_HEAD_DIM = 128
 _FUSED_ROTARY_DIM = 96
+# Interleaved-mode geometry bounds (any even rotary_dim <= head_dim <= this).
+_FUSED_MAX_HEAD_DIM = 256
 _HEADS_PER_PROGRAM = 8
+# The combined interleaved kernel: small tiles + one warp measured fastest
+# (latency-bound shapes; see the kernel docstring).
+_INTERLEAVED_HEADS_PER_PROGRAM = 4
 
 
 def _apply_rope_table(
     x: torch.Tensor,
     rope_table: torch.Tensor,
     rotary_dim: int,
+    *,
+    interleaved: bool = False,
+    dtype: torch.dtype | None = None,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Eager RoPE application: the fallback and the unit-test reference.
+
+    Shapes: ``x`` is the normalized ``[tokens, heads, head_dim]`` tensor;
+    ``rope_table`` is ``[tokens, rotary_dim]`` storing
+    ``[cos(theta) | sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``.
+
+    ``interleaved`` selects the pairing: half-split ``(d, d + rotary_dim/2)``
+    sharing ``theta_d`` (MiniMax-H3) or adjacent pairs ``(2i, 2i + 1)``
+    sharing ``theta_i`` (Boogu-Image). ``dtype`` is the precision the
+    rotation arithmetic runs in and ``output_dtype`` the precision the result
+    is cast to; both default to ``x``'s dtype (the historical half-split
+    behaviour). The interleaved caller passes ``dtype=torch.float32`` to
+    match the Triton kernel's composition — fp32 rotation of the
+    bf16-rounded normalized value, one final rounding — which is also how
+    vLLM's own ``fused_qk_norm_rope`` Triton kernel and Boogu's eager
+    ``apply_rotary_emb`` order their rounding.
+    """
+    compute_dtype = x.dtype if dtype is None else dtype
+    out_dtype = x.dtype if output_dtype is None else output_dtype
     half = rotary_dim // 2
-    cos = rope_table[..., :half].to(x.dtype).unsqueeze(1)
-    sin = rope_table[..., half:].to(x.dtype).unsqueeze(1)
-    first = x[..., :half]
-    second = x[..., half:rotary_dim]
-    return torch.cat(
-        (
-            first * cos - second * sin,
-            second * cos + first * sin,
-            x[..., rotary_dim:],
-        ),
-        dim=-1,
-    )
+    x_c = x.to(compute_dtype)
+    cos = rope_table[..., :half].to(compute_dtype).unsqueeze(1)
+    sin = rope_table[..., half:].to(compute_dtype).unsqueeze(1)
+    first = x_c[..., 0:rotary_dim:2] if interleaved else x_c[..., :half]
+    second = x_c[..., 1:rotary_dim:2] if interleaved else x_c[..., half:rotary_dim]
+    rotated_first = first * cos - second * sin
+    rotated_second = second * cos + first * sin
+    if interleaved:
+        # stack(..., dim=-1) pairs (new_even_i, new_odd_i) into
+        # [..., half, 2]; flatten(-2) lays the pairs back out as
+        # [e0, o0, e1, o1, ...] — back to their interleaved positions.
+        rotated = torch.stack(
+            (rotated_first, rotated_second),
+            dim=-1,
+        ).flatten(-2)
+        return torch.cat((rotated, x_c[..., rotary_dim:]), dim=-1).to(out_dtype)
+    else:
+        return torch.cat(
+            (
+                rotated_first,
+                rotated_second,
+                x_c[..., rotary_dim:],
+            ),
+            dim=-1,
+        ).to(out_dtype)
 
 
 if HAS_TRITON:
@@ -116,6 +166,135 @@ if HAS_TRITON:
         out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d
         tl.store(out_ptr + out_offsets, output, mask=mask)
 
+    @triton.jit
+    def _qk_norm_rope_kernel(
+        q_ptr,
+        k_ptr,
+        q_weight_ptr,
+        k_weight_ptr,
+        rope_table_ptr,
+        out_q_ptr,
+        out_k_ptr,
+        q_stride_t,
+        q_stride_h,
+        q_stride_d,
+        k_stride_t,
+        k_stride_h,
+        k_stride_d,
+        rope_stride_t,
+        out_q_stride_t,
+        out_q_stride_h,
+        out_q_stride_d,
+        out_k_stride_t,
+        out_k_stride_h,
+        out_k_stride_d,
+        num_q_heads: tl.constexpr,
+        num_kv_heads: tl.constexpr,
+        head_dim: tl.constexpr,
+        padded_dim: tl.constexpr,  # next power of 2 >= head_dim (tl.arange needs a power of 2)
+        rotary_half: tl.constexpr,
+        eps: tl.constexpr,
+        heads_per_program: tl.constexpr,
+        q_head_groups: tl.constexpr,
+        interleaved: tl.constexpr,
+    ):
+        """Fused per-head RMSNorm + RoPE for Q and K in one launch.
+
+        Grid axis 1 assigns the first ``q_head_groups`` programs of a token
+        to Q and the rest to K, so the small K grid hides inside Q's waves
+        instead of paying its own latency-bound launch. ``interleaved``
+        selects the pairing, mirroring ``_apply_rope_table``:
+
+        * ``True`` — adjacent pairs ``(2i, 2i+1)`` sharing ``theta_i``.
+        * ``False`` — half-split pairs ``(d, d + rotary_half)`` sharing
+          ``theta_d``. Production traffic still routes half-split inputs
+          through ``_rms_norm_rope_kernel``; this mode is test-covered and
+          awaits the maintainers' call before taking over that path.
+
+        The modes differ only in the pair/frequency/selector index
+        expressions; the arithmetic is shared. Beyond the rotary width
+        (partial rotary or lane padding) the output is the normalized value
+        unchanged.
+        """
+        token = tl.program_id(0)
+        head_group = tl.program_id(1)
+        dims = tl.arange(0, padded_dim)
+        dim_mask = dims < head_dim
+        if head_group < q_head_groups:
+            heads = head_group * heads_per_program + tl.arange(0, heads_per_program)
+            # Valid-domain mask: rows beyond the head count (when it is not a
+            # multiple of heads_per_program) and padded lanes beyond head_dim.
+            mask = (heads[:, None] < num_q_heads) & dim_mask[None, :]
+            in_ptr = q_ptr
+            weight_ptr = q_weight_ptr
+            out_ptr = out_q_ptr
+            in_stride_t, in_stride_h, in_stride_d = q_stride_t, q_stride_h, q_stride_d
+            out_stride_t, out_stride_h, out_stride_d = (
+                out_q_stride_t,
+                out_q_stride_h,
+                out_q_stride_d,
+            )
+        else:
+            heads = (head_group - q_head_groups) * heads_per_program + tl.arange(0, heads_per_program)
+            mask = (heads[:, None] < num_kv_heads) & dim_mask[None, :]
+            in_ptr = k_ptr
+            weight_ptr = k_weight_ptr
+            out_ptr = out_k_ptr
+            in_stride_t, in_stride_h, in_stride_d = k_stride_t, k_stride_h, k_stride_d
+            out_stride_t, out_stride_h, out_stride_d = (
+                out_k_stride_t,
+                out_k_stride_h,
+                out_k_stride_d,
+            )
+
+        offsets = token * in_stride_t + heads[:, None] * in_stride_h + dims[None, :] * in_stride_d
+        x = tl.load(in_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        weight = tl.load(weight_ptr + dims, mask=dim_mask, other=0.0).to(tl.float32)
+        # Padding lanes load 0 and do not perturb the sum; divide by the true
+        # head_dim, not the padded width.
+        inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps)
+        normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16)
+
+        # The two pairings differ only in three index expressions; the pair
+        # reload, the table gather, the rotation formulas and the write-back
+        # are shared (the same structure as _apply_rope_table).
+        rotary_dim = rotary_half * 2
+        if interleaved:
+            # Adjacent pairs (2i, 2i+1) sharing theta_i.
+            pair_dims = tl.where(dims < rotary_dim, dims ^ 1, dims)
+            freq_dims = tl.where(dims < rotary_dim, dims // 2, 0)
+            is_first = (dims % 2) == 0
+        else:
+            # Half-split pairs (d, d + rotary_half) sharing theta_d.
+            pair_dims = tl.where(
+                dims < rotary_half,
+                dims + rotary_half,
+                tl.where(dims < rotary_dim, dims - rotary_half, dims),
+            )
+            freq_dims = tl.where(
+                dims < rotary_half,
+                dims,
+                tl.where(dims < rotary_dim, dims - rotary_half, 0),
+            )
+            is_first = dims < rotary_half
+        pair_offsets = token * in_stride_t + heads[:, None] * in_stride_h + pair_dims[None, :] * in_stride_d
+        pair_x = tl.load(in_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32)
+        pair_weight = tl.load(weight_ptr + pair_dims, mask=pair_dims < head_dim, other=0.0).to(tl.float32)
+        pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16)
+        table_offsets = token * rope_stride_t + freq_dims
+        cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32)
+        sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32)
+        first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin
+        second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin
+        output = tl.where(
+            dims < rotary_dim,
+            tl.where(is_first, first, second),
+            normalized.to(tl.float32),
+        )
+
+        out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d
+        tl.store(out_ptr + out_offsets, output, mask=mask)
+
 
 def _eager_qk_norm_rope(
     q: torch.Tensor,
@@ -126,12 +305,16 @@ def _eager_qk_norm_rope(
     eps: float,
     head_dim: int,
     rotary_dim: int,
+    interleaved: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     q_norm = F.rms_norm(q, (head_dim,), q_weight, eps)
     k_norm = F.rms_norm(k, (head_dim,), k_weight, eps)
+    # fp32 rotation for the interleaved (Boogu) semantics; the half-split
+    # (H3) reference keeps its historical x-dtype arithmetic via the defaults.
+    rope_dtype = torch.float32 if interleaved else None
     return (
-        _apply_rope_table(q_norm, rope_table, rotary_dim),
-        _apply_rope_table(k_norm, rope_table, rotary_dim),
+        _apply_rope_table(q_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype),
+        _apply_rope_table(k_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype),
     )
 
 
@@ -206,17 +389,23 @@ def _fused_cuda_supported(
     k: torch.Tensor,
     head_dim: int,
     rotary_dim: int,
+    interleaved: bool = False,
 ) -> bool:
-    return (
+    if not (
         HAS_TRITON
         and current_platform.is_cuda()
         and q.is_cuda
         and k.is_cuda
         and q.dtype == torch.bfloat16
         and k.dtype == torch.bfloat16
-        and head_dim == _FUSED_HEAD_DIM
-        and rotary_dim == _FUSED_ROTARY_DIM
-    )
+    ):
+        return False
+    if interleaved:
+        return rotary_dim % 2 == 0 and 2 <= rotary_dim <= head_dim <= _FUSED_MAX_HEAD_DIM
+    # Half-split traffic keeps the pre-existing per-tensor kernel and its
+    # exact MiniMax-H3 geometry contract (that kernel is untouched by the
+    # interleaved extension).
+    return head_dim == _FUSED_HEAD_DIM and rotary_dim == _FUSED_ROTARY_DIM
 
 
 def _fused_npu_supported(
@@ -224,10 +413,12 @@ def _fused_npu_supported(
     k: torch.Tensor,
     head_dim: int,
     rotary_dim: int,
+    interleaved: bool = False,
 ) -> bool:
     """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied."""
     return (
-        current_omni_platform.is_npu()
+        not interleaved
+        and current_omni_platform.is_npu()
         and q.device.type == "npu"
         and k.device.type == "npu"
         and q.dtype == torch.bfloat16
@@ -271,6 +462,69 @@ def _launch_fused_rms_norm_rope(
     return out
 
 
+def _launch_fused_qk_norm_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    rope_table: torch.Tensor,
+    eps: float,
+    interleaved: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    tokens, q_heads, head_dim = q.shape
+    kv_heads = k.shape[1]
+    rotary_half = rope_table.shape[-1] // 2
+    out_q = torch.empty(q.shape, dtype=q.dtype, device=q.device)
+    out_k = torch.empty(k.shape, dtype=k.dtype, device=k.device)
+    if tokens == 0:
+        return out_q, out_k
+    if interleaved:
+        heads_per_program = _INTERLEAVED_HEADS_PER_PROGRAM
+        num_warps, num_stages = 1, 2
+    else:
+        # Match the historical per-tensor launch shape (8 heads per program,
+        # 8 warps) so the half-split arithmetic keeps its reduction tree.
+        heads_per_program = _HEADS_PER_PROGRAM
+        num_warps, num_stages = 8, 3
+    q_head_groups = triton.cdiv(q_heads, heads_per_program)
+    kv_head_groups = triton.cdiv(kv_heads, heads_per_program)
+    grid = (tokens, q_head_groups + kv_head_groups)
+    _qk_norm_rope_kernel[grid](
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+        out_q,
+        out_k,
+        q.stride(0),
+        q.stride(1),
+        q.stride(2),
+        k.stride(0),
+        k.stride(1),
+        k.stride(2),
+        rope_table.stride(0),
+        out_q.stride(0),
+        out_q.stride(1),
+        out_q.stride(2),
+        out_k.stride(0),
+        out_k.stride(1),
+        out_k.stride(2),
+        num_q_heads=q_heads,
+        num_kv_heads=kv_heads,
+        head_dim=head_dim,
+        padded_dim=triton.next_power_of_2(head_dim),
+        rotary_half=rotary_half,
+        eps=eps,
+        heads_per_program=heads_per_program,
+        q_head_groups=q_head_groups,
+        interleaved=interleaved,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return out_q, out_k
+
+
 def _fused_qk_norm_rope_impl(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -280,8 +534,9 @@ def _fused_qk_norm_rope_impl(
     eps: float,
     head_dim: int,
     rotary_dim: int,
+    interleaved: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    if _fused_npu_supported(q, k, head_dim, rotary_dim):
+    if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved):
         return _npu_qk_norm_rope(
             q,
             k,
@@ -291,7 +546,7 @@ def _fused_qk_norm_rope_impl(
             eps,
             rotary_dim,
         )
-    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved):
         return _eager_qk_norm_rope(
             q,
             k,
@@ -301,7 +556,12 @@ def _fused_qk_norm_rope_impl(
             eps,
             head_dim,
             rotary_dim,
+            interleaved,
         )
+    if interleaved:
+        return _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, interleaved=True)
+    # Half-split production traffic keeps the historical per-tensor kernel;
+    # _qk_norm_rope_kernel's half-split mode awaits the maintainers' call.
     return (
         _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps),
         _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps),
@@ -317,8 +577,9 @@ def _fused_qk_norm_rope_fake(
     eps: float,
     head_dim: int,
     rotary_dim: int,
+    interleaved: bool,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim
+    del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim, interleaved
     return torch.empty_like(q), torch.empty_like(k)
 
 
@@ -333,6 +594,39 @@ if not hasattr(torch.ops.vllm_omni, "fused_qk_norm_rope"):
     )
 
 
+_MIN_TOKENS_ENV = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS"
+
+
+def fused_qk_norm_rope_min_tokens(default: int) -> int:
+    """Token span below which a consumer should keep its eager chain.
+
+    Taking the fused path costs a fixed amount of host time per call (the
+    custom-op dispatch into the Python impl and the Triton launcher), while the
+    GPU time it saves grows with the number of tokens the call covers. Below a
+    host- and stack-dependent crossover a call is host-bound and the fused path
+    can cost more end-to-end than it saves; above it the fused path wins.
+    Consumers pass the default they measured on their reference hardware
+    (Boogu-Image: 2048 tokens on one H200) and a deployment on other hardware
+    overrides it once, for every consumer, through
+    ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS`` (``0`` = always fuse).
+
+    Resolved at call time — consumers call this once per forward, not per
+    attention site — so tests and operators can flip it without re-importing.
+    Unset or blank means the caller's default; anything else that is not a
+    non-negative integer raises ``ValueError``.
+    """
+    raw = envs.VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS
+    if raw is None or not raw.strip():
+        return default
+    try:
+        value = int(raw)
+    except ValueError:
+        value = -1
+    if value < 0:
+        raise ValueError(f"{_MIN_TOKENS_ENV} must be a non-negative integer, got {raw!r}")
+    return value
+
+
 def fused_qk_norm_rope(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -343,8 +637,9 @@ def fused_qk_norm_rope(
     *,
     head_dim: int | None = None,
     rotary_dim: int | None = None,
+    interleaved: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Apply Q/K RMSNorm and packed non-interleaved RoPE."""
+    """Apply Q/K RMSNorm and packed RoPE (half-split or adjacent pairs)."""
     if q.ndim != 3 or k.ndim != 3:
         raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}")
     if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]:
@@ -364,15 +659,18 @@ def fused_qk_norm_rope(
         raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}")
     if q_weight.device != q.device or k_weight.device != q.device:
         raise ValueError("Q/K norm weights must be on the activation device")
-    if rope_table.device != q.device or rope_table.dtype != q.dtype:
-        raise ValueError("rope_table must have the same dtype and device as q/k")
+    if rope_table.device != q.device or rope_table.dtype not in (
+        q.dtype,
+        torch.float32,
+    ):
+        raise ValueError("rope_table must be on q/k's device with their dtype or float32")
     if rope_table.shape != (q.shape[0], rotary_dim):
         raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}")
 
     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):
+    if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved):
         return _npu_qk_norm_rope(
             q,
             k,
@@ -382,7 +680,7 @@ def fused_qk_norm_rope(
             eps,
             rotary_dim,
         )
-    if not _fused_cuda_supported(q, k, head_dim, rotary_dim):
+    if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved):
         return _fused_qk_norm_rope_impl(
             q,
             k,
@@ -392,6 +690,7 @@ def fused_qk_norm_rope(
             eps,
             head_dim,
             rotary_dim,
+            interleaved,
         )
     return torch.ops.vllm_omni.fused_qk_norm_rope(
         q,
@@ -402,7 +701,8 @@ def fused_qk_norm_rope(
         eps,
         head_dim,
         rotary_dim,
+        interleaved,
     )
 
 
-__all__ = ["fused_qk_norm_rope"]
+__all__ = ["fused_qk_norm_rope", "fused_qk_norm_rope_min_tokens"]

vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py

@@ -31,10 +31,17 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+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.models.utils import make_attention_mask
 from vllm_omni.platforms import current_omni_platform
 
@@ -43,7 +50,10 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
-RotaryEmbedding = tuple[torch.Tensor, torch.Tensor]
+# (cos, sin) — optionally (cos, sin, packed_table) where packed_table is the
+# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused
+# qk-norm+RoPE op consumes; the eager path ignores the third element.
+RotaryEmbedding = tuple[torch.Tensor, ...]
 RotaryFrequencyTables = list[RotaryEmbedding]
 
 
@@ -56,9 +66,12 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens
 
     Args:
         x: Query or key tensor of shape [B, S, H, D].
-        rotary_emb: Cosine and sine tensors, each shaped [B, S, D].
+        rotary_emb: Cosine and sine tensors, each shaped [B, S, D]. The tuple
+            may carry a third element (the packed table for the fused
+            qk-norm+RoPE path); it is ignored here, hence the indexed access
+            instead of tuple unpacking.
     """
-    freqs_cos, freqs_sin = rotary_emb
+    freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1]
     freqs_cos = freqs_cos.unsqueeze(2)
     freqs_sin = freqs_sin.unsqueeze(2)
 
@@ -73,6 +86,87 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens
     return x_out.type_as(x)
 
 
+# The fused qk-norm+RoPE path is always on where the fast path exists; on
+# other platforms the eager chain below is the behaviour-identical fallback.
+_FUSED_QK_NORM_ROPE = HAS_TRITON and current_platform.is_cuda()
+
+# Boogu's default for the fused-path token gate: fuse only when a rotary table
+# spans at least this many positions (counted as B*S). Below the crossover a
+# request is host-bound and the fused path's per-call Python launch (custom-op
+# impl + Triton launcher) can cost more end-to-end than the kernel saves.
+# Measured on one H200: on the pre-#6871 main, 512^2 (1.1k tokens) paid
+# +191 ms/request and the crossover sat at ~2,350 tokens; on the current main
+# the 512^2 family measures as noise (+/-25 ms, sign-inconsistent) while
+# >=768^2 (2.3k tokens) gains. The crossover is host-dependent — re-benchmark
+# on your hardware and override with VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS
+# (0 = always fuse); see fused_qk_norm_rope_min_tokens().
+_FUSED_MIN_TOKENS = 2048
+
+
+def _with_packed_rope_table(rotary_emb: RotaryEmbedding) -> RotaryEmbedding:
+    """Append the packed fp32 table the fused qk-norm+RoPE op consumes.
+
+    ``freqs_cos``/``freqs_sin`` are theta-repeat-interleaved ``[B, S, D]``
+    (``cos[2i] == cos[2i+1] == cos(theta_i)``), so taking every even column
+    yields the theta-width halves and the packed table is
+    ``[B*S, D] = [cos(theta) | sin(theta)]``. Tables spanning fewer positions
+    (counted as ``B*S``) than the token gate — Boogu's ``_FUSED_MIN_TOKENS``
+    default, overridable through ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS``
+    (``0`` = always fuse) — are returned unchanged: without the third element
+    every consumer stays on the eager chain (this also keeps short tables like
+    the context refiner's off the fused path).
+    """
+    freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1]
+    if freqs_cos.shape[0] * freqs_cos.shape[1] < fused_qk_norm_rope_min_tokens(_FUSED_MIN_TOKENS):
+        return rotary_emb
+    packed = torch.cat((freqs_cos[..., ::2], freqs_sin[..., ::2]), dim=-1)
+    return freqs_cos, freqs_sin, packed.reshape(-1, freqs_cos.shape[-1]).float()
+
+
+def _qk_norm_rope(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    norm_q,
+    norm_k,
+    rotary_emb: RotaryEmbedding | None,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-head Q/K RMSNorm + RoPE + cast, fused into one kernel launch when
+    a packed rope table is present; otherwise the original eager chain."""
+    # Gate on the op's own fast-path predicate: when the fused CUDA kernel
+    # cannot run (dtype, device, or geometry outside even
+    # rotary_dim <= head_dim <= 256), the original chain below must be used
+    # bit-exactly rather than the op's internal eager fallback, whose rounding
+    # differs from vLLM's RMSNorm by one ulp on a fraction of elements.
+    if (
+        rotary_emb is not None
+        and len(rotary_emb) > 2
+        and rotary_emb[2] is not None
+        and _fused_cuda_supported(query, key, query.shape[-1], query.shape[-1], interleaved=True)
+    ):
+        batch, seq_len, num_heads, head_dim = query.shape
+        num_kv_heads = key.shape[2]
+        fused_q, fused_k = fused_qk_norm_rope(
+            query.view(batch * seq_len, num_heads, head_dim),
+            key.view(batch * seq_len, num_kv_heads, head_dim),
+            norm_q.weight,
+            norm_k.weight,
+            rotary_emb[2],
+            norm_q.variance_epsilon,
+            interleaved=True,
+        )
+        return (
+            fused_q.view(batch, seq_len, num_heads, head_dim),
+            fused_k.view(batch, seq_len, num_kv_heads, head_dim),
+        )
+    query = norm_q(query)
+    key = norm_k(key)
+    if rotary_emb is not None:
+        query = apply_rotary_emb(query, rotary_emb)
+        key = apply_rotary_emb(key, rotary_emb)
+    return query.to(dtype), key.to(dtype)
+
+
 def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
     return F.silu(x.float(), inplace=False).to(x.dtype) * y
 
@@ -559,13 +653,7 @@ class BooguImageSelfAttention(nn.Module):
         key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim))
         value = value.unflatten(-1, (self.num_local_kv_heads, self.head_dim))
 
-        query = self.norm_q(query)
-        key = self.norm_k(key)
-
-        if rotary_emb is not None:
-            query = apply_rotary_emb(query, rotary_emb)
-            key = apply_rotary_emb(key, rotary_emb)
-        query, key = query.to(dtype), key.to(dtype)
+        query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype)
 
         attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
         attn_output = self.attn(query, key, value, attn_metadata)
@@ -711,13 +799,7 @@ class BooguImageJointAttention(nn.Module):
         key = key.view(batch_size, -1, self.num_local_kv_heads, self.head_dim)
         value = value.view(batch_size, -1, self.num_local_kv_heads, self.head_dim)
 
-        query = self.norm_q(query)
-        key = self.norm_k(key)
-
-        if rotary_emb is not None:
-            query = apply_rotary_emb(query, rotary_emb)
-            key = apply_rotary_emb(key, rotary_emb)
-        query, key = query.to(dtype), key.to(dtype)
+        query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype)
 
         attn_metadata = AttentionMetadata(attn_mask=joint_attention_mask) if joint_attention_mask is not None else None
         attn_output = self.attn(query, key, value, attn_metadata)
@@ -1449,6 +1531,11 @@ class BooguImageTransformer2DModel(nn.Module):
                     shift += ref_img_len
                     idx += 1
 
+            if _FUSED_QK_NORM_ROPE:
+                # The rebuilt per-reference-image tuple needs its own packed
+                # table; the forward-level tuples do not flow into this batch.
+                batch_ref_img_rotary_emb = _with_packed_rope_table(batch_ref_img_rotary_emb)
+
             for layer in self.ref_image_refiner:
                 batch_ref_image_hidden_states = layer(
                     batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb
@@ -1534,6 +1621,17 @@ class BooguImageTransformer2DModel(nn.Module):
             device,
         )
 
+        if _FUSED_QK_NORM_ROPE:
+            # One packed [cos|sin] table per rotary embedding that reaches an
+            # attention, built once per forward; the tuples grow a third
+            # element that the fused path consumes and the eager path ignores.
+            # (ref_img_rotary_emb is packed inside img_patch_embed_and_refine,
+            # on the rebuilt per-reference-image batch tuple.)
+            context_rotary_emb = _with_packed_rope_table(context_rotary_emb)
+            noise_rotary_emb = _with_packed_rope_table(noise_rotary_emb)
+            rotary_emb = _with_packed_rope_table(rotary_emb)
+            combined_img_rotary_emb = _with_packed_rope_table(combined_img_rotary_emb)
+
         # Context refinement.
         for layer in self.context_refiner:
             instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb)