vllm-omni/vllm_omni/diffusion/attention/backends/abstract.py

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

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Generic, Literal, TypeVar

import torch

from vllm_omni.platforms import current_omni_platform


class AttentionBackend(ABC):
    """Abstract class for diffusion attention backends."""

    accept_output_buffer: bool = False
    supports_piecewise_spans: bool = False
    # A backend that supports this capability can consume the opaque paged-KV
    # context prepared by the diffusion Worker data plane.  Keeping the
    # capability on the backend class prevents a paged request from silently
    # falling back to dense attention on an incompatible implementation.
    supports_paged_kv: bool = False
    # The backend can represent a contiguous valid K/V prefix by slicing the
    # tensors instead of materializing a padding mask. Models may use this to
    # avoid a slower masked-attention plan when tail padding is not semantic.
    supports_prefix_kv_slicing: bool = False

    @classmethod
    def supports_packed_mask_free(cls) -> bool:
        """Whether [real, pad] packed layouts can run without attn_mask.

        When True, models that pack a [real, pad] two-document layout and
        provide ``AttentionMetadata.packed_padding`` alongside the packed
        cu_seqlens/max_seqlen metadata may skip constructing the padding mask
        entirely. Backends whose mask-free behavior is platform-dependent must
        check current_omni_platform.
        """
        return False

    @classmethod
    def supports_multi_doc_packed_varlen(cls) -> bool:
        """Whether this backend keeps N-document packed boundaries isolated.

        When True, the backend consumes ``AttentionMetadata.extra`` cu_seqlens
        as a genuine block-diagonal attention plan (a dedicated varlen kernel,
        not a padding-mask rebuild), so a caller may pack multiple real
        requests into one forward without attention crossing document
        boundaries. When False, callers packing more than one real document
        must run one forward per document; otherwise a backend that only
        supports a ``[real, pad]`` two-document contract, or that ignores
        cu_seqlens outright, will silently attend across request boundaries.
        Backends whose kernel selection is platform-dependent must consult
        ``current_omni_platform``.
        """
        return False

    # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
    # Platform resolution rejects an explicit selection outside this set, so a
    # hardware-specific backend fails before the model is built.
    supported_platforms: tuple[str, ...] | None = None

    @classmethod
    def validate_available(cls) -> None:
        """Raise if this backend's optional dependencies are missing.

        Called during platform resolution, i.e. before model construction, so a
        backend that probes its kernel package lazily still reports the problem
        while the user can still act on it.
        """
        return None

    @classmethod
    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
        """Return whether this backend can consume a nontrivial ``attn_mask``.

        ``attention_spec`` is the resolved per-role config when the user picked
        a backend explicitly. Implementations that depend on kernel variant
        (for example FlashInfer cute-dsl vs fa2) must consult it so capability
        probes match a runnable configuration.
        """
        del attention_spec
        return False

    @staticmethod
    @abstractmethod
    def get_name() -> str:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_impl_cls() -> type["AttentionImpl"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
    def get_supported_head_sizes() -> list[int]:
        """Get the list of supported head sizes for this backend."""
        raise NotImplementedError

    @classmethod
    def supports_head_size(cls, head_size: int) -> bool:
        supported_head_sizes = cls.get_supported_head_sizes()
        return (not supported_head_sizes) or head_size in supported_head_sizes

    @classmethod
    def indexes_kv_by_block_stride(cls) -> bool:
        """Whether this backend reads K/V pages by the runtime block stride.

        Returning ``True`` means the physical cache layout has ``num_blocks``
        as its outer stride, so native vLLM may safely use page-size padding
        when it unifies cache layouts across layers. Dense diffusion backends
        conservatively keep the default ``False``; a paged backend should
        override this only when its kernel actually follows that layout.

        Since vLLM 0.29 this no longer rides on ``AttentionSpec``.  It selects
        the block-outermost physical layout pinned on
        ``CacheConfig.kv_cache_layout`` -- see
        ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose
        ``is_block_outermost`` is what upstream reads when it sizes
        ``KVCacheTensor`` regions.
        """

        return False


@dataclass(frozen=True, slots=True)
class QueryRange:
    local_start: int
    local_end: int
    global_start: int


@dataclass(frozen=True, slots=True)
class VideoTokenSpan:
    """One physical video-grid slice in packed document 0."""

    start: int
    latent_grid: tuple[int, int, int]
    role: Literal["reference", "target"]

    @property
    def length(self) -> int:
        return self.latent_grid[0] * self.latent_grid[1] * self.latent_grid[2]


@dataclass(frozen=True, slots=True)
class VideoTokenLayout:
    """Video-grid slices in a packed multimodal sequence.

    A model that packs its sequence as ``[prefix | t*h*w video rows | padding]``
    publishes this so backends can recover spatiotemporal locality; the prefix
    holds everything that is not video (text, visual conditions, audio).
    Publishing it also asserts that any ``attn_mask`` masks only the trailing
    padding, so ``prefix_len + t*h*w`` is the used length of the sequence.

    ``prefix_len``/``latent_grid`` retain the original one-tail contract. A
    Ref2VA layout instead publishes ``used_len`` and every physical video
    span, so audio and image rows between videos remain dense.

    Plain ints, so reading it never forces a device-to-host sync.
    """

    prefix_len: int | None = None
    latent_grid: tuple[int, int, int] | None = None
    used_len: int | None = None
    video_spans: tuple[VideoTokenSpan, ...] = ()


@dataclass(frozen=True, slots=True)
class PackedPaddingMetadata:
    """Producer-validated mask-free view of padding in a [real, pad] packing.

    The cumulative-length tensors are canonical two-element ``[0, length]``
    views. Consumers may use them without reading device scalars because the
    producer owns the packing.
    """

    q_length: int
    kv_length: int
    cu_seqlens_q: torch.Tensor
    cu_seqlens_k: torch.Tensor


@dataclass
class AttentionMetadata:
    attn_mask: torch.Tensor | None = None
    joint_attn_mask: torch.Tensor | None = None
    # a joint mask for the joint query, key, and value, depends the joint_strategy
    joint_query: torch.Tensor | None = None
    # a replicated tensor among processes appended to the front or rear of query, depends the joint_strategy
    joint_key: torch.Tensor | None = None
    # a replicated tensor among processes appended to the front or rear of key, depends the joint_strategy
    joint_value: torch.Tensor | None = None
    # a replicated tensor among processes appended to the front or rear of value, depends the joint_strategy
    joint_strategy: str = "front"
    # the strategy to joint the query, key, and value, can be "front" or "rear"
    extra: dict[str, Any] = field(default_factory=dict)
    # Opaque backend-specific per-forward parameters (e.g. block masks, KV indices).
    # Backends MUST silently ignore unknown keys.
    #
    # Well-known optional keys (convention, not required on all forwards):
    #   "kv_cache_dtype": str | None — quantized KV dtype (e.g. "fp8"); backends
    #     decide whether/how to apply.
    #   "cu_seqlens_q" / "cu_seqlens_k": int32 CUDA tensors describing packed
    #     variable-length query/key sequences for FlashAttention.
    #   "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with
    #     the packed cu_seqlens tensors.
    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
    #     backend that advertises supports_prefix_kv_slicing.
    #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
    #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
    #     Requires the [real, pad] two-document packing contract; see
    #     FlashAttentionImpl._forward_varlen_packed_npu.
    #   "laser_input_scale": float — model opt-in input pre-scale for the NPU
    #     ascend_laser_attention path. The kernel stores unscaled QK^T in an
    #     fp16 workspace, so outlier activations overflow 65504 into NaN rows;
    #     with this set (>1), q/k/v are divided by the factor before the op,
    #     the kernel scale_value is multiplied by its square, and the output
    #     is scaled back (exact for power-of-two factors). Absent means no
    #     pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu.

    # Piecewise attention metadata (mixed causal/full masks).
    # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
    full_attn_spans: list[list[tuple[int, int]]] | None = None
    query_ranges: tuple[QueryRange, ...] | None = None

    # Geometry of the video segment for backends that exploit spatiotemporal
    # locality (block-sparse selection, tiled masks). Dense backends ignore it.
    video_layout: VideoTokenLayout | None = None

    # Canonical mask-free view of structural suffix padding. Backends that do
    # not advertise supports_packed_mask_free ignore it.
    packed_padding: PackedPaddingMetadata | None = None


T = TypeVar("T", bound=AttentionMetadata)


class AttentionImpl(ABC, Generic[T]):
    # Per-platform kv_cache_dtype support. Maps OmniPlatformEnum value
    # (e.g. "cuda", "npu") to the set of quantized dtypes that platform
    # handles.
    #
    # To add FP8 support for a new platform in a subclass:
    #   _supported_kv_cache_dtypes = {"cuda": {"fp8"}, "npu": {"fp8"}}
    _supported_kv_cache_dtypes: dict[str, set[str]] = {}

    @abstractmethod
    def __init__(
        self,
        num_heads: int,
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
        num_kv_heads: int | None = None,
        prefix: str = "",
        qkv_layout: str | None = None,
        backend_kwargs: dict[str, Any] | None = None,
        **extra_impl_args,
    ) -> None:
        raise NotImplementedError

    @classmethod
    def supports_kv_cache_dtype(cls, kv_cache_dtype: str | None, platform_key: str) -> bool:
        if kv_cache_dtype is None:
            return True
        return kv_cache_dtype in cls._supported_kv_cache_dtypes.get(platform_key, set())

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        """Dispatch to platform-specific forward implementation."""
        if current_omni_platform.is_rocm():
            return self.forward_hip(query, key, value, attn_metadata)
        elif current_omni_platform.is_cuda():
            return self.forward_cuda(query, key, value, attn_metadata)
        elif current_omni_platform.is_npu():
            return self.forward_npu(query, key, value, attn_metadata)
        elif current_omni_platform.is_xpu():
            return self.forward_xpu(query, key, value, attn_metadata)
        elif current_omni_platform.is_musa():
            return self.forward_musa(query, key, value, attn_metadata)
        else:
            raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")

    def forward_paged(self, paged_kv_context: Any) -> torch.Tensor:
        """Execute one Worker-prepared paged-KV attention call.

        The context is intentionally opaque to the common attention layer.
        Backends opt in by setting ``supports_paged_kv`` on their backend
        class and implementing this method.  Dense callers continue to use
        ``forward`` unchanged.
        """

        del paged_kv_context
        raise NotImplementedError(f"{type(self).__name__} does not support paged KV attention")

    def forward_cuda(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward_npu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward_xpu(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        raise NotImplementedError

    def forward_hip(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        # By default, HIP ops are compatible with CUDA ops.
        return self.forward_cuda(query, key, value, attn_metadata)

    def forward_musa(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: T | None = None,
    ) -> torch.Tensor:
        # By default, MUSA ops are compatible with CUDA ops.
        return self.forward_cuda(query, key, value, attn_metadata)