Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/diffusion/models/auk/alias_free_triton.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Triton kernels for the AuK codec's alias-free SnakeBeta activation.
+
+The eager module oversamples with a grouped transposed convolution, applies
+SnakeBeta, and decimates with a grouped strided convolution. Both grouped
+convolutions carry one 12-tap filter shared by every channel, and cuDNN and
+ATen serve them with generic depthwise kernels that dominate the decode (more
+than half of a 12 s decode on H200). Here the oversampling is written as its
+polyphase form fused with the activation, and the decimation as a direct
+strided FIR, each one pass over memory per row. The arithmetic is the same
+as the eager module in fp32; only the summation order of the taps differs.
+
+The ops are registered with ``torch.library.triton_op`` so the compiled
+decode buckets trace them instead of breaking the graph.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+__all__ = ["alias_free_snake"]
+
+_BLOCK = 1024
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _upsample_snake_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        exp_alpha_ptr,
+        inv_beta_ptr,
+        channels,
+        in_len,
+        out_len,
+        ratio: tl.constexpr,
+        taps: tl.constexpr,
+        pad: tl.constexpr,
+        crop: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[m] = snake(ratio * sum_k taps[k] * xpad[(m + crop - k) / ratio]) over the k
+        # with (m + crop - k) divisible by ratio, where xpad is x replicate-padded by pad
+        # on both sides. That is conv_transpose1d(stride=ratio) followed by the crop.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            pos = offs + (crop - k)
+            hit = (pos % ratio) == 0
+            src = pos // ratio - pad
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            w = tl.load(taps_ptr + k)
+            val = tl.load(x_row + src, mask=mask & hit, other=0.0)
+            acc += tl.where(hit, w * val, 0.0)
+        u = acc * ratio
+        channel = row % channels
+        alpha = tl.load(exp_alpha_ptr + channel)
+        inv_beta = tl.load(inv_beta_ptr + channel)
+        s = tl.sin(u * alpha)
+        u = u + inv_beta * (s * s)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, u, mask=mask)
+
+    @triton.jit
+    def _lowpass_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        in_len,
+        out_len,
+        stride: tl.constexpr,
+        taps: tl.constexpr,
+        pad_left: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[n] = sum_k taps[k] * x[clamp(stride * n + k - pad_left)], i.e. replicate
+        # padding followed by a strided conv1d.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            src = offs * stride + (k - pad_left)
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            acc += tl.load(taps_ptr + k) * tl.load(x_row + src, mask=mask, other=0.0)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, acc, mask=mask)
+
+    @torch.library.triton_op("vllm_omni::auk_alias_free_snake", mutates_args=())
+    def alias_free_snake(
+        x: torch.Tensor,
+        up_taps: torch.Tensor,
+        down_taps: torch.Tensor,
+        exp_alpha: torch.Tensor,
+        inv_beta: torch.Tensor,
+        up_ratio: int,
+        up_pad: int,
+        up_crop: int,
+        down_stride: int,
+        down_pad_left: int,
+        down_pad_right: int,
+    ) -> torch.Tensor:
+        """Upsample by ``up_ratio``, SnakeBeta, low-pass and decimate: ``[B, C, T] -> [B, C, T_out]``."""
+        batch, channels, length = x.shape
+        x = x.contiguous()
+        up_len = length * up_ratio
+        oversampled = torch.empty((batch, channels, up_len), device=x.device, dtype=torch.float32)
+        rows = batch * channels
+        torch.library.wrap_triton(_upsample_snake_kernel)[(rows, triton.cdiv(up_len, _BLOCK))](
+            x,
+            oversampled,
+            up_taps,
+            exp_alpha,
+            inv_beta,
+            channels,
+            length,
+            up_len,
+            ratio=up_ratio,
+            taps=up_taps.numel(),
+            pad=up_pad,
+            crop=up_crop,
+            block=_BLOCK,
+        )
+        taps = down_taps.numel()
+        out_len = (up_len + down_pad_left + down_pad_right - taps) // down_stride + 1
+        out = torch.empty((batch, channels, out_len), device=x.device, dtype=torch.float32)
+        torch.library.wrap_triton(_lowpass_kernel)[(rows, triton.cdiv(out_len, _BLOCK))](
+            oversampled,
+            out,
+            down_taps,
+            up_len,
+            out_len,
+            stride=down_stride,
+            taps=taps,
+            pad_left=down_pad_left,
+            block=_BLOCK,
+        )
+        return out
+
+else:  # pragma: no cover
+
+    def alias_free_snake(*args, **kwargs):
+        raise RuntimeError("Triton is not available")

