[Diffusion] Fuse ERNIE-Image Q/K RoPE for eager execution (#7502)

Signed-off-by: dongbo910220 <1275604947@qq.com>

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

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

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:

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"]