From a6570b68140cded3aa00e1fa6cddd78a869db6fe Mon Sep 17 00:00:00 2001 From: dongbo910220 <32610838+dongbo910220@users.noreply.github.com> Date: Fri, 18 Sep 2026 03:06:45 +0800 Subject: [PATCH] [Diffusion] Fuse ERNIE-Image Q/K RoPE for eager execution (#7502) Signed-off-by: dongbo910220 <1275604947@qq.com> --- .../test_ernie_image_fused_rope.py | 169 ++++++++++++++ .../ernie_image/test_ernie_image_rope.py | 31 +++ .../ernie_image/ernie_image_transformer.py | 26 ++- .../models/ernie_image/fused_rope.py | 211 ++++++++++++++++++ 4 files changed, 435 insertions(+), 2 deletions(-) create mode 100644 tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py create mode 100644 vllm_omni/diffusion/models/ernie_image/fused_rope.py diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py b/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py new file mode 100644 index 000000000..bf3dce4f3 --- /dev/null +++ b/tests/diffusion/models/ernie_image/test_ernie_image_fused_rope.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from vllm.triton_utils import HAS_TRITON + +from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import _apply_rotary_emb + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def reset_fused_rope_state(): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + fused_rope._FAILED_KEYS.clear() + yield + fused_rope._FAILED_KEYS.clear() + + +def _inputs(shape: tuple[int, int, int, int]): + batch, sequence, _, head_dim = shape + torch.manual_seed(17) + query = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + key = torch.randn_like(query) + freqs_cos = torch.randn( + batch, + sequence, + head_dim // 2, + device="cuda", + dtype=torch.float32, + ) + freqs_sin = torch.randn_like(freqs_cos) + return query, key, freqs_cos, freqs_sin + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize( + "shape", + [ + (1, 1, 1, 128), + (2, 257, 7, 128), + (2, 4608, 32, 128), + ], +) +def test_fused_qk_rope_is_bit_exact(shape): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + query, key, freqs_cos, freqs_sin = _inputs(shape) + with torch.inference_mode(): + expected_query = _apply_rotary_emb(query, freqs_cos, freqs_sin) + expected_key = _apply_rotary_emb(key, freqs_cos, freqs_sin) + actual = fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) + + assert actual is not None + actual_query, actual_key = actual + assert torch.equal(actual_query, expected_query) + assert torch.equal(actual_key, expected_key) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_launch_failure_disables_runtime_key(monkeypatch): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128)) + launch_count = 0 + + def failing_launch(*args): + nonlocal launch_count + launch_count += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr( + fused_rope, + "_launch_fused_qk_rotary_emb", + failing_launch, + ) + with torch.inference_mode(): + for _ in range(2): + assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None + assert launch_count == 1 + assert len(fused_rope._FAILED_KEYS) == 1 + + +def test_failed_runtime_key_cache_is_bounded(): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + for index in range(fused_rope._FAILED_KEYS_MAX_SIZE + 1): + fused_rope._record_failed_runtime_key((index,)) + + assert len(fused_rope._FAILED_KEYS) == fused_rope._FAILED_KEYS_MAX_SIZE + assert (0,) not in fused_rope._FAILED_KEYS + assert (fused_rope._FAILED_KEYS_MAX_SIZE,) in fused_rope._FAILED_KEYS + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_compile_path_does_not_launch_fused_kernel(monkeypatch): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128)) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + def unexpected_launch(*args): + raise AssertionError("compile path must remain native") + + monkeypatch.setattr( + fused_rope, + "_launch_fused_qk_rotary_emb", + unexpected_launch, + ) + with torch.inference_mode(): + assert fused_rope.try_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize("requires_grad_input", ["query", "key", "cos", "sin"]) +def test_gradient_inputs_do_not_launch_fused_kernel(monkeypatch, requires_grad_input): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + values = list(_inputs((1, 17, 3, 128))) + input_index = {"query": 0, "key": 1, "cos": 2, "sin": 3}[requires_grad_input] + values[input_index].requires_grad_(True) + + def unexpected_launch(*args): + raise AssertionError("autograd inputs must remain on the native path") + + monkeypatch.setattr( + fused_rope, + "_launch_fused_qk_rotary_emb", + unexpected_launch, + ) + assert ( + fused_rope.try_fused_qk_rotary_emb( + *values, + ) + is None + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_mismatched_key_dtype_or_device_does_not_launch(monkeypatch): + from vllm_omni.diffusion.models.ernie_image import fused_rope + + query, key, freqs_cos, freqs_sin = _inputs((1, 17, 3, 128)) + + def unexpected_launch(*args): + raise AssertionError("mismatched Q/K inputs must remain native") + + monkeypatch.setattr( + fused_rope, + "_launch_fused_qk_rotary_emb", + unexpected_launch, + ) + for unsupported_key in (key.float(), key.cpu()): + assert ( + fused_rope.try_fused_qk_rotary_emb( + query, + unsupported_key, + freqs_cos, + freqs_sin, + ) + is None + ) diff --git a/tests/diffusion/models/ernie_image/test_ernie_image_rope.py b/tests/diffusion/models/ernie_image/test_ernie_image_rope.py index 22d8efaad..34beb69f9 100644 --- a/tests/diffusion/models/ernie_image/test_ernie_image_rope.py +++ b/tests/diffusion/models/ernie_image/test_ernie_image_rope.py @@ -1,8 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import torch from vllm_omni.diffusion.models.ernie_image.ernie_image_transformer import ( ErnieImageEmbedND3, + _apply_qk_rotary_emb, _apply_rotary_emb, rope, ) @@ -185,3 +189,30 @@ class TestErnieImageRopePositionEmbedding: out_diffusers = x * cos_ + rotated * sin_ assert torch.allclose(out_vllm, out_diffusers, atol=1e-6), "vllm-omni and diffusers outputs should match" + + def test_apply_qk_rotary_emb_uses_native_fallback(self, monkeypatch): + """Unsupported routes preserve both existing eager calls exactly.""" + from vllm_omni.diffusion.models.ernie_image import ( + ernie_image_transformer, + ) + + B, S, H, D = 2, 17, 3, 128 + query = torch.randn(B, S, H, D) + key = torch.randn(B, S, H, D) + freqs_cos = torch.randn(B, S, D // 2) + freqs_sin = torch.randn(B, S, D // 2) + monkeypatch.setattr( + ernie_image_transformer, + "try_fused_qk_rotary_emb", + lambda *args: None, + ) + + actual_query, actual_key = _apply_qk_rotary_emb(query, key, freqs_cos, freqs_sin) + assert torch.equal( + actual_query, + _apply_rotary_emb(query, freqs_cos, freqs_sin), + ) + assert torch.equal( + actual_key, + _apply_rotary_emb(key, freqs_cos, freqs_sin), + ) diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index bd96630e8..eb26e5389 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py @@ -28,6 +28,9 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) +from vllm_omni.diffusion.models.ernie_image.fused_rope import ( + try_fused_qk_rotary_emb, +) logger = init_logger(__name__) @@ -101,6 +104,26 @@ def _apply_rotary_emb( return torch.cat((x_rot * cos_ + x_rotated * sin_, x[..., rot_dim:]), dim=-1) +def _apply_qk_rotary_emb( + query: torch.Tensor, + key: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + fused = try_fused_qk_rotary_emb( + query, + key, + freqs_cos, + freqs_sin, + ) + if fused is not None: + return fused + return ( + _apply_rotary_emb(query, freqs_cos, freqs_sin), + _apply_rotary_emb(key, freqs_cos, freqs_sin), + ) + + class ErnieImageEmbedND3(nn.Module): def __init__(self, dim: int, theta: int, axes_dim: tuple[int, int, int]): super().__init__() @@ -340,8 +363,7 @@ class ErnieImageAttention(nn.Module): if image_rotary_emb is not None: freqs_cos, freqs_sin = image_rotary_emb - query = _apply_rotary_emb(query, freqs_cos, freqs_sin) - key = _apply_rotary_emb(key, freqs_cos, freqs_sin) + query, key = _apply_qk_rotary_emb(query, key, freqs_cos, freqs_sin) attn_metadata = None if attention_mask is not None: diff --git a/vllm_omni/diffusion/models/ernie_image/fused_rope.py b/vllm_omni/diffusion/models/ernie_image/fused_rope.py new file mode 100644 index 000000000..5bdb61165 --- /dev/null +++ b/vllm_omni/diffusion/models/ernie_image/fused_rope.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Lossless eager Q/K RoPE fusion for ERNIE-Image.""" + +from __future__ import annotations + +import logging +from collections import OrderedDict + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + +_FAILED_KEYS_MAX_SIZE = 128 +_FAILED_KEYS: OrderedDict[tuple[object, ...], None] = OrderedDict() + + +if HAS_TRITON: + + @triton.jit + def _round_bf16_to_fp32(value): + """RNE-round FP32 to BF16 precision in an FP32 register.""" + + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + @triton.jit + def _qk_rope_kernel( + out_q_ptr, + out_k_ptr, + q_ptr, + k_ptr, + cos_ptr, + sin_ptr, + heads, + head_dim: tl.constexpr, + rotary_dim: tl.constexpr, + rotary_half: tl.constexpr, + heads_block: tl.constexpr, + half_block: tl.constexpr, + tail_block: tl.constexpr, + ): + row = tl.program_id(0).to(tl.int64) + is_key = tl.program_id(1) + base = row * heads * head_dim + head_offsets = tl.arange(0, heads_block)[:, None] + head_mask = head_offsets < heads + columns = tl.arange(0, half_block)[None, :] + column_mask = columns < rotary_half + mask = head_mask & column_mask + + if is_key == 0: + input_ptr = q_ptr + output_ptr = out_q_ptr + else: + input_ptr = k_ptr + output_ptr = out_k_ptr + + first_offsets = base + head_offsets * head_dim + columns + second_offsets = first_offsets + rotary_half + first = tl.load(input_ptr + first_offsets, mask=mask, other=0.0).to(tl.float32) + second = tl.load(input_ptr + second_offsets, mask=mask, other=0.0).to(tl.float32) + + # The eager path casts FP32 frequencies to BF16, duplicates adjacent + # values, rounds both BF16 multiplies, and finally rounds the add. + first_freq = columns // 2 + second_freq = (rotary_half + columns) // 2 + cos_first = _round_bf16_to_fp32(tl.load(cos_ptr + row * rotary_half + first_freq, mask=column_mask)) + cos_second = _round_bf16_to_fp32(tl.load(cos_ptr + row * rotary_half + second_freq, mask=column_mask)) + sin_first = _round_bf16_to_fp32(tl.load(sin_ptr + row * rotary_half + first_freq, mask=column_mask)) + sin_second = _round_bf16_to_fp32(tl.load(sin_ptr + row * rotary_half + second_freq, mask=column_mask)) + + out_first = _round_bf16_to_fp32(first * cos_first) + _round_bf16_to_fp32(-second * sin_first) + out_second = _round_bf16_to_fp32(second * cos_second) + _round_bf16_to_fp32(first * sin_second) + tl.store(output_ptr + first_offsets, out_first, mask=mask) + tl.store(output_ptr + second_offsets, out_second, mask=mask) + + if head_dim > rotary_dim: + tail_columns = rotary_dim + tl.arange(0, tail_block)[None, :] + tail_mask = head_mask & (tail_columns < head_dim) + tail_offsets = base + head_offsets * head_dim + tail_columns + tail = tl.load(input_ptr + tail_offsets, mask=tail_mask, other=0.0) + tl.store(output_ptr + tail_offsets, tail, mask=tail_mask) + + +def _runtime_key( + query: torch.Tensor, + key: torch.Tensor, + freqs_cos: torch.Tensor, +) -> tuple[object, ...]: + return ( + query.device.index, + tuple(query.shape), + tuple(key.shape), + tuple(freqs_cos.shape), + query.dtype, + freqs_cos.dtype, + ) + + +def _is_failed_runtime_key(runtime_key: tuple[object, ...]) -> bool: + if runtime_key not in _FAILED_KEYS: + return False + _FAILED_KEYS.move_to_end(runtime_key) + return True + + +def _record_failed_runtime_key(runtime_key: tuple[object, ...]) -> None: + _FAILED_KEYS[runtime_key] = None + _FAILED_KEYS.move_to_end(runtime_key) + while len(_FAILED_KEYS) > _FAILED_KEYS_MAX_SIZE: + _FAILED_KEYS.popitem(last=False) + + +def _supported_inputs( + query: torch.Tensor, + key: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> bool: + if torch.compiler.is_compiling(): + return False + if any(tensor.requires_grad for tensor in (query, key, freqs_cos, freqs_sin)): + return False + if not HAS_TRITON or not current_omni_platform.is_cuda(): + return False + if query.dtype is not torch.bfloat16 or not query.is_cuda: + return False + if query.ndim != 4 or key.shape != query.shape or key.dtype is not query.dtype or key.device != query.device: + return False + if not query.is_contiguous() or not key.is_contiguous(): + return False + if freqs_cos.dtype is not torch.float32 or freqs_sin.dtype is not torch.float32: + return False + if not freqs_cos.is_cuda or freqs_cos.device != query.device: + return False + if freqs_sin.device != query.device or freqs_sin.shape != freqs_cos.shape: + return False + if not freqs_cos.is_contiguous() or not freqs_sin.is_contiguous(): + return False + batch, sequence, _, head_dim = query.shape + rotary_half = freqs_cos.shape[-1] if freqs_cos.ndim == 3 else 0 + return ( + query.numel() > 0 + and freqs_cos.shape == (batch, sequence, rotary_half) + and rotary_half > 0 + and rotary_half * 2 <= head_dim + ) + + +def _launch_fused_qk_rotary_emb( + query: torch.Tensor, + key: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + batch, sequence, heads, head_dim = query.shape + rotary_half = freqs_cos.shape[-1] + rotary_dim = rotary_half * 2 + tail = head_dim - rotary_dim + output_query = torch.empty_like(query) + output_key = torch.empty_like(key) + with torch.accelerator.device_index(query.device.index): + _qk_rope_kernel[(batch * sequence, 2)]( + output_query, + output_key, + query, + key, + freqs_cos, + freqs_sin, + heads, + head_dim=head_dim, + rotary_dim=rotary_dim, + rotary_half=rotary_half, + heads_block=triton.next_power_of_2(heads), + half_block=triton.next_power_of_2(rotary_half), + tail_block=triton.next_power_of_2(max(tail, 1)), + ) + return output_query, output_key + + +def try_fused_qk_rotary_emb( + query: torch.Tensor, + key: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return fused Q/K outputs, or ``None`` to request the eager fallback.""" + + if not _supported_inputs(query, key, freqs_cos, freqs_sin): + return None + + runtime_key = _runtime_key(query, key, freqs_cos) + if _is_failed_runtime_key(runtime_key): + return None + + try: + return _launch_fused_qk_rotary_emb(query, key, freqs_cos, freqs_sin) + except Exception: + _record_failed_runtime_key(runtime_key) + logger.exception("ERNIE-Image fused Q/K RoPE failed; using the eager fallback") + return None + + +__all__ = ["try_fused_qk_rotary_emb"]