vllm-omni/vllm_omni/diffusion/attention/layer.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
# Copyright (c) Microsoft Corporation and Jiarui Fang
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team & Jiarui Fang
# Adapted from
# https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py


from collections.abc import Mapping
from dataclasses import replace
from typing import cast

import torch
import torch.nn as nn
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.models.utils import extract_layer_index
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec

from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
from vllm_omni.diffusion.forward_context import (
    get_forward_context,
    get_ulysses_mode,
    is_forward_context_available,
)
from vllm_omni.platforms import current_omni_platform

logger = init_logger(__name__)


def _try_extract_layer_index(prefix: str) -> int | None:
    if not prefix:
        return None
    try:
        return extract_layer_index(prefix)
    except (AssertionError, ValueError):
        return None


class Attention(nn.Module):
    _scheduler_paged_kv = False
    _has_custom_attention = False

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        causal: bool,
        softmax_scale: float,
        num_kv_heads: int | None = None,
        prefix: str = "",
        # Per-role backend selection (RFC: per-role attention backend)
        role: str = "self",
        role_category: str | None = None,
        # Model-defined Q/K/V tensor layout hint for backend execution.
        qkv_layout: str | None = None,
        # ulysses attention
        scatter_idx: int = 2,
        gather_idx: int = 1,
        use_sync: bool = False,
        skip_sequence_parallel: bool = False,
        # Opt-out for KV-cache quantization at this specific attention layer.
        # Set by the model author when quant is known to degrade quality or
        # perf for this layer (e.g. Wan2.2 cross-attn has short sequences and
        # block-FP8 quant offers no win). Default False = follow global config.
        disable_kv_quant: bool = False,
        # Opt-in marker for Scheduler-managed paged KV. Unmarked diffusion
        # attention remains dense and contributes no native KVCacheSpec.
        paged_kv_cache_role: str | None = None,
        paged_kv_cache_dtype: torch.dtype | None = None,
        # Model-owned kernel for architectures whose attention contract cannot
        # be represented by the generic backend interface (for example packed
        # varlen attention with learned sink logits). The shared Attention
        # layer still owns parallel dispatch and compile boundaries.
        custom_attention: nn.Module | None = None,
        # Preserve dense FP32 inference for models opting into CUDA auto fallback.
        allow_fp32_fallback: bool = False,
        # Model-owned implementation subclasses, keyed by backend name.
        # Keep the platform-selected backend and its capabilities unchanged.
        # Each override must preserve the selected implementation's contract.
        impl_overrides: Mapping[str, type[AttentionImpl]] | None = None,
    ):
        super().__init__()

        self.role = role
        self.role_category = role_category
        self.qkv_layout = qkv_layout
        # ``prefix`` is also the stable layer identity used by vLLM's native
        # KV-cache metadata.  Keep it on the Omni layer so the active paged
        # adapter can dispatch the already-resharded Q/K/V to the matching
        # native cache tensor without replacing this Omni execution path.
        self.prefix = prefix
        if paged_kv_cache_role == "":
            raise ValueError("paged_kv_cache_role must be non-empty when provided")
        self.paged_kv_cache_role = paged_kv_cache_role
        self.paged_kv_cache_dtype = paged_kv_cache_dtype
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
        self.head_size = head_size

        self._has_custom_attention = custom_attention is not None

        # Resolve backend via role-aware config.
        # The global diffusion config is set during model init via
        # set_current_diffusion_config(); no env-var re-parsing needed here.
        backend_kwargs: dict | None = None
        self.backend_pref = None
        self.backend_explicit = False

        config = get_current_diffusion_config_or_none()
        attention_config = config.diffusion_attention_config if config is not None else None

        from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata

        model_class_name = getattr(config, "model_class_name", None) if config is not None else None
        allow_trtllm_default = get_diffusion_model_metadata(model_class_name).attention_mask_free

        scheduler_paged_kv = (
            config is not None
            and getattr(config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
            is DiffusionKVCacheMode.PAGED_SCHEDULER
        )
        self._scheduler_paged_kv = scheduler_paged_kv
        if custom_attention is None:
            attn_backend_cls, spec = get_attn_backend_for_role(
                role=role,
                head_size=head_size,
                attention_config=attention_config,
                role_category=role_category,
                allow_trtllm_default=allow_trtllm_default,
            )
            if (
                scheduler_paged_kv
                and paged_kv_cache_role is not None
                and spec is None
                and not attn_backend_cls.supports_paged_kv
            ):
                # FLASH_ATTN is an Omni selector, not a device-specific kernel.
                # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
                # attention on NPU. A platform may legitimately default dense
                # attention to CUDNN/SDPA (for example when optional dense FA
                # extras are absent), but a paged layer must advertise the
                # capability so the Worker can register its native cache view.
                # Load that selector from the registry: this is not a user-explicit
                # dense FLASH_ATTN request, so skip platform explicit-backend
                # validation (Blackwell would otherwise require CuTe FA4). Formal
                # paged execution still delegates to vLLM's native paged backend.
                from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum

                dense_backend_name = attn_backend_cls.get_name()
                attn_backend_cls = DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
                logger.info(
                    "Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
                    "does not support Scheduler-owned KV",
                    role,
                    attn_backend_cls.get_name(),
                    dense_backend_name,
                )
            parallel_config = getattr(config, "parallel_config", None)
            allgather_degree = getattr(parallel_config, "allgather_degree", 1)
            # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
            # so validation does not depend on backend names.
            if not skip_sequence_parallel and allgather_degree > 1 and attn_backend_cls.get_name() == "TRTLLM_ATTN":
                raise ValueError(
                    "TRTLLM_ATTN does not support AllGather-KV sequence parallelism. "
                    "Set --allgather-degree 1 or select another diffusion attention backend."
                )
            self.attn_spec = spec
            if spec is not None:
                backend_kwargs = spec.backend_kwargs()
                self.backend_pref = spec.backend
                self.backend_explicit = True
                logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
            else:
                # Propagate the resolved platform default so Ring Attention can
                # make a compatible automatic selection on the current GPU.
                self.backend_pref = attn_backend_cls.get_name()
                logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref)

            self.attn_backend: type[AttentionBackend] | None = attn_backend_cls
            self.attn_impl_cls = self.attn_backend.get_impl_cls()
            if impl_overrides is not None:
                override = impl_overrides.get(attn_backend_cls.get_name())
                if override is not None:
                    if not issubclass(override, self.attn_impl_cls):
                        raise TypeError(
                            f"Attention implementation override {override.__qualname__} must subclass "
                            f"the selected implementation {self.attn_impl_cls.__qualname__} "
                            f"for backend {attn_backend_cls.__qualname__}"
                        )
                    self.attn_impl_cls = override
            self.attention = self.attn_impl_cls(
                num_heads=num_heads,
                head_size=head_size,
                softmax_scale=softmax_scale,
                causal=causal,
                num_kv_heads=num_kv_heads,
                qkv_layout=qkv_layout,
                prefix=prefix,
                backend_kwargs=backend_kwargs,
                role=role,
                backend_explicit=self.backend_explicit,
            )
            # Compatibility kernels run inside shared dispatch, between the
            # parallel strategy's input preparation and output restoration.
            self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()(
                num_heads=num_heads,
                head_size=head_size,
                softmax_scale=softmax_scale,
                causal=causal,
                num_kv_heads=num_kv_heads,
                qkv_layout=qkv_layout,
            )
        else:
            if paged_kv_cache_role is not None:
                raise ValueError("custom_attention does not support Scheduler-managed paged KV")
            if not skip_sequence_parallel:
                raise ValueError("custom_attention must own its communication and requires skip_sequence_parallel=True")
            self.attn_spec = None
            self.attn_backend = None
            self.attn_impl_cls = type(custom_attention)
            self.attention = custom_attention
            self.sdpa_fallback = None
            logger.debug("Attention(role=%s) → custom kernel=%s", role, type(custom_attention).__name__)

        self.softmax_scale = softmax_scale
        self.scatter_idx = scatter_idx
        self.gather_idx = gather_idx
        self.use_sync = use_sync
        self.causal = causal
        self.skip_sequence_parallel = skip_sequence_parallel
        self.allow_fp32_fallback = allow_fp32_fallback

        self.use_ring = False
        self.ring_pg = None
        self.ring_runner = None

        if config is not None:
            if config.parallel_config.ring_degree > 1:
                self.use_ring = True
                sp_group = get_sp_group()
                self.ring_pg = sp_group.ring_group
                self.ring_runner = RingParallelAttention(
                    sp_group,
                    attn_backend_pref=self.backend_pref,
                    attn_backend_explicit=self.backend_explicit,
                )

        self.parallel_strategy = build_parallel_attention_strategy(
            scatter_idx=scatter_idx,
            gather_idx=gather_idx,
            use_sync=use_sync,
            causal=causal,
        )
        # Local strategy when SP is intentionally inactive outside sharded regions.
        self._no_parallel_strategy = NoParallelAttention()

        self.layer_idx: int | None = _try_extract_layer_index(prefix)

        self._kv_cache_dtype: str | None = None
        self._kv_cache_skip_steps: set[int] | None = None
        self._kv_cache_skip_layers: set[int] | None = None
        # Per-layer opt-out from KV-cache quantization (set by model author).
        self._disable_kv_quant: bool = disable_kv_quant
        self._init_kv_cache_quantization(config)

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
        """Return native rank-local geometry for an opted-in paged cache."""

        if self.paged_kv_cache_role is None:
            return None
        dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
        # Keep backend layout discovery under the same config context used by
        # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
        # block-stride decision off the spec and onto the single physical
        # ``CacheConfig.kv_cache_layout`` resolved before the cache is built,
        # so the backend's preference is enforced there (see
        # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than
        # carried per layer.
        with set_current_vllm_config(vllm_config):
            assert_backend_layout_supported(vllm_config, self.attn_backend)
        return FullAttentionSpec(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=self.num_kv_heads,
            head_size=self.head_size,
            dtype=dtype,
            non_causal=not self.causal,
        )

    def _get_active_parallel_strategy(self):
        """Get the parallel strategy based on current SP active state.

        Returns NoParallelAttention if we're outside an SP sharded region
        (e.g., in noise_refiner/context_refiner before unified_prepare in Z-Image).
        This avoids unnecessary SP communication for layers not covered by _sp_plan.
        """
        if self.skip_sequence_parallel:
            return self._no_parallel_strategy
        if is_forward_context_available():
            ctx = get_forward_context()
            if not ctx.sp_active:
                return self._no_parallel_strategy
        return self.parallel_strategy

    def _init_kv_cache_quantization(self, config) -> None:
        if config is None or self._has_custom_attention:
            return
        assert self.attn_backend is not None
        dtype = getattr(config, "diffusion_kv_cache_dtype", None)
        if dtype == "auto":
            dtype = None
        parallel_config = getattr(config, "parallel_config", None)
        ring_degree = getattr(parallel_config, "ring_degree", 1)
        if dtype and dtype != "float":
            if ring_degree > 1:
                raise ValueError(
                    "KV quantization is not compatible with ring attention "
                    "(ring_degree > 1). Ring kernels do not propagate quantization descale "
                    "factors. Use Ulysses SP instead."
                )
            platform_key = current_omni_platform.device_name
            if not self.attention.supports_kv_cache_dtype(dtype, platform_key):
                raise ValueError(
                    f"Attention backend {self.attn_backend.get_name()} does not support "
                    f"kv_cache_dtype={dtype!r} on {platform_key}. Select a compatible "
                    "backend or set diffusion_kv_cache_dtype='auto'."
                )
        self._kv_cache_dtype = dtype
        self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None)
        self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None)
        if self._kv_cache_skip_layers and self.layer_idx is None and not self._disable_kv_quant:
            raise ValueError("Attention quantization skip_layers requires a parseable transformer block index.")

    def _should_apply_kv_cache_quant(self) -> bool:
        skip_steps = self._kv_cache_skip_steps
        skip_layers = self._kv_cache_skip_layers
        if skip_steps is not None:
            step_idx = get_forward_context().denoise_step_idx if is_forward_context_available() else None
            if skip_steps and (step_idx is None or step_idx in skip_steps):
                return False
        if skip_layers is not None:
            if self.layer_idx is not None and self.layer_idx in skip_layers:
                return False
        return True

    def _with_kv_cache_dtype(self, attn_metadata: AttentionMetadata | None) -> AttentionMetadata | None:
        disabled = self._disable_kv_quant or not self._should_apply_kv_cache_quant()
        dtype = self._kv_cache_dtype
        if dtype in (None, "float"):
            dtype = None
        elif disabled:
            dtype = "float"
        if dtype is None and (attn_metadata is None or "kv_cache_dtype" not in attn_metadata.extra):
            return attn_metadata
        extra = dict(attn_metadata.extra) if attn_metadata is not None else {}
        # Recompute per forward so shared metadata cannot retain another step's policy.
        extra.pop("kv_cache_dtype", None)
        if dtype is not None:
            extra["kv_cache_dtype"] = dtype
        if attn_metadata is None:
            return AttentionMetadata(extra=extra) if extra else None
        return replace(attn_metadata, extra=extra)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        if torch.compiler.is_compiling() and is_forward_context_available():
            od_config = get_forward_context().omni_diffusion_config
            parallel_config = getattr(od_config, "parallel_config", None)
            if getattr(parallel_config, "use_hsdp", False):
                # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
                # attention graph; otherwise scheduler dependency analysis can
                # fail on the fused attention region.
                return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)

        return self._forward_impl(query, key, value, attn_metadata)

    @torch.compiler.disable
    def _forward_hsdp_compile_boundary(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        return self._forward_impl(query, key, value, attn_metadata)

    def _forward_impl(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attn_metadata: AttentionMetadata | None = None,
    ) -> torch.Tensor:
        # Get the appropriate parallel strategy based on SP active state
        strategy = self._get_active_parallel_strategy()
        paged_adapter = self._active_paged_kv_adapter()
        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
        if (
            self._scheduler_paged_kv
            and self.paged_kv_cache_role is not None
            and paged_adapter is None
            and not in_kv_memory_profile
        ):
            raise RuntimeError(
                "Scheduler-paged diffusion attention reached model forward without an active Worker adapter. "
                "Only the startup KV memory profile may execute before paged KV initialization."
            )
        use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
        if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
            assert self.attn_backend is not None
            raise NotImplementedError(
                f"Diffusion paged KV requires an Omni backend with paged support; "
                f"selected {self.attn_backend.get_name()}"
            )
        if use_paged_attention and strategy is not self._no_parallel_strategy:
            strategy_name = strategy.name
            if self.use_ring or strategy_name == "ring":
                raise NotImplementedError(
                    "paged Scheduler KV is not supported with Ring attention; use strict Ulysses or no SP"
                )
            if strategy_name == "allgather_kv":
                raise NotImplementedError("paged Scheduler KV is not supported with AllGather-KV sequence parallelism")
            if strategy_name == "ulysses" and get_ulysses_mode(default="strict") != "strict":
                raise NotImplementedError("paged Scheduler KV currently supports only strict Ulysses")

        # 1. Prepare inputs (Communication / Resharding)
        # For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v
        # For Ring: Concat joint_q
        query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata)

        # Scheduler rows describe the logical sequence, while strict Ulysses
        # may append synthetic tokens solely to make the image shard divisible.
        # Remove those tokens after the all-to-all and put zero placeholders
        # back before the reverse all-to-all.
        paged_sp_padding = 0
        paged_sp_padding_offset = query.shape[1]
        if use_paged_attention and strategy is not self._no_parallel_strategy and strategy.name == "ulysses":
            forward_ctx = get_forward_context() if is_forward_context_available() else None
            paged_sp_padding = int(getattr(forward_ctx, "sp_padding_size", 0))
            if paged_sp_padding:
                joint_len = int(getattr(ctx, "joint_len", 0))
                joint_strategy = str(getattr(ctx, "joint_strategy", "front"))
                paged_sp_padding_offset = query.shape[1] - (joint_len if joint_strategy == "rear" else 0)
                padding_start = paged_sp_padding_offset - paged_sp_padding
                if padding_start < 0:
                    raise ValueError(
                        "Paged Ulysses padding exceeds the post-all-to-all sequence: "
                        f"padding={paged_sp_padding}, sequence={query.shape[1]}"
                    )

                def _remove_paged_sp_padding(tensor: torch.Tensor) -> torch.Tensor:
                    return torch.cat(
                        (tensor[:, :padding_start], tensor[:, paged_sp_padding_offset:]),
                        dim=1,
                    ).contiguous()

                query = _remove_paged_sp_padding(query)
                key = _remove_paged_sp_padding(key)
                value = _remove_paged_sp_padding(value)

        # 2. This is the shared GPU/NPU boundary. The Worker adapter prepares
        # the native page-table context after SP has produced rank-local Q/K/V;
        # backend resolution below selects CUDA or Ascend execution.
        if use_paged_attention:
            assert paged_adapter is not None
            paged_kv_context = paged_adapter.prepare_layer_context(
                self.prefix,
                query,
                key,
                value,
                omni_attn_metadata=attn_metadata,
            )
            out = self.attention.forward_paged(paged_kv_context)
        else:
            attn_metadata = self._with_kv_cache_dtype(attn_metadata)
            if self.use_ring and strategy is not self._no_parallel_strategy:
                out = self._run_ring_attention(query, key, value, attn_metadata)
            else:
                out = self._run_local_attention(query, key, value, attn_metadata)

        if paged_sp_padding:
            padding_start = paged_sp_padding_offset - paged_sp_padding
            output_padding = out.new_zeros((out.shape[0], paged_sp_padding, *out.shape[2:]))
            out = torch.cat((out[:, :padding_start], output_padding, out[:, padding_start:]), dim=1)

        # 3. Post-processing (Reverse Communication)
        # For Ulysses: AllToAll Output, and AllGather Joint Output
        out = strategy.post_attention(out, ctx)

        return out

    @staticmethod
    def _active_paged_kv_adapter():
        """Return the Worker adapter selected by Runner-owned metadata."""

        if not is_forward_context_available():
            return None
        return getattr(get_forward_context(), "paged_kv_adapter", None)

    def is_paged_kv_active(self) -> bool:
        """Return whether this layer will use Scheduler-managed paged KV."""

        return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None

    def _run_local_attention(self, query, key, value, attn_metadata):
        if self._has_custom_attention:
            return cast(nn.Module, self.attention)(query, key, value, attn_metadata)

        self._assert_metadata_compatible(attn_metadata)

        if (
            self.allow_fp32_fallback
            and query.is_cuda
            and query.dtype == torch.float32
            and query.ndim == 4
            and not self.backend_explicit
            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
            and (
                attn_metadata is None
                or (
                    attn_metadata.full_attn_spans is None
                    and attn_metadata.query_ranges is None
                    and attn_metadata.video_layout is None
                    and attn_metadata.packed_padding is None
                    and not attn_metadata.extra
                )
            )
        ):
            logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.")
            return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata)

        in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
        # The startup KV-capacity profile needs tensor shapes, not a paged
        # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD
        # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged
        # requests never use this branch because their Worker adapter is active.
        # A user-explicit backend must run or raise.
        if (
            self._scheduler_paged_kv
            and self.paged_kv_cache_role is not None
            and in_kv_memory_profile
            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
            and not current_omni_platform.supports_diffusion_dense_flash_attention()
        ):
            logger.warning_once(
                "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                "Formal paged requests still use the platform-native paged attention backend."
            )
            return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata)

        return self.attention.forward(query, key, value, attn_metadata)

    def _assert_metadata_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
        if attn_metadata is None:
            return
        if self.attn_backend is None:
            return
        backend_name = self.attn_backend.get_name()
        if attn_metadata.attn_mask is not None and not self.attn_backend.supports_attention_mask(
            getattr(self, "attn_spec", None)
        ):
            raise ValueError(
                f"Attention backend '{backend_name}' does not support attn_mask. Select a mask-capable backend."
            )
        if attn_metadata.full_attn_spans is None:
            return
        if attn_metadata.attn_mask is not None and attn_metadata.attn_mask.ndim == 4:
            return
        if not self.attn_backend.supports_piecewise_spans:
            raise ValueError(
                f"Attention backend '{backend_name}' does not support "
                f"piecewise attention (full_attn_spans without a 4D attn_mask). "
                f"Use a Flash backend (FLASH_ATTN / FLASH_ATTN_HUB / FLASH_ATTN_3_HUB), "
                f"or provide a 4D attn_mask that encodes the mixed causal/full pattern."
            )

    def _run_ring_attention(self, query, key, value, attn_metadata):
        if attn_metadata is not None and attn_metadata.attn_mask is not None:
            raise ValueError("Ring attention does not support attn_mask; use Ulysses SP or disable mask_sp_padding.")
        skip = getattr(self.attention, "skip", None)
        if skip is not None and getattr(skip, "configured", False):
            raise NotImplementedError(
                "Skip-Softmax (TRTLLM_ATTN) is not supported with ring sequence parallelism: "
                "the ring path bypasses the backend, so the skip config would be silently ignored. "
                "Use Ulysses SP instead, or remove the skip_softmax config."
            )
        # Delegate to RingParallelAttention strategy if available
        if self.ring_runner is not None:
            return self.ring_runner.run_attention(
                query, key, value, attn_metadata, softmax_scale=self.softmax_scale, causal=self.causal
            )

        raise RuntimeError("Ring attention is enabled but strategy is not RingParallelAttention")