vllm_omni/diffusion/models/auk/auk_transformer.py

@@ -153,6 +153,33 @@ class TimeEmbedding(nn.Module):
         return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype))
 
 
+def _grouped_conv1d_btc(x: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor:
+    """``conv`` applied to channels-last ``x`` ``[B, T, C]``, returning ``[B, T, C]``.
+
+    A same-padded grouped convolution written as one im2col copy and one
+    batched GEMM over the groups. cuDNN runs a grouped conv of this shape as
+    one kernel per group plus two layout transposes, which at AuK's sizes
+    (1536 channels in 16 groups, kernel 31, a few hundred frames) costs about
+    ten times the arithmetic.
+    """
+    batch, frames, channels = x.shape
+    groups = conv.groups
+    kernel = conv.kernel_size[0]
+    per_group = channels // groups
+    pad = conv.padding[0]
+    padded = F.pad(x, (0, 0, pad, pad))
+    # unfold gives a [B, T, C, K] view; one copy lays it out as [B, T, G, C/G * K].
+    cols = padded.unfold(1, kernel, 1).reshape(batch * frames, groups, per_group * kernel)
+    # conv.weight [C_out, C/G, K] -> [G, C/G * K, C_out/G]; the reduction order matches the columns.
+    weight = conv.weight.reshape(groups, channels // groups, per_group * kernel).transpose(1, 2).to(cols.dtype)
+    # [G, B*T, C/G*K] @ [G, C/G*K, C_out/G] -> [G, B*T, C_out/G]
+    out = torch.bmm(cols.transpose(0, 1), weight)
+    out = out.transpose(0, 1).reshape(batch, frames, channels)
+    if conv.bias is not None:
+        out = out + conv.bias.to(out.dtype)
+    return out
+
+
 class ConvPosEmbedding(nn.Module):
     """Depthwise-grouped conv stack that adds local position information."""
 
@@ -169,17 +196,19 @@ class ConvPosEmbedding(nn.Module):
         )
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        keep = None if mask is None else mask.unsqueeze(1)
-        x = x.transpose(1, 2)
+        keep = None if mask is None else mask.unsqueeze(-1)
         if keep is not None:
             x = x.masked_fill(~keep, 0.0)
         for layer in self.conv1d:
-            x = layer(x)
-            # Re-zero padding after each convolution so it cannot bleed into
-            # valid frames through the kernel window.
-            if keep is not None and isinstance(layer, nn.Conv1d):
-                x = x.masked_fill(~keep, 0.0)
-        return x.transpose(1, 2)
+            if isinstance(layer, nn.Conv1d):
+                x = _grouped_conv1d_btc(x, layer)
+                # Re-zero padding after each convolution so it cannot bleed into
+                # valid frames through the kernel window.
+                if keep is not None:
+                    x = x.masked_fill(~keep, 0.0)
+            else:
+                x = layer(x)
+        return x
 
 
 class AudioEmbedding(nn.Module):
@@ -205,9 +234,12 @@ class AdaLayerNorm(nn.Module):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
     def forward(
-        self, x: torch.Tensor, emb: torch.Tensor
+        self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1)
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
@@ -221,8 +253,11 @@ class AdaLayerNormFinal(nn.Module):
         self.linear = nn.Linear(dim, dim * 2)
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
-    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1)
+    def forward(self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None) -> torch.Tensor:
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        scale, shift = mod.chunk(2, dim=1)
         return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
 
 
@@ -383,9 +418,11 @@ class DoubleBlock(nn.Module):
         c_rope: RopeCache,
         c_mask: torch.Tensor | None,
         bias: torch.Tensor | None = None,
+        mod_c: torch.Tensor | None = None,
+        mod_x: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t)
-        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t)
+        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t, mod_c)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t, mod_x)
 
         x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask, bias=bias)
 
@@ -416,8 +453,9 @@ class SingleBlock(nn.Module):
         mask: torch.Tensor | None,
         rope: RopeCache,
         bias: torch.Tensor | None = None,
+        mod: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t)
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t, mod)
         x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope, bias=bias)
         norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
         return x + gate_mlp[:, None] * self.ff(norm)
