vllm-omni/vllm_omni/diffusion/attention/parallel/ring.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING

import torch
from vllm.logger import init_logger

# import torch.distributed as dist # Not used directly here, but good practice if needed
from vllm_omni.diffusion.attention.backends.ring.ring_globals import (
    FA3_SUPPORTED_CUDA_MAJORS,
    HAS_AITER,
    HAS_FA3,
    HAS_FA4,
    HAS_FLASH_ATTN,
)
from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
from vllm_omni.diffusion.attention.parallel.base import (
    ParallelAttentionContext,
    # ParallelAttentionStrategy, # Not used in type hint below currently
)
from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator

if TYPE_CHECKING:
    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata

logger = init_logger(__name__)


def _can_use_fa3(device: torch.device) -> bool:
    """Return whether the installed FA3 kernels support ``device``.

    Importing an extension only proves that its Python module is installed.
    When the source publishes a supported-major contract, also check it before
    entering a CUDA launcher that may abort instead of raising an exception.
    """
    if not HAS_FA3 or device.type != "cuda":
        return False
    if FA3_SUPPORTED_CUDA_MAJORS is None:
        return True
    major, _minor = torch.cuda.get_device_capability(device)
    return major in FA3_SUPPORTED_CUDA_MAJORS


def _can_use_fa4(device: torch.device) -> bool:
    """Return whether the installed FA4 kernels support ``device``."""
    if not HAS_FA4 or device.type != "cuda":
        return False
    major, _minor = torch.cuda.get_device_capability(device)
    return major >= 10


def _can_use_fa2(device: torch.device) -> bool:
    """Return whether the FA2 backend supports ``device``.

    FA2's CUDA backend supports Ampere, Ada, and Hopper.  Blackwell support is
    provided by the separate FA4 implementation.
    """
    if not HAS_FLASH_ATTN or device.type != "cuda":
        return False
    major, _minor = torch.cuda.get_device_capability(device)
    return major in (8, 9)


@dataclass(frozen=True, slots=True)
class _RingCtx(ParallelAttentionContext):
    """Per-forward context for Ring sequence-parallel attention."""

    # Ring attention typically doesn't need complex context for post-processing
    # as the output is already correctly sharded along sequence dimension.
    pass


