190 lines
7.3 KiB
Python
190 lines
7.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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
|
|
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import 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.distributed.parallel_state import get_sp_group
|
|
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
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,
|
|
):
|
|
super().__init__()
|
|
|
|
self.role = role
|
|
self.role_category = role_category
|
|
self.qkv_layout = qkv_layout
|
|
|
|
# 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
|
|
|
|
config = get_current_diffusion_config_or_none()
|
|
attention_config = config.diffusion_attention_config if config is not None else None
|
|
|
|
attn_backend_cls, spec = get_attn_backend_for_role(
|
|
role=role,
|
|
head_size=head_size,
|
|
attention_config=attention_config,
|
|
role_category=role_category,
|
|
)
|
|
if spec is not None:
|
|
backend_kwargs = spec.extra or None
|
|
self.backend_pref = spec.backend
|
|
logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
|
|
else:
|
|
logger.debug("Attention(role=%s) → platform default", role)
|
|
|
|
self.attn_backend = attn_backend_cls
|
|
self.attn_impl_cls = self.attn_backend.get_impl_cls()
|
|
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,
|
|
backend_kwargs=backend_kwargs,
|
|
)
|
|
# Instantiate fallback backend for float32 support
|
|
self.sdpa_fallback = 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,
|
|
)
|
|
|
|
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.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
|
|
try:
|
|
sp_group = get_sp_group()
|
|
self.ring_pg = sp_group.ring_group
|
|
self.ring_runner = RingParallelAttention(
|
|
sp_group,
|
|
attn_backend_pref=self.backend_pref,
|
|
)
|
|
except Exception:
|
|
self.use_ring = False
|
|
self.ring_runner = None
|
|
|
|
self.parallel_strategy = build_parallel_attention_strategy(
|
|
scatter_idx=scatter_idx,
|
|
gather_idx=gather_idx,
|
|
use_sync=use_sync,
|
|
)
|
|
# Fallback strategy when SP is not active (outside sharded regions)
|
|
self._no_parallel_strategy = NoParallelAttention()
|
|
|
|
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 forward(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AttentionMetadata = None,
|
|
) -> torch.Tensor:
|
|
# Get the appropriate parallel strategy based on SP active state
|
|
strategy = self._get_active_parallel_strategy()
|
|
|
|
# 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)
|
|
|
|
# 2. Kernel Execution (Computation)
|
|
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)
|
|
|
|
# 3. Post-processing (Reverse Communication)
|
|
# For Ulysses: AllToAll Output, and AllGather Joint Output
|
|
out = strategy.post_attention(out, ctx)
|
|
|
|
return out
|
|
|
|
def _run_local_attention(self, query, key, value, attn_metadata):
|
|
if query.dtype == torch.float32:
|
|
logger.warning_once(
|
|
f"Only SDPA supports float32. Overriding user config {type(self.attention)} "
|
|
f"attention_backend='{self.backend_pref}' to 'sdpa' for dtype={query.dtype}."
|
|
)
|
|
return self.sdpa_fallback.forward(query, key, value, attn_metadata)
|
|
|
|
# Fallback to standard attention
|
|
return self.attention.forward(query, key, value, attn_metadata)
|
|
|
|
def _run_ring_attention(self, query, key, value, attn_metadata):
|
|
# 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")
|