From f7d9deb45ab56e6a2ccc1690279bd9e6bdefbfe3 Mon Sep 17 00:00:00 2001 From: KANG Qihan <3149604185@qq.com> Date: Sun, 13 Sep 2026 09:39:27 +0800 Subject: [PATCH] [Kernel][Boogu-Image] Fuse Q/K RMSNorm + interleaved RoPE via fused_qk_norm_rope (#6982) Signed-off-by: Qihan Kang --- docs/configuration/environment_variables.md | 1 + tests/config/test_environment_variables.py | 2 +- .../layers/test_fused_qk_norm_rope.py | 95 ++++- .../test_boogu_image_transformer.py | 153 +++++++- .../config/environment_variable_inventory.py | 1 + vllm_omni/diffusion/envs.py | 6 + .../diffusion/layers/fused_qk_norm_rope.py | 366 ++++++++++++++++-- .../boogu_image/boogu_image_transformer.py | 132 ++++++- 8 files changed, 703 insertions(+), 53 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 90b9ffae0..4ddc2939c 100644 --- a/docs/configuration/environment_variables.md +++ b/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 | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index cc37cd6b1..911cfbc80 100644 --- a/tests/config/test_environment_variables.py +++ b/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, diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py index 40888556a..84edce0dc 100644 --- a/tests/diffusion/layers/test_fused_qk_norm_rope.py +++ b/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) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 43e2f12f4..79954f1db 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/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 diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 4a4966d9c..fe439457d 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/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", diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index 32e006f7e..4cbc0907c 100644 --- a/vllm_omni/diffusion/envs.py +++ b/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), } diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py index f7e27675c..9f6605726 100644 --- a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py +++ b/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"] diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 66031ca3a..93d52b960 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/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)