126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.parallel.base import (
|
|
ParallelAttentionContext,
|
|
)
|
|
from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
|
|
from vllm_omni.diffusion.distributed.group_coordinator import (
|
|
SequenceParallelGroupCoordinator,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _VideoToAudioContext(ParallelAttentionContext):
|
|
ulysses_pg: dist.ProcessGroup
|
|
use_sync: bool
|
|
|
|
|
|
class LTX2VideoToAudioParallelAttention:
|
|
"""LTX video-to-audio SP with replicated audio queries.
|
|
|
|
Video K/V arrive sequence-sharded, while the much shorter audio query is
|
|
replicated. Redistribute only K/V across sequence/head dimensions, slice
|
|
audio query heads locally, then gather output heads. This uses three
|
|
collectives instead of standard Ulysses' four and keeps audio replicated.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sp_group: SequenceParallelGroupCoordinator,
|
|
*,
|
|
use_sync: bool = False,
|
|
) -> None:
|
|
if sp_group.ring_world_size != 1:
|
|
raise NotImplementedError(
|
|
"LTX video-to-audio sequence parallelism currently supports pure Ulysses only (ring_degree must be 1)."
|
|
)
|
|
self._sp_group = sp_group
|
|
self._ulysses_pg = sp_group.ulysses_group
|
|
self._use_sync = use_sync
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return True
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return "ltx2_video_to_audio"
|
|
|
|
def pre_attention(
|
|
self,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
attn_metadata: AttentionMetadata | None,
|
|
):
|
|
if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
|
|
raise ValueError(
|
|
"LTX video-to-audio SP expects Q/K/V in (B, S, H, D) layout, "
|
|
f"got {query.shape=}, {key.shape=}, {value.shape=}."
|
|
)
|
|
if key.shape != value.shape:
|
|
raise ValueError(
|
|
"LTX video-to-audio SP requires matching K/V shapes, "
|
|
f"got key={tuple(key.shape)} and value={tuple(value.shape)}."
|
|
)
|
|
|
|
world_size = self._sp_group.ulysses_world_size
|
|
rank = self._sp_group.ulysses_rank
|
|
query_heads = int(query.shape[2])
|
|
if query_heads % world_size != 0:
|
|
raise ValueError(
|
|
"LTX video-to-audio SP requires query heads divisible by "
|
|
f"ulysses_degree, got heads={query_heads}, ulysses_degree={world_size}."
|
|
)
|
|
if int(key.shape[2]) % world_size != 0:
|
|
raise ValueError(
|
|
"LTX video-to-audio SP requires K/V heads divisible by "
|
|
f"ulysses_degree, got heads={key.shape[2]}, ulysses_degree={world_size}."
|
|
)
|
|
|
|
heads_per_rank = query_heads // world_size
|
|
head_start = rank * heads_per_rank
|
|
query = query[:, :, head_start : head_start + heads_per_rank, :].contiguous()
|
|
key = SeqAllToAll4D.apply(self._ulysses_pg, key, 2, 1, self._use_sync)
|
|
value = SeqAllToAll4D.apply(self._ulysses_pg, value, 2, 1, self._use_sync)
|
|
|
|
ctx = _VideoToAudioContext(
|
|
name=self.name,
|
|
ulysses_pg=self._ulysses_pg,
|
|
use_sync=self._use_sync,
|
|
)
|
|
return query, key, value, attn_metadata, ctx
|
|
|
|
def post_attention(
|
|
self,
|
|
attn_output: torch.Tensor,
|
|
ctx: ParallelAttentionContext | None,
|
|
) -> torch.Tensor:
|
|
if not isinstance(ctx, _VideoToAudioContext):
|
|
raise TypeError(f"Unexpected LTX video-to-audio context: {type(ctx)!r}.")
|
|
|
|
# Gather along heads by moving H to the leading dimension required by
|
|
# all_gather_into_tensor, then restore the backend's (B, S, H, D) layout.
|
|
local_heads_first = attn_output.permute(2, 0, 1, 3).contiguous()
|
|
world_size = dist.get_world_size(ctx.ulysses_pg)
|
|
gathered = torch.empty(
|
|
(local_heads_first.shape[0] * world_size, *local_heads_first.shape[1:]),
|
|
dtype=attn_output.dtype,
|
|
device=attn_output.device,
|
|
)
|
|
dist.all_gather_into_tensor(gathered, local_heads_first, group=ctx.ulysses_pg)
|
|
if ctx.use_sync:
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
current_omni_platform.synchronize()
|
|
return gathered.permute(1, 2, 0, 3).contiguous()
|