261 lines
10 KiB
Python
261 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Per-shape CUDA graph for one AuK DiT denoise step."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import OrderedDict
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.math_utils import round_up
|
|
|
|
from vllm_omni.diffusion.models.auk.auk_transformer import AuKStepContext, AuKTransformer
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class _GraphEntry:
|
|
graph: torch.cuda.CUDAGraph
|
|
static_x: torch.Tensor
|
|
static_ctx: AuKStepContext
|
|
static_timestep: torch.Tensor
|
|
static_cfg: torch.Tensor | None
|
|
static_step: torch.Tensor | None
|
|
static_out: torch.Tensor
|
|
|
|
|
|
class AuKCUDAGraphWrapper:
|
|
"""Replay one DiT denoise step and leave Euler scheduling to the caller.
|
|
|
|
The graph is keyed by the target, text and reference sequence lengths,
|
|
whether the CFG branch is enabled and, when the caller passes the time
|
|
grid, its length. Timestep (or the step index), and the CFG strength are
|
|
mutable scalar buffers, so all Euler steps and CFG values within one path
|
|
reuse one graph. With the time grid, the adaLN modulations of every step
|
|
are part of the per-request context and the graph selects its row.
|
|
|
|
Only :meth:`AuKTransformer.step` is captured. The per-request context
|
|
(text projection, reference embedding, padding biases, rotary tables) is
|
|
built by :meth:`AuKTransformer.prepare` on the first step of a request and
|
|
copied into the graph's static context, so the replayed steps skip it.
|
|
"""
|
|
|
|
_TARGET_ALIGNMENT = 32
|
|
_TEXT_ALIGNMENT = 32
|
|
_REF_ALIGNMENT = 50
|
|
|
|
def __init__(self, dit: AuKTransformer, *, enabled: bool = True, max_graphs: int = 32) -> None:
|
|
self.dit = dit
|
|
self.enabled = bool(enabled)
|
|
self.max_graphs = max(1, int(max_graphs))
|
|
self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict()
|
|
self._pool_handle: int | None = None
|
|
# Key whose static context holds the current request, and the context
|
|
# the eager fallback reuses across the steps of one request.
|
|
self._active_key: tuple | None = None
|
|
self._eager_ctx: AuKStepContext | None = None
|
|
logger.info("Initialized AuK DiT lazy CUDA graph cache: enabled=%s slots=%d", self.enabled, self.max_graphs)
|
|
|
|
@classmethod
|
|
def _bucket_inputs(
|
|
cls,
|
|
x: torch.Tensor,
|
|
text: torch.Tensor,
|
|
c_mask: torch.Tensor,
|
|
ref: torch.Tensor,
|
|
ref_mask: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
target_bucket = round_up(x.shape[1], cls._TARGET_ALIGNMENT)
|
|
text_bucket = round_up(text.shape[1], cls._TEXT_ALIGNMENT)
|
|
ref_bucket = round_up(ref.shape[1], cls._REF_ALIGNMENT)
|
|
x_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device)
|
|
return (
|
|
F.pad(x, (0, 0, 0, target_bucket - x.shape[1])),
|
|
F.pad(x_mask, (0, target_bucket - x_mask.shape[1]), value=False),
|
|
F.pad(text, (0, 0, 0, text_bucket - text.shape[1])),
|
|
F.pad(c_mask, (0, text_bucket - c_mask.shape[1]), value=False),
|
|
F.pad(ref, (0, 0, 0, ref_bucket - ref.shape[1])),
|
|
F.pad(ref_mask, (0, ref_bucket - ref_mask.shape[1]), value=False),
|
|
)
|
|
|
|
@staticmethod
|
|
def _key(
|
|
x: torch.Tensor,
|
|
text: torch.Tensor,
|
|
ref: torch.Tensor,
|
|
uses_cfg: bool,
|
|
steps: int = 0,
|
|
) -> tuple[int, int, int, bool, int]:
|
|
return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg, steps)
|
|
|
|
def _retire_graph_generation_if_full(self) -> None:
|
|
"""Retire all graphs together so none outlive shared workspaces."""
|
|
if len(self._cache) >= self.max_graphs:
|
|
self._cache.clear()
|
|
self._active_key = None
|
|
|
|
def _prepare(
|
|
self,
|
|
x: torch.Tensor,
|
|
x_mask: torch.Tensor | None,
|
|
text: torch.Tensor,
|
|
c_mask: torch.Tensor,
|
|
ref: torch.Tensor,
|
|
ref_mask: torch.Tensor,
|
|
uses_cfg: bool,
|
|
timesteps: torch.Tensor | None = None,
|
|
) -> AuKStepContext:
|
|
return self.dit.prepare(
|
|
text,
|
|
target_len=x.shape[1],
|
|
mask=x_mask,
|
|
c_mask=c_mask,
|
|
ref=ref,
|
|
ref_mask=ref_mask,
|
|
cfg_infer=uses_cfg,
|
|
timesteps=timesteps,
|
|
)
|
|
|
|
def _step(
|
|
self,
|
|
x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
ctx: AuKStepContext,
|
|
cfg_strength: torch.Tensor | float,
|
|
step_index: torch.Tensor | int | None = None,
|
|
) -> torch.Tensor:
|
|
"""One guided (or unguided) velocity; the CFG combine is part of the graph."""
|
|
velocity = self.dit.step(x, timestep, ctx, step_index=step_index)
|
|
if ctx.branches == 1:
|
|
return velocity
|
|
conditional, unconditional = velocity.chunk(2, dim=0)
|
|
return conditional + (conditional - unconditional) * cfg_strength
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
*,
|
|
x: torch.Tensor,
|
|
text: torch.Tensor,
|
|
c_mask: torch.Tensor,
|
|
ref: torch.Tensor,
|
|
ref_mask: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
cfg_strength: float,
|
|
new_request: bool = True,
|
|
timesteps: torch.Tensor | None = None,
|
|
step_index: int | None = None,
|
|
) -> torch.Tensor:
|
|
"""Velocity of ``x`` at ``timestep``.
|
|
|
|
``new_request`` marks the first step of a request, whose conditioning
|
|
(``text``, ``c_mask``, ``ref``, ``ref_mask``) is then prepared anew; the
|
|
later steps of the request pass ``False`` and reuse it. The default
|
|
prepares on every call, which is always correct.
|
|
|
|
``timesteps`` (the times every step of the request runs at) and
|
|
``step_index`` (this step's position in them) let the adaLN
|
|
modulations be computed once per request instead of once per step.
|
|
"""
|
|
uses_cfg = cfg_strength >= 1e-5
|
|
if timesteps is None or step_index is None:
|
|
timesteps, step_index = None, None
|
|
if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
|
|
if new_request or self._eager_ctx is None:
|
|
self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
|
|
return self._step(x, timestep, self._eager_ctx, cfg_strength, step_index)
|
|
|
|
target_frames = x.shape[1]
|
|
x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask)
|
|
steps = 0 if timesteps is None else int(timesteps.numel())
|
|
key = self._key(x, text, ref, uses_cfg, steps)
|
|
entry = self._cache.get(key)
|
|
ctx = None
|
|
if entry is None or new_request or key != self._active_key:
|
|
ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
|
|
if entry is None:
|
|
# Same Retirement as #6587, details discussed in #7469
|
|
self._retire_graph_generation_if_full()
|
|
try:
|
|
assert ctx is not None
|
|
entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength, step_index=step_index)
|
|
except Exception:
|
|
self._cache.clear()
|
|
self._active_key = None
|
|
self.enabled = False
|
|
logger.exception(
|
|
"Disabling AuK DiT CUDA graphs after capture failure for key=%s; "
|
|
"subsequent requests will use eager execution.",
|
|
key,
|
|
)
|
|
raise
|
|
self._cache[key] = entry
|
|
else:
|
|
self._cache.move_to_end(key)
|
|
if ctx is not None:
|
|
entry.static_ctx.copy_(ctx)
|
|
self._active_key = key
|
|
|
|
entry.static_x.copy_(x)
|
|
entry.static_timestep.copy_(timestep)
|
|
if entry.static_step is not None:
|
|
entry.static_step.fill_(step_index)
|
|
if entry.static_cfg is not None:
|
|
entry.static_cfg.fill_(cfg_strength)
|
|
entry.graph.replay()
|
|
return entry.static_out[:, :target_frames].clone()
|
|
|
|
def _capture(
|
|
self,
|
|
x: torch.Tensor,
|
|
ctx: AuKStepContext,
|
|
timestep: torch.Tensor,
|
|
*,
|
|
cfg_strength: float,
|
|
step_index: int | None = None,
|
|
) -> _GraphEntry:
|
|
static_x = x.clone()
|
|
static_ctx = ctx.clone()
|
|
static_timestep = timestep.clone()
|
|
static_step = None
|
|
if static_ctx.modulation is not None and step_index is not None:
|
|
static_step = torch.full((), step_index, device=static_x.device, dtype=torch.long)
|
|
static_cfg = None
|
|
if ctx.branches == 2:
|
|
static_cfg = torch.empty((), device=static_x.device, dtype=torch.float32)
|
|
static_cfg.fill_(cfg_strength)
|
|
scale = static_cfg if static_cfg is not None else cfg_strength
|
|
# Warm-up runs any lazy torch.compile outside the capture.
|
|
for _ in range(3):
|
|
self._step(static_x, static_timestep, static_ctx, scale, static_step)
|
|
if self._pool_handle is None:
|
|
self._pool_handle = torch.cuda.graph_pool_handle()
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(graph, pool=self._pool_handle):
|
|
# The CFG scalar stays a mutable buffer, so cfg=2 and cfg=3 share this graph.
|
|
static_out = self._step(static_x, static_timestep, static_ctx, scale, static_step)
|
|
|
|
logger.info(
|
|
"Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s",
|
|
x.shape[1],
|
|
ctx.c.shape[1],
|
|
0 if ctx.prompt is None else ctx.prompt.shape[1],
|
|
cfg_strength,
|
|
)
|
|
return _GraphEntry(
|
|
graph=graph,
|
|
static_x=static_x,
|
|
static_ctx=static_ctx,
|
|
static_timestep=static_timestep,
|
|
static_cfg=static_cfg,
|
|
static_step=static_step,
|
|
static_out=static_out,
|
|
)
|
|
|
|
|
|
__all__ = ["AuKCUDAGraphWrapper"]
|