@@ -491,6 +529,28 @@ class AuKTransformer(nn.Module):
 
         self.norm_out = AdaLayerNormFinal(dim)
         self.proj_out = nn.Linear(dim, latent_dim)
+        # Widths of the adaLN modulations in the order modulation_table() lays them out.
+        self._modulation_widths = [module.linear.out_features for module in self._adaln_modules()]
+
+    def _adaln_modules(self) -> list[nn.Module]:
+        """Every timestep-conditioned norm, in the order :meth:`step` consumes them."""
+        modules: list[nn.Module] = []
+        for block in self.transformer_blocks:
+            modules += [block.attn_norm_c, block.attn_norm_x]
+        modules += [block.attn_norm for block in self.single_transformer_blocks]
+        modules.append(self.norm_out)
+        return modules
+
+    def modulation_table(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """adaLN modulations for every timestep, ``[len(timesteps), sum of widths]``.
+
+        The modulations depend only on the timestep, and a request's time grid
+        is known before its first step. Computing them for the whole grid at
+        once reads each modulation weight once per request instead of once per
+        step, which at batch 1 is what those matrix-vector products cost.
+        """
+        emb = F.silu(self.time_embed(timesteps.reshape(-1).float()))
+        return torch.cat([module.linear(emb) for module in self._adaln_modules()], dim=1)
 
     def project_text(self, text: torch.Tensor) -> torch.Tensor:
         """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width."""
@@ -520,6 +580,7 @@ class AuKTransformer(nn.Module):
         drop_text: bool = False,
         cfg_infer: bool = False,
         cache: bool = False,
+        timesteps: torch.Tensor | None = None,
     ) -> "AuKStepContext":
         """Everything a denoise step needs that does not depend on ``x`` or the time.
 
@@ -528,6 +589,9 @@ class AuKTransformer(nn.Module):
         prompt embedding, the padding masks and biases and the rotary tables
         are not recomputed on every step. Arguments are those of
         :meth:`forward`; ``target_len`` is the target frame count ``n``.
+        ``timesteps`` are the times the steps will run at (the grid without
+        its end point); when given, their adaLN modulations are precomputed
+        and :meth:`step` selects them by ``step_index``.
         """
         batch = text.shape[0]
         if c_mask is None:
@@ -593,25 +657,44 @@ class AuKTransformer(nn.Module):
             rope_text=_rope_cos_sin(self.rotary_embed(text_len, c_mask)),
             rope_single=_rope_cos_sin(self.rotary_embed(text_len + audio_len, single_mask)),
             branches=branches,
+            modulation=None if timesteps is None else self.modulation_table(timesteps),
         )
 
-    def step(self, x: torch.Tensor, time: torch.Tensor, ctx: "AuKStepContext") -> torch.Tensor:
+    def step(
+        self,
+        x: torch.Tensor,
+        time: torch.Tensor,
+        ctx: "AuKStepContext",
+        step_index: torch.Tensor | int | None = None,
+    ) -> torch.Tensor:
         """Predict the velocity of ``x`` at ``time`` under a :meth:`prepare` context.
 
         The target embedding does not depend on the CFG branch, so it is
-        computed once and shared by both branches.
+        computed once and shared by both branches. When the context carries a
+        modulation table, ``step_index`` selects this step's row of it and
+        ``time`` is unused; the row broadcasts over the CFG branches.
         """
-        if time.ndim == 0:
-            time = time.repeat(x.shape[0])
-        t = self.time_embed(time)
         target = self.audio_embed(x, mask=ctx.target_mask)
+        if ctx.modulation is not None and step_index is not None:
+            if isinstance(step_index, int):
+                row = ctx.modulation[step_index : step_index + 1]
+            else:
+                row = ctx.modulation.index_select(0, step_index.reshape(1))
+            mods = list(torch.split(row, self._modulation_widths, dim=1))
+            t = None
+        else:
+            if time.ndim == 0:
+                time = time.repeat(x.shape[0])
+            t = self.time_embed(time)
+            if ctx.branches == 2:
+                t = torch.cat((t, t), dim=0)
+            mods = [None] * len(self._modulation_widths)
         if ctx.branches == 2:
-            t = torch.cat((t, t), dim=0)
             target = torch.cat((target, target), dim=0)
         audio = target if ctx.prompt is None else torch.cat([ctx.prompt, target], dim=1)
         c = ctx.c
 
