From 4025bb30f670c83649a4dce6779227d01f2b58ec Mon Sep 17 00:00:00 2001 From: Yueqian Lin Date: Sat, 3 Oct 2026 01:43:55 +1000 Subject: [PATCH] [Perf][AuK] Graph the encoder prefill, cache reference latents and fuse the codec activation (#8305) Signed-off-by: Yueqian Lin --- .../models/auk/test_auk_vae_fastpath.py | 69 ++++++- .../diffusion/models/auk/test_pipeline_auk.py | 46 +++++ .../test_auk_encoder_cudagraph.py | 175 ++++++++++++++++++ vllm_omni/deploy/auk.yaml | 14 +- .../diffusion/models/auk/alias_free_triton.py | 148 +++++++++++++++ vllm_omni/diffusion/models/auk/auk_vae.py | 39 ++++ .../diffusion/models/auk/pipeline_auk.py | 31 +++- 7 files changed, 515 insertions(+), 7 deletions(-) create mode 100644 tests/e2e/offline_inference/test_auk_encoder_cudagraph.py create mode 100644 vllm_omni/diffusion/models/auk/alias_free_triton.py diff --git a/tests/diffusion/models/auk/test_auk_vae_fastpath.py b/tests/diffusion/models/auk/test_auk_vae_fastpath.py index 8a1b39aa1..bec028bc7 100644 --- a/tests/diffusion/models/auk/test_auk_vae_fastpath.py +++ b/tests/diffusion/models/auk/test_auk_vae_fastpath.py @@ -12,7 +12,8 @@ import pytest import torch from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample +from vllm_omni.diffusion.models.auk import auk_vae +from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles from vllm_omni.platforms import current_omni_platform @@ -246,6 +247,72 @@ def test_tiled_decode_matches_the_whole_decode() -> None: AuKVAEDecodeGraph(vae, tile_frames=60) +@pytest.mark.cpu +@torch.inference_mode() +def test_fused_activation_is_skipped_off_cuda() -> None: + activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)) + x = torch.randn(1, 4, 16) + assert not activation._fused_eligible(x) + assert activation(x).shape == (1, 4, 16) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA") +@torch.inference_mode() +def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None: + activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda") + x = torch.randn(1, 4, 16, device="cuda") + assert activation._fused_eligible(x) + monkeypatch.setattr(auk_vae, "HAS_TRITON", False) + assert not activation._fused_eligible(x) + assert activation(x).shape == (1, 4, 16) + + +def _activation_input(layout: str, length: int) -> torch.Tensor: + """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer.""" + if layout == "contiguous": + x = torch.randn(2, 6, length, device="cuda") + elif layout == "transposed": + x = torch.randn(2, length, 6, device="cuda").transpose(1, 2) + elif layout == "sliced": + x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2] + else: + raise ValueError(layout) + assert x.shape == (2, 6, length) + if length > 1: # contiguity of a single-sample view depends on the layout + assert x.is_contiguous() == (layout == "contiguous") + return x * 3 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA") +@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"]) +@pytest.mark.parametrize("causal", [False, True]) +# The kernels tile the output in blocks of 1024 samples: the 2x oversampled +# signal crosses a block boundary at 512 input samples and the decimated one at +# 1024, so both sides of each boundary are covered, plus short and odd lengths. +@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025]) +@torch.inference_mode() +def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None: + torch.manual_seed(0) + activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda") + with torch.no_grad(): + activation.act.alpha.normal_(0.0, 0.3) + activation.act.beta.normal_(0.0, 0.3) + x = _activation_input(layout, length) + original = x.clone() + + eager = activation.downsample(activation.act(activation.upsample(x))) + assert activation._fused_eligible(x) + fused = activation(x) + + # Same fp32 arithmetic, only the order in which the taps are summed differs. + assert fused.shape == eager.shape + torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5) + # Neither path writes into its (possibly strided) input. + torch.testing.assert_close(x, original, atol=0, rtol=0) + + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") @torch.inference_mode() diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py index 67d38e09a..01cc035f3 100644 --- a/tests/diffusion/models/auk/test_pipeline_auk.py +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -291,6 +291,18 @@ class TestRequestParsing: with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"): build_pipeline(model_config={"max_dit_graphs": max_dit_graphs}) + @pytest.mark.parametrize( + "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)] + ) + def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected): + pipeline, _ = build_pipeline(model_config=model_config) + assert pipeline._ref_cache_size == expected + + @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None]) + def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size): + with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"): + build_pipeline(model_config={"auk_ref_cache_size": size}) + def test_pipeline_declares_audio_output(self, build_pipeline): pipeline, _ = build_pipeline() @@ -492,6 +504,40 @@ class TestRequestParsing: assert call["sample"] is True assert isinstance(call["generator"], torch.Generator) + def test_reference_latents_are_cached_by_content(self, build_pipeline): + pipeline, calls = build_pipeline() + clip = _silence(1.0) + + for _ in range(3): + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + # A different clip is a new entry, the first one still hits. + other = (np.full_like(clip[0], 0.25), clip[1]) + pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1)) + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + + assert len(pipeline.vae.encode_calls) == 2 + assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"] + + def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline): + pipeline, _ = build_pipeline() + for _ in range(2): + pipeline.forward( + _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1) + ) + assert len(pipeline.vae.encode_calls) == 2 + + pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0}) + for _ in range(2): + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1)) + assert len(pipeline.vae.encode_calls) == 2 + + def test_reference_cache_is_bounded(self, build_pipeline): + pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2}) + for level in (0.1, 0.2, 0.3): + clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE) + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + assert len(pipeline._ref_cache) == 2 + def test_one_request_per_forward(self, build_pipeline): pipeline, _ = build_pipeline() batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1) diff --git a/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py b/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py new file mode 100644 index 000000000..72c1a73dd --- /dev/null +++ b/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request. + +``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few +token buckets. A replay reuses static input buffers and the graph pool, so a +request could in principle read conditioning left behind by an earlier, longer +prompt. This test runs the same request sequence through an eager encoder and +through the graphed one and compares the text condition stage 0 hands to +stage 1 (the fused thinker hidden states) for every request. The sequence +alternates prompts of different lengths inside one capture bucket, crosses +into a larger bucket, and then returns to the first prompt. + +Like ``test_auk.py`` it needs an assembled checkpoint directory: + + VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py +""" + +from __future__ import annotations + +import bisect +import os +from pathlib import Path +from typing import Any + +import pytest +import torch +import yaml + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor +from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params + +MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR" + +# Text-only instruct prompts: the token count, and so the capture bucket, +# follows the text length. SHORT_A and SHORT_B differ in length but share a +# bucket; LONG lands in a larger one. The test checks these properties on the +# recorded condition lengths, so a tokenizer change fails loudly instead of +# silently shrinking the coverage. +_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".' +SHORT_A = _TEMPLATE.format( + desc="A calm young woman speaking warmly.", + text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.", +) +SHORT_B = _TEMPLATE.format( + desc="An older man with a deep, slightly raspy voice, speaking slowly.", + text=( + "The train to the coast leaves at seven, so please be on time. Bring a warm coat, " + "some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain." + ), +) +LONG = _TEMPLATE.format( + desc=( + "A cheerful radio host in her thirties, bright and energetic, with clear articulation, " + "a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone" + ), + text=( + "Good morning and welcome to the weekend edition. Today we are talking about small gardens, " + "the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into " + "a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. " + "Stay with us, because after the news we will open the phone lines and hear from you." + ), +) +SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A] + +# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also +# runs eager in both configs, see _deploy). +GEN_SECONDS = 1.0 +NFE = 1 + +# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs +# the same kernels on the real tokens, so the two paths agree to bf16 rounding. +# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the +# condition's max magnitude for every request; the bounds leave about 3x +# headroom for other GPUs. +MIN_COSINE = 0.9999 +MAX_REL_ABS = 1e-2 + +_model_dir = os.environ.get(MODEL_DIR_ENV) +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"), +] + + +def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str: + """``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped. + + Only the encoder differs between the two runs. The conditions compared here + are recorded before stage 1 runs, so its DiT compile, graph capture and + codec bucket warm-up would only add startup time; eager stage 1 skips them + and still takes the real stage handoff. + """ + with open(get_deploy_config_path("auk.yaml")) as f: + deploy = yaml.safe_load(f) + for stage in deploy["stages"]: + if stage["stage_id"] == 1: + stage["enforce_eager"] = True + elif not graph_encoder: + stage["enforce_eager"] = True + stage.pop("compilation_config", None) + path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml" + path.write_text(yaml.safe_dump(deploy)) + return str(path) + + +def _graph_capture_sizes() -> list[int]: + with open(get_deploy_config_path("auk.yaml")) as f: + deploy = yaml.safe_load(f) + stage = next(s for s in deploy["stages"] if s["stage_id"] == 0) + assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test" + return sorted(stage["compilation_config"]["cudagraph_capture_sizes"]) + + +def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]: + """Run SEQUENCE one request at a time and return each request's text condition.""" + conditions: list[torch.Tensor] = [] + original = auk_input_processor.encoder2dit + + def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None: + result = original(*args, **kwargs) + if result is not None: + conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone()) + return result + + # The orchestrator runs in this process and resolves the stage input + # function by module attribute at startup, so the patch must precede Omni. + with monkeypatch.context() as patch: + patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit) + with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner: + for text in SEQUENCE: + prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS) + outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests" + return conditions + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + capture_sizes = _graph_capture_sizes() + eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch) + graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch) + + lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)} + for i, (ref, got) in enumerate(zip(eager, graphed)): + assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}" + assert torch.isfinite(got).all(), f"request {i}: non-finite condition" + cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item() + rel_abs = ((got - ref).abs().max() / ref.abs().max()).item() + assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, ( + f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} " + "of the condition's magnitude" + ) + + # The sequence must exercise what it claims to: two different lengths in + # one bucket, a second bucket, and every prompt inside the graphed range. + assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size" + bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()} + assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket) + assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket) + + # Repeats of a prompt reproduce its first condition exactly, eager or graphed: + # nothing from the requests in between leaks into the replay. + for run in (eager, graphed): + first: dict[str, torch.Tensor] = {} + for text, cond in zip(SEQUENCE, run): + if text in first: + torch.testing.assert_close(cond, first[text], atol=0, rtol=0) + else: + first[text] = cond diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml index 989e64965..555342027 100644 --- a/vllm_omni/deploy/auk.yaml +++ b/vllm_omni/deploy/auk.yaml @@ -20,14 +20,20 @@ enable_chunked_prefill: false stages: # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder. - # Prefill only — one sampled token, discarded downstream. eager because the - # fused-hidden forward walks the decoder layers itself and so bypasses the - # compiled graph. + # Prefill only: one sampled token, discarded downstream. The fused-hidden + # forward walks the decoder layers itself, so torch.compile has nothing to + # hook (mode 0); the whole prefill is instead captured as FULL CUDA graphs + # padded to the sizes below, which removes most of its launch overhead + # (about 20 ms per request on H200). Longer prompts run eager. - stage_id: 0 devices: "0" max_num_seqs: 1 gpu_memory_utilization: 0.45 - enforce_eager: true + enforce_eager: false + compilation_config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [64, 128, 192, 256, 384, 512] max_num_batched_tokens: 32768 # AuK takes text plus at most one clip; no image or video tower is built. limit_mm_per_prompt: {audio: 1} diff --git a/vllm_omni/diffusion/models/auk/alias_free_triton.py b/vllm_omni/diffusion/models/auk/alias_free_triton.py new file mode 100644 index 000000000..2d37e859e --- /dev/null +++ b/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") diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py index 2322cfbc9..4ad20dd87 100644 --- a/vllm_omni/diffusion/models/auk/auk_vae.py +++ b/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.""" diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py index da5d8d1e5..21b639411 100644 --- a/vllm_omni/diffusion/models/auk/pipeline_auk.py +++ b/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 @@ -64,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 @@ -157,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 @@ -218,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", @@ -399,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."""