class RingParallelAttention:
    """Ring sequence-parallel strategy.

    This strategy prepares inputs for Ring Attention.
    Key responsibilities:
    - Concatenate joint_query (Text) to query (Image) if present.
    - Keep joint_key/value separate in metadata for the Ring kernel to handle as static prefix.
    """

    def __init__(
        self,
        sp_group: SequenceParallelGroupCoordinator,
        attn_backend_pref: str | None = None,
        attn_backend_explicit: bool = False,
    ) -> None:
        self._sp_group = sp_group
        self.attn_backend_pref = attn_backend_pref
        self.attn_backend_explicit = attn_backend_explicit

    @property
    def enabled(self) -> bool:
        return True

    @property
    def name(self) -> str:
        return "ring"

    def pre_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None,
    ):
        joint_tensor_query = None
        joint_strategy = "front"

        if attn_metadata is not None:
            joint_tensor_query = attn_metadata.joint_query
            joint_strategy = attn_metadata.joint_strategy

        if joint_tensor_query is not None:
            supported_joint_strategy = ["front", "rear"]
            if joint_strategy not in supported_joint_strategy:
                raise ValueError(f"joint_strategy: {joint_strategy} not supported.")

            if joint_strategy == "front":
                query = torch.cat([joint_tensor_query, query], dim=1)
            else:
                query = torch.cat([query, joint_tensor_query], dim=1)

            # Note: We do NOT concatenate joint_key/value here.
            # They are preserved in attn_metadata and will be passed
            # explicitly to ring_flash_attn_func.

        ctx = _RingCtx(name=self.name)
        return query, key, value, attn_metadata, ctx

    def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor:
        # Ring attention output is already sharded correctly along sequence dimension.
        return attn_output

    def run_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None,
        softmax_scale: float | None = None,
        causal: bool = False,
    ) -> torch.Tensor:
        """Run the actual Ring Attention kernel."""
        if softmax_scale is None:
            softmax_scale = query.shape[-1] ** -0.5

        backend_pref = self.attn_backend_pref
        if backend_pref is not None:
            backend_pref = backend_pref.lower()

        # Extract joint tensors
        joint_key, joint_value = None, None
        joint_strategy = "front"
        valid_kv_length = None
        if attn_metadata is not None:
            joint_key = attn_metadata.joint_key
            joint_value = attn_metadata.joint_value
            valid_kv_length = attn_metadata.extra.get("valid_kv_length")
            if attn_metadata.joint_strategy is not None:
                joint_strategy = attn_metadata.joint_strategy

        def _run_sdpa() -> torch.Tensor:
            from vllm_omni.diffusion.attention.backends.ring_pytorch_attn import ring_pytorch_attn_func

            return ring_pytorch_attn_func(
                query,
                key,
                value,
                softmax_scale=softmax_scale,
                causal=causal,
                group=self._sp_group.ring_group,
                op_type="efficient",
                joint_tensor_key=joint_key,
                joint_tensor_value=joint_value,
                joint_strategy=joint_strategy,
                valid_kv_length=valid_kv_length,
            )

        # Ring only implements local Flash/AITER and SDPA kernels. HuggingFace
        # Hub FA and other non-local backends can use the SDPA ring path only
        # when they came from automatic platform selection. An explicit
        # request must never silently substitute local FA4/FA3/FA2.
        _sdpa_prefs = {
            "sdpa",
            "torch",
            "torch_sdpa",
        }
        _non_ring_prefs = {
            "cudnn_attn",
            "flashinfer_attn",
            "trtllm_attn",
            "sage_attn",
            "sage_attn_3",
            "flash_attn_hub",
            "flash_attn_3_hub",
        }
        if backend_pref in _sdpa_prefs:
            return _run_sdpa()
        if backend_pref in _non_ring_prefs:
            if self.attn_backend_explicit:
                raise ValueError(
                    f"{self.attn_backend_pref} was explicitly selected, but ring sequence parallelism "
                    "has no implementation for that backend. Select TORCH_SDPA/FLASH_ATTN or use Ulysses SP."
                )
            return _run_sdpa()

        if query.dtype == torch.float32:
            if self.attn_backend_explicit:
                raise ValueError(
                    f"{self.attn_backend_pref} was explicitly selected for ring attention, "
                    "but its ring kernel does not support float32. Select TORCH_SDPA or use a supported dtype."
                )
            return _run_sdpa()

        can_use_fa4 = _can_use_fa4(query.device)
        can_use_fa3 = _can_use_fa3(query.device)
        can_use_fa2 = _can_use_fa2(query.device)
        if not can_use_fa4 and not can_use_fa3 and not can_use_fa2 and not HAS_AITER:
            if self.attn_backend_explicit:
                raise RuntimeError(
                    f"{self.attn_backend_pref} was explicitly selected, but no compatible ring kernel "
                    f"is available for device {query.device}."
                )
            logger.warning_once(
                "Automatic ring backend selection chose TORCH_SDPA because no compatible "
                "FA2/FA3/FA4/AITER ring kernel is available for this device."
            )
            return _run_sdpa()

        from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func

        # Prefer FA4 on Blackwell. An importable Hopper-only FA3 wheel can
        # otherwise be selected and fail at launch with "no kernel image".
        if can_use_fa4:
            attn_type = AttnType.FA4
        # Prefer FA3 over FA2 on Ampere/Ada/Hopper. On ROCm, use AITER.
        elif can_use_fa3:
            attn_type = AttnType.FA3
        elif HAS_AITER:
            attn_type = AttnType.AITER
        elif can_use_fa2:
            attn_type = AttnType.FA
        else:
            raise RuntimeError("No compatible Flash Attention backend is available for ring attention.")

        return ring_flash_attn_func(
            query,
            key,
            value,
            dropout_p=0.0,
            softmax_scale=softmax_scale,
            causal=causal,
            window_size=(-1, -1),
            softcap=0.0,
            alibi_slopes=None,
            deterministic=False,
            group=self._sp_group.ring_group,
            attn_type=attn_type,
            joint_tensor_key=joint_key,
            joint_tensor_value=joint_value,
            joint_strategy=joint_strategy,
            valid_kv_length=valid_kv_length,
        )