-        for block in self.transformer_blocks:
+        for i, block in enumerate(self.transformer_blocks):
             c, audio = block(
                 audio,
                 c,
@@ -621,14 +704,24 @@ class AuKTransformer(nn.Module):
                 c_rope=ctx.rope_text,
                 c_mask=ctx.c_mask,
                 bias=ctx.joint_bias,
+                mod_c=mods[2 * i],
+                mod_x=mods[2 * i + 1],
             )
 
         h = torch.cat([c, audio], dim=1)
-        for block in self.single_transformer_blocks:
-            h = block(h, t, mask=ctx.single_mask, rope=ctx.rope_single, bias=ctx.single_bias)
+        offset = 2 * len(self.transformer_blocks)
+        for i, block in enumerate(self.single_transformer_blocks):
+            h = block(
+                h,
+                t,
+                mask=ctx.single_mask,
+                rope=ctx.rope_single,
+                bias=ctx.single_bias,
+                mod=mods[offset + i],
+            )
 
         h = h[:, h.shape[1] - x.shape[1] :]
-        return self.proj_out(self.norm_out(h, t))
+        return self.proj_out(self.norm_out(h, t, mods[-1]))
 
     def forward(
         self,
@@ -706,6 +799,8 @@ class AuKStepContext:
     rope_text: RopeCache
     rope_single: RopeCache
     branches: int
+    #: adaLN modulations per step, ``[steps, widths]``, when the time grid was known.
+    modulation: torch.Tensor | None = None
 
     def tensors(self) -> list[torch.Tensor | None]:
         """Every tensor field in a fixed order, rope pairs flattened."""
@@ -721,6 +816,7 @@ class AuKStepContext:
             *self.rope_audio,
             *self.rope_text,
             *self.rope_single,
+            self.modulation,
         ]
 
     def copy_(self, other: "AuKStepContext") -> None:
@@ -752,6 +848,7 @@ class AuKStepContext:
             rope_text=(self.rope_text[0].clone(), self.rope_text[1].clone()),
             rope_single=(self.rope_single[0].clone(), self.rope_single[1].clone()),
             branches=self.branches,
+            modulation=_clone(self.modulation),
         )
 
 
@@ -881,6 +978,8 @@ def sample_latents(
                     timestep=t[i],
                     cfg_strength=cfg_strength,
                     new_request=i == 0,
+                    timesteps=t[:-1],
+                    step_index=i,
                 )
                 x = x + (t[i + 1] - t[i]) * velocity
         finally:
@@ -922,9 +1021,10 @@ def _sample_latents(
             ref_mask=ref_mask,
             cfg_infer=guided,
             cache=guided,
+            timesteps=timesteps[:-1],
         )
         for i in range(timesteps.shape[0] - 1):
-            v = dit.step(x, timesteps[i], ctx)
+            v = dit.step(x, timesteps[i], ctx, step_index=i)
             if guided:
                 v_cond, v_uncond = v.chunk(2, dim=0)
                 v = v_cond + (v_cond - v_uncond) * cfg_strength

vllm_omni/diffusion/models/auk/auk_vae.py

@@ -32,7 +32,9 @@ from torch import nn
 from torch.nn.utils import remove_weight_norm as _fold_weight_norm
 from torch.nn.utils import weight_norm as _apply_weight_norm
 from vllm.logger import init_logger
+from vllm.triton_utils import HAS_TRITON
 
+from vllm_omni.diffusion.models.auk.alias_free_triton import alias_free_snake
 from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta as _SharedSnakeBeta
 
 logger = init_logger(__name__)
@@ -319,8 +321,45 @@ class AliasFreeActivation(nn.Module):
         self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._fused_eligible(x):
+            return self._fused_forward(x)
         return self.downsample(self.act(self.upsample(x)))
 
