diff --git a/recipes/Tencent/AuK-H100.md b/recipes/Tencent/AuK-H100.md index 5acd01cb0..56ee2c18d 100644 --- a/recipes/Tencent/AuK-H100.md +++ b/recipes/Tencent/AuK-H100.md @@ -132,15 +132,23 @@ transcripts, speaker similarity to the upstream output 0.993 to 0.999, log-mel L1 0.11 to 0.47 (the upstream VAE-resampling noise floor is 0.39). Flash: 0.989 to 0.999 and 0.24 to 0.51. Wall per request on the second call: base 1.0 to 2.0 s for 3.5 to 10.9 s of audio, Flash 0.17 to 0.30 s; engine start -about 47 s with a warm page cache. +about 47 s with a warm page cache, plus the codec decode compile on the +diffusion stage (about a minute on a cold Inductor cache, ~30 s warm). ## Notes - Memory usage: about 25 GB peak on the device for a 12 s generation with both stages resident (deploy defaults: encoder 0.45, diffusion stage 0.35 - of device memory). -- Key flags: `enforce_eager` on both stages (the encoder walks the decoder - layers itself for the layer fusion). `enable_prefix_caching` must stay off + of device memory), plus about 1 GB of CUDA graph pools for the compiled + codec decode buckets. +- Key flags: `enforce_eager` on the encoder stage (it walks the decoder + layers itself for the layer fusion); the diffusion stage runs with + `enforce_eager: false` so the codec decode is compiled into bucketed CUDA + graphs at startup (128/256/512 latent frames, i.e. up to 10.24 s; longer + clips are decoded in overlapping 512-frame tiles of the same graph; + override with `model_config.auk_vae_compile_shapes` and + `auk_vae_tile_frames`). The DiT itself stays eager unless + `diffusion_compile_granularity: full`. `enable_prefix_caching` must stay off for the encoder: a cache hit skips prompt positions that the fused condition needs. `enable_chunked_prefill` is off by default: forcing it (128-token chunks, so two to three chunks per prompt) reproduces the diff --git a/tests/diffusion/models/auk/test_auk_vae_fastpath.py b/tests/diffusion/models/auk/test_auk_vae_fastpath.py new file mode 100644 index 000000000..b25381097 --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_vae_fastpath.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AuK codec's decode fast paths keep the reference numerics. + +Five paths are checked against the plain eager decode: the Snake activation +with precomputed exponent caches, the per-channel FIR filter cache, the +per-length CUDA graph tier, the torch.compile bucket tier (both fall back to +eager off CUDA) and the tiled decode of long clips. +""" + +import pytest +import torch + +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample +from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _small_vae() -> AuKVAE: + torch.manual_seed(3) + # Six halvings of the initial width must leave at least two channels. + vae = AuKVAE( + upsample_initial_channel=128, + downsample_channels=(4, 4, 8, 8, 8, 16, 16), + latent_dim=8, + ).eval() + # Give the activations non-trivial parameters so the paths are exercised. + with torch.no_grad(): + for module in vae.modules(): + if isinstance(module, SnakeBeta): + module.alpha.normal_(0.0, 0.3) + module.beta.normal_(0.0, 0.3) + vae.remove_weight_norm() + return vae + + +def _reference_snake(module: SnakeBeta, x: torch.Tensor) -> torch.Tensor: + """The formula as the reference implementation writes it.""" + alpha = module.alpha.unsqueeze(0).unsqueeze(-1) + beta = module.beta.unsqueeze(0).unsqueeze(-1) + if module.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) + + +@torch.inference_mode() +def test_window_padding_decodes_to_raw_zero() -> None: + """Bucket padding is the normalized latent whose denormalized value is zero.""" + vae = _small_vae() + with torch.no_grad(): + vae.global_mean.normal_() + vae.global_log_std.uniform_(0.5, 2.0) + pad = AuKVAEDecodeGraph(vae)._pad_for(torch.device("cpu")) + raw = pad * torch.sqrt(vae.global_log_std) + vae.global_mean + torch.testing.assert_close(raw, torch.zeros_like(raw), atol=1e-6, rtol=0.0) + + +@torch.inference_mode() +def test_eager_snake_matches_the_reference_formula() -> None: + module = SnakeBeta(6, alpha_logscale=True) + with torch.no_grad(): + module.alpha.normal_() + module.beta.normal_() + module.precompute_exp_cache() + x = torch.randn(2, 6, 17) + + assert torch.equal(module(x), _reference_snake(module, x)) + + +@torch.inference_mode() +def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None: + lowpass = LowPass(cutoff=0.25, half_width=0.3, stride=2, kernel_size=12, causal=True) + upsample = Upsample(ratio=2, kernel_size=12) + x = torch.randn(1, 5, 40) + + for module in (lowpass, upsample): + module.cache_filters = False + plain = module(x) + module.cache_filters = True + cached = module(x) + assert torch.equal(cached, plain) + assert module._expanded is not None and module._expanded.shape == (5, 1, 12) + assert module._expanded.is_contiguous() + # A different channel count rebuilds the cache rather than reusing it. + module(torch.randn(1, 3, 40)) + assert module._expanded.shape == (3, 1, 12) + + +@torch.inference_mode() +def test_decode_fast_paths_reproduce_the_plain_decode() -> None: + vae = _small_vae() + latents = torch.randn(1, 6, vae.latent_dim) + + vae.set_decode_fast_paths(cached_filters=False) + plain = vae.decode(latents) + assert plain.shape == (1, 6 * vae.hop_size) + + vae.set_decode_fast_paths(cached_filters=True) + fast = vae.decode(latents) + assert torch.equal(fast, plain) + assert all(module._cached for module in vae.modules() if isinstance(module, SnakeBeta)) + + +@torch.inference_mode() +def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None: + vae = _small_vae() + wrapper = AuKVAEDecodeGraph(vae, frame_alignment=64, compile_shapes=(8,)) + capture_spy = mocker.spy(wrapper, "_capture") + compile_spy = mocker.spy(torch, "compile") + latents = torch.randn(1, 6, vae.latent_dim) + + wrapper.warmup(torch.device("cpu")) + assert torch.equal(wrapper(latents), vae.decode(latents)) + assert wrapper.last_mode == "eager" + capture_spy.assert_not_called() + compile_spy.assert_not_called() + assert not wrapper._cache and not wrapper._compiled + + +def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None: + wrapper = AuKVAEDecodeGraph(_small_vae(), compile_shapes=(16, 8, 8)) + assert wrapper.compile_shapes == [8, 16] + # Nothing captured yet: every length goes to the per-length graph tier. + assert wrapper.compiled_bucket(6) is None + wrapper._compiled = {8: object(), 16: object()} # type: ignore[dict-item] + assert [wrapper.compiled_bucket(frames) for frames in (6, 8, 9, 16, 17)] == [8, 8, 16, 16, None] + # A bucket whose capture failed is skipped, not padded to. + wrapper._compiled = {16: object()} # type: ignore[dict-item] + assert wrapper.compiled_bucket(6) == 16 + + +def test_plan_tiles_covers_the_clip_once_with_full_context() -> None: + assert plan_tiles(40, 64, 10, 3) == [(0, 40, 0, 40)] + windows = plan_tiles(1000, 512, 53, 9) + assert windows == [(0, 512, 0, 503), (450, 512, 503, 953), (488, 512, 953, 1000)] + # Every emitted frame sits at least the context away from a window edge + # that is not the clip's own edge, and the emitted ranges tile the clip. + emitted = 0 + for start, width, emit_start, emit_end in windows: + assert emit_start == emitted and start + width <= 1000 + assert start == 0 or emit_start - start >= 53 + assert start + width == 1000 or start + width - emit_end >= 9 + emitted = emit_end + assert emitted == 1000 + # With the compiled buckets known, the last window shrinks to the smallest + # one that holds the remainder plus the left context: 600 frames cost a + # 512 and a 256 window instead of two 512s. + assert plan_tiles(600, 512, 53, 9, sizes=(128, 256, 512)) == [(0, 512, 0, 503), (344, 256, 503, 600)] + assert plan_tiles(1000, 512, 53, 9, sizes=(128, 256, 512))[-1] == (872, 128, 953, 1000) + with pytest.raises(ValueError, match="must exceed"): + plan_tiles(100, 60, 53, 9) + + +@torch.inference_mode() +def test_decode_context_bounds_the_measured_receptive_field() -> None: + vae = _small_vae() + left, right = vae.decode_context_frames() + # The production geometry: causal stack, conv_pre and the upsamplers look ahead a little. + assert (left, right) == (53, 9) + frames, start, end = 240, 70, 190 + latents = torch.randn(1, frames, vae.latent_dim) + whole = vae.decode(latents)[:, start * vae.hop_size : end * vae.hop_size] + window = vae.decode(latents[:, start:end]) + per_frame = (window - whole).abs().reshape(-1, vae.hop_size).amax(dim=1) + differing = (per_frame > 1e-5).nonzero().flatten().tolist() + # Only frames within the context of a fake edge may differ, and some do: + # the bound is tight enough that the halo is not wasted. + assert differing and all(index < left or end - start - index <= right for index in differing) + + +@torch.inference_mode() +def test_tiled_decode_matches_the_whole_decode() -> None: + vae = _small_vae() + # Off the accelerator every window decodes eagerly, so this isolates the stitching. + wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(64,), tile_frames=64) + assert wrapper.tile_frames == 64 and wrapper.context_frames == (53, 9) + for frames in (64, 65, 200, 331): + latents = torch.randn(1, frames, vae.latent_dim) + whole = vae.decode(latents) + tiled = wrapper(latents) + assert tiled.shape == whole.shape + torch.testing.assert_close(tiled, whole, atol=1e-5, rtol=0.0) + assert wrapper.last_mode == ("eager" if frames == 64 else "tiled") + # The streaming form yields the same audio in order. + latents = torch.randn(1, 200, vae.latent_dim) + pieces = [(start, chunk.clone()) for start, chunk in wrapper.decode_tiles(latents)] + # With a 64-frame tile and a 62-frame halo each interior tile adds two frames. + assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 64, 53, 9)] + assert [start for start, _ in pieces][:3] == [0, 55, 57] + torch.testing.assert_close( + torch.cat([chunk for _, chunk in pieces], dim=1), vae.decode(latents), atol=1e-5, rtol=0.0 + ) + # tile_frames=0 turns tiling off; a tile smaller than the context is refused. + assert AuKVAEDecodeGraph(vae, tile_frames=0).tile_frames == 0 + with pytest.raises(ValueError, match="must exceed"): + AuKVAEDecodeGraph(vae, tile_frames=60) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@torch.inference_mode() +def test_graph_replay_matches_eager_per_length() -> None: + vae = _small_vae().to("cuda") + wrapper = AuKVAEDecodeGraph(vae, max_graphs=2) + pools = [] + for frames in (6, 9, 6, 12, 6): + latents = torch.randn(1, frames, vae.latent_dim, device="cuda") + eager = vae.decode(latents) + replay = wrapper(latents) + # Exact-length graphs replay the very same kernels: bit-identical, + # including after the generation was retired. + assert torch.equal(replay, eager), frames + pools.append(wrapper._plain_pool) + # The third distinct length found the cache full, so the whole generation + # (6 and 9) was retired together with its pool rather than one graph at a + # time; 12 and the re-captured 6 share the new pool. + assert list(wrapper._cache) == [12, 6] + assert pools[0] is pools[2] and pools[3] is pools[4] and pools[2] is not pools[3] + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@torch.inference_mode() +def test_bucketed_graph_only_disturbs_the_tail() -> None: + vae = _small_vae().to("cuda") + wrapper = AuKVAEDecodeGraph(vae, frame_alignment=8) + latents = torch.randn(1, 6, vae.latent_dim, device="cuda") + eager = vae.decode(latents) + replay = wrapper(latents) + assert replay.shape == eager.shape and list(wrapper._cache) == [8] + # The padding decodes to raw zero, which is what the eager conv_pre pads + # with, but it still leaks in through the alias-free upsamplers, whose + # lookahead accumulates through the stack, so a bucketed replay is close + # to but not identical with the eager decode. That is why frame_alignment + # defaults to 1. + assert torch.isfinite(replay).all() + assert not torch.equal(replay, eager) + torch.testing.assert_close(replay, eager, atol=0.1, rtol=0.0) + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA") +@torch.inference_mode() +def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_to_plain_graphs() -> None: + vae = _small_vae().to("cuda") + wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(8,)) + wrapper.warmup(torch.device("cuda")) + assert list(wrapper._compiled) == [8] + + exact = torch.randn(1, 8, vae.latent_dim, device="cuda") + eager = vae.decode(exact) + replay = wrapper(exact) + assert wrapper.last_mode == "compiled" + assert replay.shape == eager.shape + # Same formula, different fusion order: not bit-identical, but close. + torch.testing.assert_close(replay, eager, atol=1e-4, rtol=0.0) + + short = torch.randn(1, 5, vae.latent_dim, device="cuda") + padded = wrapper(short) + assert wrapper.last_mode == "compiled" and padded.shape == (1, 5 * vae.hop_size) + torch.testing.assert_close(padded, vae.decode(short), atol=0.1, rtol=0.0) + + # An 8-frame bucket cannot hold the decoder context, so tiling is off and + # a longer clip gets its own plain graph. + assert wrapper.tile_frames == 0 + long = torch.randn(1, 12, vae.latent_dim, device="cuda") + assert torch.equal(wrapper(long), vae.decode(long)) + assert wrapper.last_mode == "graph" and list(wrapper._cache) == [12] + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA") +@torch.inference_mode() +def test_tiles_replay_the_compiled_bucket_for_long_clips() -> None: + vae = _small_vae().to("cuda") + wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(72,)) + wrapper.warmup(torch.device("cuda")) + assert wrapper.tile_frames == 72 and list(wrapper._compiled) == [72] + + latents = torch.randn(1, 150, vae.latent_dim, device="cuda") + tiled = wrapper(latents) + assert wrapper.last_mode == "tiled" and not wrapper._cache + assert tiled.shape == (1, 150 * vae.hop_size) + torch.testing.assert_close(tiled, vae.decode(latents), atol=1e-4, rtol=0.0) diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py index 261c5240d..954491e9e 100644 --- a/tests/diffusion/models/auk/test_pipeline_auk.py +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -157,6 +157,10 @@ class _StubVAE(nn.Module): self.decode_calls += 1 return torch.zeros(1, latents.shape[1] * HOP, device=latents.device) + def decode_context_frames(self) -> tuple[int, int]: + # A stateless stub: no context, so any tile size is valid. + return 0, 0 + def _stub_sampler(calls: list[dict[str, Any]]): """Record the ODE arguments and return noise drawn from the request generator.""" @@ -399,6 +403,28 @@ class TestRequestParsing: assert calls[0]["seed"] is None assert isinstance(calls[0]["generator"], torch.Generator) + def test_setup_compile_warms_the_decode_buckets_on_the_pipeline_device(self, build_pipeline, mocker): + pipeline, _ = build_pipeline() + warmup = mocker.patch.object(pipeline.vae_decode, "warmup") + compile_dit = mocker.patch.object(pipeline.dit, "compile") + + pipeline.setup_compile() + + warmup.assert_called_once_with(pipeline.device) + # Regional is the default and a no-op for the DiT, which has no repeated blocks. + compile_dit.assert_not_called() + + def test_setup_compile_honours_full_dit_granularity(self, build_pipeline, mocker): + pipeline, _ = build_pipeline() + pipeline.od_config.diffusion_compile_granularity = "full" + pipeline.od_config.diffusion_compile_dynamic = False + mocker.patch.object(pipeline.vae_decode, "warmup") + compile_dit = mocker.patch.object(pipeline.dit, "compile") + + pipeline.setup_compile() + + compile_dit.assert_called_once_with(dynamic=False) + def test_latent_output_type_skips_the_decoder(self, build_pipeline): pipeline, _ = build_pipeline() diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml index bbe544919..43c403bf3 100644 --- a/vllm_omni/deploy/auk.yaml +++ b/vllm_omni/deploy/auk.yaml @@ -40,11 +40,18 @@ stages: # Stage 1: AuKPipeline (DiT + BigVGAN-flow VAE encoder and decoder). # Shares GPU 0 with the encoder; 32 Euler steps at CFG 2.0 is the Base # recipe, and AuK-Flash overrides both with its distilled 4-step grid. + # Not eager: the codec decode is compiled into bucketed CUDA graphs at + # startup (about a minute more on a cold Inductor cache, ~30 s warm, and + # about 1 GB of graph pools); clips longer than the largest bucket are + # decoded in overlapping tiles of it. The DiT stays eager unless + # diffusion_compile_granularity is set to full. Buckets, tile size and the + # plain-graph cache are tunable through model_config: + # model_config: {auk_vae_compile_shapes: [128, 256, 512], auk_vae_tile_frames: 512, auk_vae_max_graphs: 32} - stage_id: 1 devices: "0" max_num_seqs: 1 gpu_memory_utilization: 0.35 - enforce_eager: true + enforce_eager: false default_sampling_params: seed: 0 num_inference_steps: 32 diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py index 1fabeb327..2322cfbc9 100644 --- a/vllm_omni/diffusion/models/auk/auk_vae.py +++ b/vllm_omni/diffusion/models/auk/auk_vae.py @@ -22,6 +22,7 @@ from __future__ import annotations import inspect import math import warnings +from fractions import Fraction from pathlib import Path import torch @@ -32,6 +33,8 @@ 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_omni.model_executor.models.common.snake_activation import SnakeBeta as _SharedSnakeBeta + logger = init_logger(__name__) __all__ = ["AuKVAE"] @@ -195,27 +198,50 @@ class Encoder(nn.Module): return self.generator(x) -class SnakeBeta(nn.Module): - """``x + sin^2(alpha * x) / beta`` with per-channel learned alpha and beta (log-scale in AuK).""" +class SnakeBeta(_SharedSnakeBeta): + """x + sin^2(alpha * x) / beta with per-channel learned alpha and beta (log-scale in AuK). + + Built on the shared speech-decoder activation for its precomputed + exp(alpha) and 1 / (exp(beta) + eps) buffers, which are materialised once + instead of on every call. The forward always uses the eager formula in + the reference's operation order, so it reproduces the reference bit for + bit; the shared class's fused Triton kernel is not used, because inside + the compiled decode buckets Inductor fuses the formula itself and on the + remaining eager paths (the reference-audio encode and clips longer than + the largest bucket) the kernel measured slower than eager. + """ def __init__(self, channels: int, alpha_logscale: bool = False) -> None: - super().__init__() - self.alpha_logscale = alpha_logscale - init = torch.zeros(channels) if alpha_logscale else torch.ones(channels) - self.alpha = nn.Parameter(init.clone()) - self.beta = nn.Parameter(init.clone()) - self.eps = 1e-9 + super().__init__(channels, alpha_logscale=alpha_logscale) def forward(self, x: torch.Tensor) -> torch.Tensor: - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) - beta = self.beta.unsqueeze(0).unsqueeze(-1) - if self.alpha_logscale: - alpha = torch.exp(alpha) - beta = torch.exp(beta) - return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + return self._eager_forward(x) -class LowPass(nn.Module): +class _CachedFilter: + """Mixin for the FIR modules: expand the shared taps per channel once, not per call. + + filter.expand(channels, -1, -1) is a stride-0 view, and the convolution + then copies it to a contiguous weight on every call. Caching the contiguous + copy per channel count removes one launch per activation; the taps are the + same, so the output is unchanged. + """ + + filter: torch.Tensor + cache_filters: bool = True + _expanded: torch.Tensor | None = None + + def _taps(self, channels: int, x: torch.Tensor) -> torch.Tensor: + if not self.cache_filters: + return self.filter.expand(channels, -1, -1) + cached = self._expanded + if cached is None or cached.shape[0] != channels or cached.device != x.device: + cached = self.filter.expand(channels, -1, -1).contiguous() + self._expanded = cached + return cached + + +class LowPass(_CachedFilter, nn.Module): """Fixed FIR low-pass with replicate padding, applied per channel.""" def __init__( @@ -241,10 +267,10 @@ class LowPass(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: channels = x.size(1) x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") - return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) + return F.conv1d(x, self._taps(channels, x), stride=self.stride, groups=channels) -class Upsample(nn.Module): +class Upsample(_CachedFilter, nn.Module): """Band-limited interpolation by ``ratio``. Always non-causal, matching how AuK builds it.""" def __init__(self, ratio: int = 2, kernel_size: int = 12) -> None: @@ -258,7 +284,7 @@ class Upsample(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: channels = x.size(1) x = F.pad(x, (self.pad, self.pad), mode="replicate") - x = self.ratio * F.conv_transpose1d(x, self.filter.expand(channels, -1, -1), stride=self.ratio, groups=channels) + x = self.ratio * F.conv_transpose1d(x, self._taps(channels, x), stride=self.ratio, groups=channels) return x[..., self.pad_left : -self.pad_right] @@ -329,6 +355,51 @@ class AmpBlock(nn.Module): return x +def _conv_context(conv: Conv) -> tuple[Fraction, Fraction]: + """(left, right) input samples a causal or same-padded convolution reaches.""" + if conv.causal: + return Fraction(conv.left_padding), Fraction(0) + reach = Fraction(conv.dilation[0] * (conv.kernel_size[0] - 1), 2) + return reach, reach + + +def _transpose_context(up: ConvTranspose) -> tuple[Fraction, Fraction]: + """(left, right) input frames a transposed convolution reaches, in its input's units.""" + if up.causal: + return Fraction(1), Fraction(0) + reach = Fraction(up.kernel_size[0], 2 * up.stride[0]) + return reach, reach + + +def _activation_context(act: AliasFreeActivation) -> tuple[Fraction, Fraction]: + """(left, right) samples an alias-free activation reaches, in its input's units. + + The upsampler's taps reach half a kernel each side, and the low-pass + behind the activation runs at the oversampled rate, so its padding + counts at ``1 / ratio``. + """ + upsample, lowpass = act.upsample, act.downsample.lowpass + reach = Fraction(upsample.filter.shape[-1], 2 * upsample.ratio) + left = reach + Fraction(lowpass.pad_left, upsample.ratio) + right = reach + Fraction(lowpass.pad_right, upsample.ratio) + return left, right + + +def _block_context(block: AmpBlock) -> tuple[Fraction, Fraction]: + """(left, right) samples one AMP block reaches: its convolutions and activations run in series.""" + left = Fraction(0) + right = Fraction(0) + for conv in list(block.convs1) + list(block.convs2): + conv_left, conv_right = _conv_context(conv) + left += conv_left + right += conv_right + for act in block.activations: + act_left, act_right = _activation_context(act) + left += act_left + right += act_right + return left, right + + class AuKVAE(nn.Module): """AuK's audio codec: :meth:`encode` waveform to normalized latents, :meth:`decode` back. @@ -458,6 +529,26 @@ class AuKVAE(nn.Module): folded += 1 self._norm_folded = True logger.debug("Folded weight norm on %d AuK VAE convolutions", folded) + # The activations' exp(alpha) / 1/(exp(beta)+eps) caches depend only on + # the loaded parameters; materialise them here so no decode call, and + # in particular no CUDA graph capture, computes them lazily. + for module in self.modules(): + if isinstance(module, SnakeBeta): + module.precompute_exp_cache() + + def set_decode_fast_paths(self, *, cached_filters: bool) -> None: + """Toggle the FIR modules' per-channel filter caching. + + On by default. The switch exists so the cache's cost can be measured + against the plain expand-per-call path. Flip it before any CUDA graph + is captured: it reallocates the cached taps, and a captured graph + keeps reading the old buffers. + """ + + for module in self.modules(): + if isinstance(module, _CachedFilter): + module.cache_filters = cached_filters + module._expanded = None def encode( self, wav: torch.Tensor, *, sample: bool = False, generator: torch.Generator | None = None @@ -480,6 +571,36 @@ class AuKVAE(nn.Module): latents = latents.transpose(1, 2).float() return (latents - self.global_mean.float()) / torch.sqrt(self.global_log_std.float()) + def decode_context_frames(self) -> tuple[int, int]: + """Latent frames of (left, right) context a decoded frame can depend on. + + Summed layer by layer from the modules' own padding and kernel + geometry, so it is an upper bound that holds for any weights: a + window decoded on its own reproduces the full decode exactly except + within these many frames of a window edge that is not a true clip + edge. The decoder is causal apart from ``conv_pre`` and the + band-limited upsamplers, so the right context is a few frames while + the left one is dominated by the dilated convolutions of the first + (lowest-rate) stage. + """ + left, right = _conv_context(self.conv_pre) + rate = Fraction(1) # samples per latent frame at the current depth + for i in range(self.num_upsamples): + for up in self.ups[i]: + up_left, up_right = _transpose_context(up) + left += up_left / rate + right += up_right / rate + rate *= up.stride[0] + first = i * self.num_kernels + blocks = [_block_context(self.resblocks[first + j]) for j in range(self.num_kernels)] + left += max(block_left for block_left, _ in blocks) / rate + right += max(block_right for _, block_right in blocks) / rate + act_left, act_right = _activation_context(self.activation_post) + post_left, post_right = _conv_context(self.conv_post) + left += (act_left + post_left) / rate + right += (act_right + post_right) / rate + return math.ceil(left), math.ceil(right) + def decode(self, latents: torch.Tensor) -> torch.Tensor: """Decode normalized ``[B, Np, latent_dim]`` latents into a ``[B, Np * hop_size]`` waveform in [-1, 1].""" if latents.dim() != 3 or latents.size(-1) != self.latent_dim: diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py index 6ad12b02d..b4b379894 100644 --- a/vllm_omni/diffusion/models/auk/pipeline_auk.py +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -27,6 +27,7 @@ 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_vae import AuKVAE +from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph from vllm_omni.diffusion.models.interface import ( SupportAudioInput, SupportAudioOutput, @@ -182,6 +183,19 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True) self.dit = self.dit.to(device=self.device).eval() self.dit.requires_grad_(False) + # The compiled decode buckets are warmed by setup_compile(), which the + # model runner calls at startup unless the stage is enforce_eager. The + # bucket list and the plain-graph cache size come from the stage's + # model_config, like the other codec graph wrappers' knobs. + model_config = getattr(od_config, "model_config", None) or {} + vae_decode_kwargs: dict[str, Any] = {} + if model_config.get("auk_vae_compile_shapes") is not None: + vae_decode_kwargs["compile_shapes"] = [int(size) for size in model_config["auk_vae_compile_shapes"]] + if model_config.get("auk_vae_max_graphs") is not None: + vae_decode_kwargs["max_graphs"] = int(model_config["auk_vae_max_graphs"]) + 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) logger.info( "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d", @@ -192,6 +206,27 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp self.sample_rate, ) + def setup_compile(self) -> None: + """Compile the DiT as configured, then capture the codec decode buckets. + + Defining this hook replaces the runner's generic transformer compile, + so the DiT's ``diffusion_compile_granularity`` is honoured here: full + compiles the whole transformer lazily, exactly as the runner would; + regional is a no-op because the DiT declares no repeated blocks. The + startup cost worth paying is the VAE decode, whose buckets are + compiled and captured before the first request. + """ + granularity = self.od_config.diffusion_compile_granularity + if granularity == "full": + self.dit.compile(dynamic=self.od_config.diffusion_compile_dynamic) + logger.info( + "AuK DiT configured for lazy full torch.compile with dynamic=%s", + self.od_config.diffusion_compile_dynamic, + ) + else: + logger.info("AuK DiT declares no repeated blocks; regional torch.compile is a no-op, running eager") + self.vae_decode.warmup(self.device) + # The assembled checkpoint is not a diffusers layout: __init__ reads # auk.safetensors and vae.safetensors directly, so the loader has no # component sources to stream and load_weights only reports what is loaded. @@ -435,7 +470,7 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp if output_type == "latent": return [DiffusionOutput(output=latents.detach().cpu())] - wav = self.vae.decode(latents) + wav = self.vae_decode(latents) # One mono waveform per request; the formatter expects [T]. wav = wav.detach().to(device="cpu", dtype=torch.float32).reshape(-1) diff --git a/vllm_omni/diffusion/models/auk/vae_cudagraph.py b/vllm_omni/diffusion/models/auk/vae_cudagraph.py new file mode 100644 index 000000000..1c679bfb4 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/vae_cudagraph.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graph replay for the AuK codec decode. + +The decoder launches about a thousand small kernels per clip, so it is +bound by launch overhead rather than by the GPU. This wrapper replays the +whole decode as CUDA graphs, in three tiers: + +Compiled bucket graphs are built at startup by warmup(): decode is passed +through torch.compile so Inductor fuses the elementwise chains, then one +graph is captured per bucket in compile_shapes. Shorter clips are +right-padded to their bucket with the normalized latent that decodes to +raw zero, which is what the eager decode's own convolution padding sees, +so the padding only reaches the last frames through the upsamplers' +lookahead. The fused kernels differ from eager in rounding order, so these +graphs are close to but not bit-identical with the eager decode. + +Clips longer than tile_frames are decoded in tiles of that many frames +(by default the largest bucket, so every tile replays the same compiled +graph). Neighbouring tiles overlap by the decoder's receptive field, taken +from AuKVAE.decode_context_frames(), and each tile only contributes the +frames that lie outside that halo, so the stitched waveform equals the +whole-clip decode up to floating-point order. Memory and startup cost stay +constant however long the clip; the price is the halo recomputed per tile +(about an eighth of a 512-frame tile). decode_tiles() yields the tiles as +they finish so a caller can stream them. + +Plain graphs are captured on demand, one per exact latent length, for +windows no compiled bucket serves (tiling off, or a bucket whose compile +failed). They replay the same kernels as eager and are bit-identical with +it. They share one wrapper-private graph pool, so they are never evicted +one at a time: when max_graphs is reached the whole generation is retired +and a fresh pool is started, and no surviving graph can hold an address +that a destroyed graph released. +""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger +from vllm.utils.math_utils import round_up + +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, SnakeBeta +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +# Latent-frame buckets for the compiled graphs, at 50 Hz: 2.56, 5.12 and +# 10.24 s. Each bucket costs one Inductor compilation (tens of seconds) at +# startup and one private CUDA graph pool; longer clips are tiled over the +# largest one. Deployments override them with +# ``model_config.auk_vae_compile_shapes`` on the diffusion stage. +DEFAULT_COMPILE_SHAPES: tuple[int, ...] = (128, 256, 512) + + +def plan_tiles( + frames: int, tile: int, left: int, right: int, sizes: Sequence[int] = () +) -> list[tuple[int, int, int, int]]: + """Split a clip into decode windows: ``(window_start, window_frames, emit_start, emit_end)``. + + Windows are ``tile`` frames wide. A window contributes the frames that + are at least ``left`` frames after its start and ``right`` frames before + its end, unless that edge is the clip's own edge, so each emitted frame + has its full decoder context inside its window. The last window is + pinned to the clip's end; it shrinks to the smallest of ``sizes`` (the + compiled buckets) that still holds the remaining frames plus the left + context, so a clip just past one tile does not pay for a second full + one. A clip that fits in one tile is a single window of its own length. + """ + if frames <= tile: + return [(0, frames, 0, frames)] + if tile <= left + right: + raise ValueError(f"tile_frames={tile} must exceed the decoder context {left} + {right}") + windows: list[tuple[int, int, int, int]] = [] + emit = 0 + while True: + start = 0 if emit == 0 else emit - left + if start + tile >= frames: + width = tile + for size in sorted(sizes): + if left + frames - emit <= size <= tile: + width = size + break + windows.append((frames - width, width, emit, frames)) + return windows + end = start + tile - right + windows.append((start, tile, emit, end)) + emit = end + + +@dataclass +class _GraphEntry: + graph: torch.cuda.CUDAGraph + static_latents: torch.Tensor + static_wav: torch.Tensor + + +class AuKVAEDecodeGraph: + """Replay AuKVAE.decode for one clip per call.""" + + def __init__( + self, + vae: AuKVAE, + *, + enabled: bool = True, + frame_alignment: int = 1, + max_graphs: int = 32, + compile_shapes: Sequence[int] = DEFAULT_COMPILE_SHAPES, + tile_frames: int | None = None, + ) -> None: + """``tile_frames``: window for clips longer than it; None = the largest bucket, 0 = no tiling.""" + self.vae = vae + self.enabled = bool(enabled) + self.frame_alignment = max(1, int(frame_alignment)) + self.max_graphs = max(1, int(max_graphs)) + self.compile_shapes = sorted({int(size) for size in compile_shapes if int(size) > 0}) + self.context_frames = vae.decode_context_frames() + if tile_frames is None: + tile_frames = self.compile_shapes[-1] if self.compile_shapes else 0 + if tile_frames <= sum(self.context_frames): + logger.warning( + "AuK codec decode: largest bucket %d does not exceed the decoder context %s; not tiling", + tile_frames, + self.context_frames, + ) + tile_frames = 0 + self.tile_frames = max(0, int(tile_frames)) + if self.tile_frames and self.tile_frames <= sum(self.context_frames): + raise ValueError(f"tile_frames={self.tile_frames} must exceed the decoder context {self.context_frames}") + # Plain graphs keyed by latent length. They all live in _plain_pool and + # are retired together (see _retire_plain_graphs). + self._cache: OrderedDict[int, _GraphEntry] = OrderedDict() + self._plain_pool = None + # Normalized latent that decodes to raw zero, per device; pads short windows. + self._pad_latent: torch.Tensor | None = None + # Compiled graphs keyed by bucket, filled by warmup(). + self._compiled: dict[int, _GraphEntry] = {} + self._compiled_decode: Callable[[torch.Tensor], torch.Tensor] | None = None + # Which path served the last call: "compiled", "graph" or "eager". + self.last_mode: str | None = None + + @torch.no_grad() + def __call__(self, latents: torch.Tensor) -> torch.Tensor: + """Decode [1, frames, latent_dim] latents into a [1, frames * hop] waveform.""" + + # Graphs are captured for one clip at a time. Anything else runs the + # plain eager decode, whole. + single_clip = latents.ndim == 3 and latents.shape[0] == 1 + if not self.enabled or not single_clip: + self.last_mode = "eager" + return self.vae.decode(latents) + + frames = int(latents.shape[1]) + if not self.tile_frames or frames <= self.tile_frames: + return self._decode_window(latents).clone() + + hop = self.vae.hop_size + wav = latents.new_empty((1, frames * hop)) + for emit_start, chunk in self.decode_tiles(latents): + wav[:, emit_start * hop : emit_start * hop + chunk.shape[1]].copy_(chunk) + self.last_mode = "tiled" + return wav + + def decode_tiles(self, latents: torch.Tensor): + """Yield ``(emit_start_frame, waveform)`` per tile, in order, for a [1, frames, latent_dim] clip. + + The waveforms are views into graph buffers where a graph served the + tile; copy before the next tile if they must outlive the iteration. + """ + frames = int(latents.shape[1]) + tile = self.tile_frames or frames + left, right = self.context_frames + hop = self.vae.hop_size + for start, width, emit_start, emit_end in plan_tiles(frames, tile, left, right, self._compiled): + wav = self._decode_window(latents[:, start : start + width]) + yield emit_start, wav[:, (emit_start - start) * hop : (emit_end - start) * hop] + + def _decode_window(self, latents: torch.Tensor) -> torch.Tensor: + """Decode one window through the best available tier; graph results are views into static buffers.""" + + # A call made while an outer graph is being captured, or off the + # accelerator, runs eager. The capture query is asked last because it + # needs a CUDA runtime. + if not latents.is_cuda or torch.cuda.is_current_stream_capturing(): + self.last_mode = "eager" + return self.vae.decode(latents) + + frames = int(latents.shape[1]) + bucket = self.compiled_bucket(frames) + if bucket is not None: + entry = self._compiled[bucket] + self.last_mode = "compiled" + else: + bucket = round_up(frames, self.frame_alignment) + entry = self._cache.get(bucket) + if entry is None: + if len(self._cache) >= self.max_graphs: + self._retire_plain_graphs() + if self._plain_pool is None: + self._plain_pool = torch.cuda.graph_pool_handle() + entry = self._capture(bucket, latents.device, self.vae.decode, warm_iters=2, pool=self._plain_pool) + self._cache[bucket] = entry + self.last_mode = "graph" + + if bucket == frames: + entry.static_latents.copy_(latents) + else: + entry.static_latents[:, :frames].copy_(latents) + entry.static_latents[:, frames:].copy_(self._pad_for(latents.device).expand(1, bucket - frames, -1)) + entry.graph.replay() + return entry.static_wav[:, : frames * self.vae.hop_size] + + def _retire_plain_graphs(self) -> None: + """Drop every plain graph and their shared pool; the next capture starts a new one.""" + logger.info("AuK codec decode: retiring %d plain CUDA graphs", len(self._cache)) + self._cache.clear() + self._plain_pool = None + + def _pad_for(self, device: torch.device) -> torch.Tensor: + """The [latent_dim] normalized latent that AuKVAE.decode maps to raw zero.""" + pad = self._pad_latent + if pad is None or pad.device != device: + mean = self.vae.global_mean.detach().float().to(device) + scale = torch.sqrt(self.vae.global_log_std.detach().float().to(device)) + pad = -mean / scale + self._pad_latent = pad + return pad + + def compiled_bucket(self, frames: int) -> int | None: + """The smallest compiled bucket that holds frames latents, or None past the largest.""" + for size in self.compile_shapes: + if frames <= size and size in self._compiled: + return size + return None + + def warmup(self, device: torch.device | str) -> None: + """Compile the decode and capture every bucket in compile_shapes. + + Meant for service startup, since each bucket costs one Inductor + compilation. A failure only logs a warning and the bucket is served + by a plain graph instead. + """ + device = torch.device(device) + if not self.enabled or not self.compile_shapes or self._compiled: + return + on_accelerator = current_omni_platform.is_cuda_alike() and device.type == current_omni_platform.device_type + if not on_accelerator or torch.cuda.is_current_stream_capturing(): + return + + # The traced graph must read exp(alpha) from a buffer, not recompute it. + for module in self.vae.modules(): + if isinstance(module, SnakeBeta): + module.precompute_exp_cache() + + try: + self._compiled_decode = torch.compile(self.vae.decode, mode="default", fullgraph=False, dynamic=False) + except Exception: + logger.warning("torch.compile of the AuK codec decode failed; using plain CUDA graphs", exc_info=True) + self._compiled_decode = None + return + + # Inductor fuses the plain Snake formula with the ops around it. + for size in self.compile_shapes: + try: + self._compiled[size] = self._capture(size, device, self._compiled_decode, warm_iters=5) + logger.info("Compiled and captured AuK codec decode: latent_frames=%d", size) + except Exception: + logger.warning( + "Compiled AuK codec decode failed for latent_frames=%d; falling back to plain CUDA graphs", + size, + exc_info=True, + ) + logger.info( + "AuK codec decode compile warmup done: %d/%d buckets", len(self._compiled), len(self.compile_shapes) + ) + + def _capture( + self, + bucket: int, + device: torch.device, + decode: Callable[[torch.Tensor], torch.Tensor], + *, + warm_iters: int, + pool=None, + ) -> _GraphEntry: + """Capture one graph of decode on zero latents of bucket frames. + + The warm iterations let cuDNN pick its algorithms and, for the + compiled decode, Inductor finish tracing and autotuning before + anything is recorded. + """ + static_latents = torch.zeros(1, bucket, self.vae.latent_dim, device=device, dtype=torch.float32) + with torch.inference_mode(): + self.vae.decode(static_latents) + for _ in range(warm_iters): + decode(static_latents) + torch.accelerator.synchronize(device) + graph = torch.cuda.CUDAGraph() + if pool is None: + pool = current_omni_platform.get_global_graph_pool() + with torch.cuda.graph(graph, pool=pool): + static_wav = decode(static_latents) + logger.info("Captured AuK codec decode CUDA graph: latent_frames=%d", bucket) + return _GraphEntry(graph=graph, static_latents=static_latents, static_wav=static_wav) + + +__all__ = ["DEFAULT_COMPILE_SHAPES", "AuKVAEDecodeGraph", "plan_tiles"]