[Perf][AuK] Graph the encoder prefill, cache reference latents and fuse the codec activation (#8305)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

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()

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)

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

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}

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_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/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."""