+    def _fused_eligible(self, x: torch.Tensor) -> bool:
+        """Whether ``x`` can take the Triton kernels in alias_free_triton (CUDA fp32, AuK's filter geometry)."""
+        up = self.upsample
+        taps = up.filter.numel()
+        return (
+            x.is_cuda
+            and x.dtype == torch.float32
+            and x.dim() == 3
+            and isinstance(self.act, SnakeBeta)
+            and HAS_TRITON
+            # The kernel indexes the crop from the start of the transposed
+            # conv output, so the crop must cover the filter's reach, and the
+            # cropped length must be exactly ratio * T.
+            and up.pad_left >= taps - 1
+            and (2 * up.pad - 1) * up.ratio + taps == up.pad_left + up.pad_right
+        )
+
+    def _fused_forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, low, act = self.upsample, self.downsample.lowpass, self.act
+        if not act._cached:
+            act.precompute_exp_cache()
+        return alias_free_snake(
+            x,
+            up.filter.reshape(-1),
+            low.filter.reshape(-1),
+            act._exp_alpha,
+            act._inv_beta,
+            up.ratio,
+            up.pad,
+            up.pad_left,
+            low.stride,
+            low.pad_left,
+            low.pad_right,
+        )
+
 
 class AmpBlock(nn.Module):
     """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs."""

vllm_omni/diffusion/models/auk/cudagraph_wrapper.py

@@ -24,16 +24,19 @@ class _GraphEntry:
     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 plus
-    whether the CFG branch is enabled. Timestep and the CFG strength are
+    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.
+    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
@@ -85,8 +88,9 @@ class AuKCUDAGraphWrapper:
         text: torch.Tensor,
         ref: torch.Tensor,
         uses_cfg: bool,
-    ) -> tuple[int, int, int, bool]:
-        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg)
+        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."""
@@ -103,6 +107,7 @@ class AuKCUDAGraphWrapper:
         ref: torch.Tensor,
         ref_mask: torch.Tensor,
         uses_cfg: bool,
+        timesteps: torch.Tensor | None = None,
     ) -> AuKStepContext:
         return self.dit.prepare(
             text,
@@ -112,6 +117,7 @@ class AuKCUDAGraphWrapper:
             ref=ref,
             ref_mask=ref_mask,
             cfg_infer=uses_cfg,
+            timesteps=timesteps,
         )
 
     def _step(
@@ -120,9 +126,10 @@ class AuKCUDAGraphWrapper:
         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)
+        velocity = self.dit.step(x, timestep, ctx, step_index=step_index)
         if ctx.branches == 1:
             return velocity
         conditional, unconditional = velocity.chunk(2, dim=0)
@@ -140,6 +147,8 @@ class AuKCUDAGraphWrapper:
         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``.
 
@@ -147,26 +156,33 @@ class AuKCUDAGraphWrapper:
         (``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)
-            return self._step(x, timestep, self._eager_ctx, cfg_strength)
+                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)
-        key = self._key(x, text, ref, uses_cfg)
+        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)
+            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)
+                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength, step_index=step_index)
             except Exception:
                 self._cache.clear()
                 self._active_key = None
@@ -186,6 +202,8 @@ class AuKCUDAGraphWrapper:
 
         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()
@@ -198,10 +216,14 @@ class AuKCUDAGraphWrapper:
         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)
@@ -209,13 +231,13 @@ class AuKCUDAGraphWrapper:
         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)
+            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_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",
@@ -230,6 +252,7 @@ class AuKCUDAGraphWrapper:
             static_ctx=static_ctx,
             static_timestep=static_timestep,
             static_cfg=static_cfg,
+            static_step=static_step,
             static_out=static_out,
         )
 

vllm_omni/diffusion/models/auk/pipeline_auk.py

@@ -11,9 +11,11 @@ both, and the BigVGAN-flow decoder renders the waveform.
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
 import time
+from collections import OrderedDict
 from collections.abc import Iterable
 from typing import Any, ClassVar
 
@@ -27,7 +29,12 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    build_time_grid,
+    dit_state_dict,
+    sample_latents,
+)
 from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
 from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph
@@ -59,6 +66,10 @@ _FLASH_CFG = 0.0
 _WARMUP_TARGET_FRAMES = (150, 300, 600)
 _WARMUP_REF_FRAMES = (0, 150)
 _WARMUP_TEXT_TOKENS = 96
+# Reference clips whose posterior-mean latents are kept. Voice cloning reuses a
+# small set of reference voices, so repeat requests skip the VAE encode.
+# ``model_config.auk_ref_cache_size`` overrides it; 0 disables the cache.
+_REF_CACHE_SIZE = 16
 # Decorrelates the VAE posterior draw from the initial latent for the same request seed.
 _VAE_SEED_OFFSET = 0x5EED_0A0C
 
@@ -152,6 +163,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32)
         if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1:
             raise ValueError("AuK max_dit_graphs must be a positive integer")
+        ref_cache_size = od_config.model_config.get("auk_ref_cache_size", _REF_CACHE_SIZE)
+        if isinstance(ref_cache_size, bool) or not isinstance(ref_cache_size, int) or ref_cache_size < 0:
+            raise ValueError("AuK auk_ref_cache_size must be a non-negative integer (0 disables the cache)")
         self.device = get_local_device()
         self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16
 
@@ -213,6 +227,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         if model_config.get("auk_vae_tile_frames") is not None:
             vae_decode_kwargs["tile_frames"] = int(model_config["auk_vae_tile_frames"])
         self.vae_decode = AuKVAEDecodeGraph(self.vae, enabled=not od_config.enforce_eager, **vae_decode_kwargs)
+        self._ref_cache_size = ref_cache_size
+        self._ref_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
 
         logger.info(
             "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d",
@@ -267,6 +283,12 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         text = torch.zeros(1, _WARMUP_TEXT_TOKENS, self.text_hidden_dim, device=self.device, dtype=self.dtype)
         c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device)
         timestep = torch.zeros((), device=self.device, dtype=torch.float32)
+        # The graphs are keyed by the number of steps, so warm the variant's default grid.
+        if self.is_flash:
+            grid = torch.tensor(self.flash_t_grid, device=self.device, dtype=torch.float32)
+        else:
+            sway = None if self.default_sway is None else float(self.default_sway)
+            grid = build_time_grid(nfe=self.default_nfe, sway_sampling_coef=sway, t_grid=None, device=self.device)
         start = time.perf_counter()
         with torch.inference_mode(), self._dit_autocast():
             for ref_frames in _WARMUP_REF_FRAMES:
@@ -283,6 +305,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
                         timestep=timestep,
                         cfg_strength=cfg,
                         new_request=True,
+                        timesteps=grid[:-1],
+                        step_index=0,
                     )
         if self.device.type != "cpu":
             torch.accelerator.synchronize(self.device)
@@ -386,8 +410,24 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
 
         if audio is None:
             return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32)
-        wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device)
-        return self.vae.encode(wav, sample=sample, generator=generator)
+        wav = self._prepare_waveform(_unwrap_single(audio))
+        # The posterior mean is a pure function of the prepared waveform, so it
+        # is cached by content. A posterior draw depends on the generator and
+        # is always recomputed. Callers must not modify the returned tensor.
+        key = None
+        if not sample and self._ref_cache_size:
+            digest = hashlib.blake2b(wav.numpy().tobytes(), digest_size=16).hexdigest()
+            key = f"{wav.shape[-1]}:{digest}"
+            cached = self._ref_cache.get(key)
+            if cached is not None:
+                self._ref_cache.move_to_end(key)
+                return cached
+        latents = self.vae.encode(wav.to(self.device), sample=sample, generator=generator)
+        if key is not None:
+            self._ref_cache[key] = latents
+            while len(self._ref_cache) > self._ref_cache_size:
+                self._ref_cache.popitem(last=False)
+        return latents
 
     def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int:
         """Target latent length, capped by the remaining context."""

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload_strategy,
@@ -1084,6 +1085,7 @@ class Cosmos3LanguageModel(nn.Module):
         rope_theta: float,
         mrope_section: list[int],
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -1093,21 +1095,21 @@ class Cosmos3LanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3UndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3UndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         # TODO: Not used right now, will be used in the future for prompt upsampler.
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
@@ -1228,6 +1230,10 @@ class Cosmos3VFMTransformer(nn.Module):
 
     _hsdp_shard_conditions = [_is_transformer_block]
 
+    # Standard unquantized vLLM linears need no value-dependent post-load
+    # processing on accelerators. Edge inherits the same loading contract.
+    _hsdp_pre_sharded_meta_post_load = True
+
     # Modules whose parameters must NOT be FSDP-sharded at the root level.
     # time_embedder is cast to fp32 by post_load_weights for precision; if it
     # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16),
@@ -1346,6 +1352,10 @@ class Cosmos3VFMTransformer(nn.Module):
         language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs(
             getattr(od_config, "quantization_config", None)
         )
+        release_completed_blocks_to_meta = bool(
+            getattr(getattr(od_config, "parallel_config", None), "use_hsdp", False)
+            and getattr(od_config, "hsdp_weight_load_strategy", "full") == "pre_sharded"
+        )
         mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config)
         if mixed_precision_config is None:
             if mixed_precision_source == "additional_config_disabled":
@@ -1372,6 +1382,7 @@ class Cosmos3VFMTransformer(nn.Module):
             rope_theta=self.rope_theta,
             mrope_section=self.mrope_section,
             quant_config=language_model_quant_config,
+            release_completed_blocks_to_meta=release_completed_blocks_to_meta,
             prefix="language_model",
             **self._language_model_kwargs(),
         )
@@ -1399,24 +1410,24 @@ class Cosmos3VFMTransformer(nn.Module):
             self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
             self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
 
-        self.gen_layers = nn.ModuleList(
-            [
-                Cosmos3GenDecoderLayer(
-                    layer_idx=i,
-                    hidden_size=self.hidden_size,
-                    intermediate_size=self.intermediate_size,
-                    num_attention_heads=self.num_attention_heads,
-                    num_key_value_heads=self.num_key_value_heads,
-                    head_dim=self.head_dim,
-                    rms_norm_eps=self.rms_norm_eps,
-                    quant_config=gen_layers_quant_config,
-                    mlp_cls=self._gen_mlp_cls,
-                    qk_norm=self.qk_norm_for_diffusion,
-                    prefix=f"gen_layers.{i}",
-                )
-                for i in range(self.num_hidden_layers)
-            ]
-        )
+        self.gen_layers = nn.ModuleList()
+        for i in range(self.num_hidden_layers):
+            layer = Cosmos3GenDecoderLayer(
+                layer_idx=i,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                rms_norm_eps=self.rms_norm_eps,
+                quant_config=gen_layers_quant_config,
+                mlp_cls=self._gen_mlp_cls,
+                qk_norm=self.qk_norm_for_diffusion,
+                prefix=f"gen_layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.gen_layers.append(layer)
 
         self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None
         if mixed_precision_config is not None:

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelL
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 
 from .transformer_cosmos3 import (
     Cosmos3VFMTransformer,
@@ -258,6 +259,7 @@ class Cosmos3EdgeLanguageModel(nn.Module):
         mrope_section: list[int],
         use_und_k_norm_for_gen: bool,
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -267,22 +269,22 @@ class Cosmos3EdgeLanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3EdgeUndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    use_und_k_norm_for_gen=use_und_k_norm_for_gen,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3EdgeUndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                use_und_k_norm_for_gen=use_und_k_norm_for_gen,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -288,13 +288,13 @@ def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
             output_key = output_keys[0]
             frames = frames[output_key]
         output_type = getattr(sampling_params, "output_type", None) or "pt"
-        # The image serving path currently accepts PIL images or NumPy arrays,
-        # while LingBot decodes T2I outputs to tensors. Keep that compatibility
-        # conversion without discarding the model's image/video payload key.
+        # Decoded LingBot pixels are already in [0, 1]. The video API treats
+        # floating tensors as [-1, 1], so use its NumPy path to avoid normalizing
+        # twice. Images also need NumPy for serving; keep latent tensors intact.
         if (
             isinstance(frames, torch.Tensor)
             and output_type != "latent"
-            and (output_type == "np" or output_key == "image")
+            and (output_type == "np" or output_key in {"image", "video"})
         ):
             frames = frames.float().cpu().numpy()
         return {output_key: frames} if output_key is not None else frames

vllm_omni/diffusion/models/utils.py

@@ -168,6 +168,28 @@ def init_parameters(
         init_parameters(child, dtype, device)
 
 
+def release_module_parameters_to_meta(module: nn.Module) -> None:
+    """Release a fully constructed module's parameter storage to meta.
+
+    The module must finish construction first so initialization code and vLLM
+    weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps
+    each registered Parameter object's identity while dropping only its storage.
+    Persistent buffers intentionally remain materialized.
+    """
+    for param in module.parameters():
+        if param.device.type == "meta":
+            continue
+
+        meta_tensor = torch.empty_like(param, device="meta")
+        if type(param) is nn.Parameter:
+            meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad)
+        else:
+            meta_param = param.__class__.__new__(param.__class__, meta_tensor)
+            meta_param.requires_grad_(param.requires_grad)
+        meta_param.__dict__.update(getattr(param, "__dict__", {}))
+        torch.utils.swap_tensors(param, meta_param)
+
+
 def init_parameters_preserving_vllm_linear(
     module: nn.Module,
     dtype: torch.dtype | None,