[Perf][AuK] Compile the DiT step, use cuDNN attention and warm up at startup (#8300)

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

recipes/Tencent/AuK-H100.md

@@ -12,7 +12,7 @@
   both with `Qwen/Qwen2.5-Omni-3B` as the frozen encoder
 - Task: zero-shot TTS, instruct TTS, content and acoustic editing,
   paralinguistic editing, speech enhancement, speaker and music separation
-- Mode: offline inference (online chat serving is not yet qualified)
+- Mode: offline inference and online `/v1/audio/speech` (online chat serving is not yet qualified)
 - Hardware: 1x H100 80GB
 - Maintainer: Community
 
@@ -130,10 +130,16 @@ Measured on one H100 against the upstream implementation with the same seeds
 and mean VAE latents (four cookbook cases, base checkpoint): identical
 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, plus the codec decode compile on the
-diffusion stage (about a minute on a cold Inductor cache, ~30 s warm).
+0.989 to 0.999 and 0.24 to 0.51.
+
+Latency with the DiT and codec CUDA graphs, measured end to end on one
+H200-class GPU with both stages on it, concurrency 1, median of 10 warm
+requests: base (32 steps, CFG 2) zero-shot 0.25 / 0.31 / 0.40 s for 3 / 6 /
+12 s of audio, instruct TTS 0.17 / 0.23 / 0.32 s; Flash zero-shot 0.08 /
+0.10 / 0.14 s. The first request after start is 1.2 to 1.6 s, because the
+common shapes are compiled at startup. Engine start adds the codec decode
+compile and the DiT warmup to model loading (a few minutes on a cold
+Inductor cache, about 30 s less warm).
 
 ## Notes
 
@@ -144,11 +150,15 @@ diffusion stage (about a minute on a cold Inductor cache, ~30 s warm).
 - 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;
+  graphs at startup (the deploy config sets 160/320/640 latent frames, i.e.
+  up to 12.8 s; longer clips are decoded in overlapping 640-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
+  `auk_vae_tile_frames`). Each DiT denoise step replays a per-shape CUDA
+  graph of the regionally compiled double- and single-stream blocks; the
+  per-request conditioning (text projection, reference embedding, padding
+  biases, rotary tables) is prepared once per request outside the graph.
+  `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
@@ -169,7 +179,7 @@ diffusion stage (about a minute on a cold Inductor cache, ~30 s warm).
 | --- | --- | --- |
 | Offline `Omni.generate` | supported (base and Flash) | `docs/getting_started/quickstart.md` |
 | Online `/v1/chat/completions` with audio | not yet qualified | `docs/serving/` |
-| `/v1/audio/speech` | not supported (needs a TTS adapter) | `docs/contributing/model/adding_tts_model.md` |
+| `/v1/audio/speech` | supported (base and Flash, `instructions` or `task_type`) | `docs/user_guide/examples/online_serving/text_to_speech.md` |
 | Streaming / async chunk | not supported | `docs/design/feature/async_chunk.md` |
 | Batching across requests | one request per DiT forward | `docs/user_guide/diffusion/` |
 | Tensor / sequence parallelism | not supported | `docs/configuration/composable_parallel.md` |

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -29,13 +29,14 @@ def _make_dit(device: str) -> AuKTransformer:
     )
 
 
-def _sample_inputs(device: str, offset: float = 0.0) -> dict[str, torch.Tensor]:
+def _sample_inputs(device: str, offset: float = 0.0, ref_frames: int = 4) -> dict[str, torch.Tensor]:
+    """``ref_frames=0`` is the text-only (instruct TTS) request shape."""
     generator = torch.Generator(device=device).manual_seed(42)
     return {
         "text": torch.randn(1, 7, 8, generator=generator, device=device) + offset,
         "c_mask": torch.ones(1, 7, dtype=torch.bool, device=device),
-        "ref": torch.randn(1, 4, 4, generator=generator, device=device) - offset,
-        "ref_mask": torch.ones(1, 4, dtype=torch.bool, device=device),
+        "ref": (torch.randn(1, 4, 4, generator=generator, device=device) - offset)[:, :ref_frames],
+        "ref_mask": torch.ones(1, ref_frames, dtype=torch.bool, device=device),
     }
 
 
@@ -44,8 +45,8 @@ def _sample_inputs(device: str, offset: float = 0.0) -> dict[str, torch.Tensor]:
 def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: float, mocker) -> None:
     dit = _make_dit("cpu")
     wrapper = AuKCUDAGraphWrapper(dit)
-    run_spy = mocker.spy(wrapper, "_run")
-    run_cfg_spy = mocker.spy(wrapper, "_run_cfg")
+    prepare_spy = mocker.spy(wrapper, "_prepare")
+    step_spy = mocker.spy(wrapper, "_step")
     capture_spy = mocker.spy(wrapper, "_capture")
     inputs = _sample_inputs("cpu")
     common = dict(
@@ -58,13 +59,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa
     graph = sample_latents(dit, **common, generator=torch.Generator().manual_seed(7), sampler=wrapper)
     torch.testing.assert_close(graph, eager)
 
-    if cfg_strength >= 1e-5:
-        run_spy.assert_not_called()
-        assert run_cfg_spy.call_count == 2
-    else:
-        assert run_spy.call_count == 2
-        run_cfg_spy.assert_not_called()
-
+    # The conditioning is prepared once per request and reused by every step.
+    assert prepare_spy.call_count == 1
+    assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5)
+    assert step_spy.call_count == 2
     capture_spy.assert_not_called()
     assert not wrapper._cache
 
@@ -79,14 +77,14 @@ def test_graph_inputs_use_bounded_length_buckets() -> None:
 
     padded = wrapper._bucket_inputs(x, text, c_mask, ref, ref_mask)
 
-    assert padded[0].shape == (1, 128, 4)
-    assert padded[1].shape == (1, 128)
-    assert padded[2].shape == (1, 128, 8)
-    assert padded[3].shape == (1, 128)
+    assert padded[0].shape == (1, 96, 4)
+    assert padded[1].shape == (1, 96)
+    assert padded[2].shape == (1, 96, 8)
+    assert padded[3].shape == (1, 96)
     assert padded[4].shape == (1, 100, 4)
     assert padded[5].shape == (1, 100)
     assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51]
-    assert wrapper._key(padded[0], padded[2], padded[4], False) == (128, 128, 100, False)
+    assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False)
     assert wrapper.max_graphs == 32
 
 
@@ -105,43 +103,24 @@ def test_full_graph_cache_retires_as_one_generation() -> None:
 
 
 @torch.inference_mode()
+@pytest.mark.parametrize("ref_frames", [0, 4])
 @pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
-def test_bucket_padding_preserves_real_frame_outputs(cfg_strength: float) -> None:
+def test_bucket_padding_preserves_real_frame_outputs(cfg_strength: float, ref_frames: int) -> None:
     dit = _make_dit("cpu")
     wrapper = AuKCUDAGraphWrapper(dit)
-    inputs = _sample_inputs("cpu")
+    inputs = _sample_inputs("cpu", ref_frames=ref_frames)
     x = torch.randn(1, 9, 4)
     timestep = torch.tensor(0.4)
     cfg = torch.tensor(cfg_strength)
-    if cfg_strength >= 1e-5:
-        eager = wrapper._run_cfg(
-            x,
-            None,
-            inputs["text"],
-            inputs["c_mask"],
-            inputs["ref"],
-            inputs["ref_mask"],
-            timestep,
-            cfg_strength=cfg,
-        )
-    else:
-        eager = wrapper._run(
-            x,
-            None,
-            inputs["text"],
-            inputs["c_mask"],
-            inputs["ref"],
-            inputs["ref_mask"],
-            timestep,
-        )
-    dit.clear_cache()
+    uses_cfg = cfg_strength >= 1e-5
+    text, c_mask, ref, ref_mask = inputs["text"], inputs["c_mask"], inputs["ref"], inputs["ref_mask"]
 
-    bucketed = wrapper._bucket_inputs(x, inputs["text"], inputs["c_mask"], inputs["ref"], inputs["ref_mask"])
-    if cfg_strength >= 1e-5:
-        padded = wrapper._run_cfg(*bucketed, timestep, cfg_strength=cfg)
-    else:
-        padded = wrapper._run(*bucketed, timestep)
-    dit.clear_cache()
+    ctx = wrapper._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg)
+    eager = wrapper._step(x, timestep, ctx, cfg)
+
+    bucketed = wrapper._bucket_inputs(x, text, c_mask, ref, ref_mask)
+    padded_ctx = wrapper._prepare(*bucketed, uses_cfg)
+    padded = wrapper._step(bucketed[0], timestep, padded_ctx, cfg)
 
     torch.testing.assert_close(padded[:, : x.shape[1]], eager)
 
@@ -149,13 +128,14 @@ def test_bucket_padding_preserves_real_frame_outputs(cfg_strength: float) -> Non
 @pytest.mark.cuda
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
 @torch.inference_mode()
+@pytest.mark.parametrize("ref_frames", [0, 4])
 @pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
-def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_strength: float) -> None:
+def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_strength: float, ref_frames: int) -> None:
     dit = _make_dit("cuda")
     wrapper = AuKCUDAGraphWrapper(dit)
 
     for offset in (0.0, 0.25):
-        inputs = _sample_inputs("cuda", offset)
+        inputs = _sample_inputs("cuda", offset, ref_frames=ref_frames)
         common = dict(
             **inputs,
             gen_frames=9,
@@ -176,12 +156,14 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng
         key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
         assert key in wrapper._cache
 
+        # The static context was refreshed for this request's conditioning.
         entry = wrapper._cache[key]
-        torch.testing.assert_close(entry.static_x_mask, bucketed[1])
-        torch.testing.assert_close(entry.static_text, bucketed[2])
-        torch.testing.assert_close(entry.static_c_mask, bucketed[3])
-        torch.testing.assert_close(entry.static_ref, bucketed[4])
-        torch.testing.assert_close(entry.static_ref_mask, bucketed[5])
+        fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5)
+        for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True):
+            if want is None:
+                assert static is None
+            else:
+                torch.testing.assert_close(static, want)
         torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
 
     assert len(wrapper._cache) == 1
@@ -194,7 +176,7 @@ def test_graph_capture_failure_is_propagated(mocker) -> None:
     dit = _make_dit("cuda")
     wrapper = AuKCUDAGraphWrapper(dit)
     capture = mocker.patch.object(wrapper, "_capture", side_effect=RuntimeError("capture failed"))
-    eager = mocker.spy(wrapper, "_run")
+    eager = mocker.spy(wrapper, "_step")
     inputs = _sample_inputs("cuda")
     common = dict(
         **inputs,

tests/diffusion/models/auk/test_auk_step_context.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The AuK DiT's per-request prepare() / per-step step() split.
+
+The Euler loop prepares the conditioning once and then only steps; these
+tests pin that the split, the additive key-padding bias and the shared CFG
+target embedding reproduce the one-shot forward exactly.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    _key_padding_bias,
+    _sdpa,
+    build_time_grid,
+    sample_latents,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(attn_mask_enabled: bool = True) -> AuKTransformer:
+    torch.manual_seed(3)
+    return AuKTransformer(
+        dim=32,
+        heads=2,
+        dim_head=16,
+        ff_mult=2,
+        latent_dim=4,
+        text_hidden_dim=8,
+        num_layers=2,
+        num_single_layers=2,
+        attn_mask_enabled=attn_mask_enabled,
+    ).eval()
+
+
+def _inputs(ref_frames: int, pad: bool) -> dict[str, torch.Tensor | None]:
+    generator = torch.Generator().manual_seed(5)
+    text = torch.randn(1, 7, 8, generator=generator)
+    c_mask = torch.ones(1, 7, dtype=torch.bool)
+    mask = None
+    ref = torch.randn(1, ref_frames, 4, generator=generator)
+    ref_mask = torch.ones(1, ref_frames, dtype=torch.bool)
+    if pad:
+        # Padded tails in every stream, as the CUDA graph buckets produce.
+        c_mask[:, 5:] = False
+        text[:, 5:] = 0.0
+        mask = torch.ones(1, 9, dtype=torch.bool)
+        mask[:, 7:] = False
+        if ref_frames:
+            ref_mask[:, ref_frames - 1 :] = False
+    return {"text": text, "c_mask": c_mask, "mask": mask, "ref": ref, "ref_mask": ref_mask}
+
+
+@torch.inference_mode()
+def test_key_padding_bias_matches_boolean_mask() -> None:
+    generator = torch.Generator().manual_seed(0)
+    q, k, v = (torch.randn(2, 3, 6, 8, generator=generator) for _ in range(3))
+    keep = torch.tensor([[True] * 6, [True, True, True, True, False, False]])
+    expanded = keep[:, None, None, :].expand(2, 3, 6, 6)
+    want = F.scaled_dot_product_attention(q, k, v, attn_mask=expanded)
+    torch.testing.assert_close(_sdpa(q, k, v, _key_padding_bias(keep, q.dtype)), want)
+    assert _key_padding_bias(None, q.dtype) is None
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("attn_mask_enabled", [True, False])
+@pytest.mark.parametrize("cfg_infer", [False, True])
+@pytest.mark.parametrize("ref_frames", [0, 4])
+@pytest.mark.parametrize("pad", [False, True])
+def test_prepare_then_step_matches_forward(attn_mask_enabled: bool, cfg_infer: bool, ref_frames: int, pad: bool):
+    dit = _make_dit(attn_mask_enabled)
+    inputs = _inputs(ref_frames, pad)
+    x = torch.randn(1, 9, 4, generator=torch.Generator().manual_seed(9))
+    time = torch.tensor(0.3)
+
+    want = dit(x, time=time, cfg_infer=cfg_infer, **inputs)
+    ctx = dit.prepare(inputs.pop("text"), target_len=x.shape[1], cfg_infer=cfg_infer, **inputs)
+    got = dit.step(x, time, ctx)
+
+    torch.testing.assert_close(got, want, rtol=0, atol=0)
+    assert got.shape == (2 if cfg_infer else 1, 9, 4)
+    assert ctx.branches == (2 if cfg_infer else 1)
+    if not attn_mask_enabled:
+        assert ctx.joint_bias is None and ctx.single_bias is None
+
+
+@torch.inference_mode()
+def test_cfg_branches_match_separate_cond_and_uncond_forwards() -> None:
+    """The shared target embedding keeps each CFG branch equal to its own forward."""
+    dit = _make_dit()
+    inputs = _inputs(4, pad=True)
+    x = torch.randn(1, 9, 4, generator=torch.Generator().manual_seed(9))
+    time = torch.tensor(0.6)
+
+    both = dit(x, time=time, cfg_infer=True, **inputs)
+    cond = dit(x, time=time, **inputs)
+    uncond = dit(x, time=time, drop_audio_cond=True, drop_text=True, **inputs)
+
+    torch.testing.assert_close(both[:1], cond)
+    torch.testing.assert_close(both[1:], uncond)
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> None:
+    """sample_latents prepares once; stepping the full forward every time agrees."""
+    dit = _make_dit()
+    inputs = _inputs(4, pad=False)
+    inputs.pop("mask")
+    grid = [0.0, 0.3, 0.7, 1.0]
+
+    got = sample_latents(
+        dit, **inputs, gen_frames=9, t_grid=grid, cfg_strength=cfg_strength, generator=torch.Generator().manual_seed(1)
+    )
+
+    x = torch.randn(9, 4, generator=torch.Generator().manual_seed(1)).unsqueeze(0)
+    timesteps = build_time_grid(nfe=len(grid) - 1, sway_sampling_coef=None, t_grid=grid, device="cpu")
+    for i in range(len(grid) - 1):
+        if cfg_strength >= 1e-5:
+            cond, uncond = dit(x, time=timesteps[i], cfg_infer=True, **inputs).chunk(2, dim=0)
+            v = cond + (cond - uncond) * cfg_strength
+        else:
+            v = dit(x, time=timesteps[i], **inputs)
+        x = x + (timesteps[i + 1] - timesteps[i]) * v
+
+    torch.testing.assert_close(got, x, rtol=0, atol=0)
+    assert dit.text_cond is None and dit.text_uncond is None
+
+
+@torch.inference_mode()
+def test_context_copy_rejects_structural_mismatch() -> None:
+    dit = _make_dit()
+    inputs = _inputs(4, pad=True)
+    text = inputs.pop("text")
+    cfg = dit.prepare(text, target_len=9, cfg_infer=True, **inputs)
+    plain = dit.prepare(text, target_len=9, **inputs)
+    with pytest.raises(ValueError, match="branches"):
+        cfg.copy_(plain)
+
+    clone = cfg.clone()
+    other = dit.prepare(text * 2.0, target_len=9, cfg_infer=True, **inputs)
+    clone.copy_(other)
+    torch.testing.assert_close(clone.c, other.c)
+    assert clone.c.data_ptr() != other.c.data_ptr()

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -424,24 +424,56 @@ class TestRequestParsing:
     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")
+        regional = mocker.patch.object(pipeline_auk, "regionally_compile")
+        mocker.patch.object(pipeline, "_warmup_dit")
 
         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()
+        # Regional is the default: the DiT's repeated blocks are compiled.
+        regional.assert_called_once_with(pipeline.dit, dynamic=pipeline.od_config.diffusion_compile_dynamic)
 
     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")
+        mocker.patch.object(pipeline, "_warmup_dit")
+        step = mocker.Mock()
+        pipeline.dit.step = step
+        compile_fn = mocker.patch.object(pipeline_auk.torch, "compile", side_effect=lambda fn, **_: fn)
 
         pipeline.setup_compile()
 
-        compile_dit.assert_called_once_with(dynamic=False)
+        # The samplers call step() directly, so that is what gets compiled.
+        compile_fn.assert_called_once_with(step, dynamic=False)
+
+    @pytest.mark.parametrize("variant, cfg", [("base", 2.0), ("flash", 0.0)])
+    def test_setup_compile_warms_the_dit_graph_buckets(self, build_pipeline, mocker, variant, cfg):
+        pipeline, _ = build_pipeline(variant)
+        mocker.patch.object(pipeline.vae_decode, "warmup")
+        mocker.patch.object(pipeline_auk, "regionally_compile")
+        wrapper = mocker.Mock(enabled=True)
+        pipeline.cudagraph_wrapper = wrapper
+
+        pipeline.setup_compile()
+
+        shapes = [(call.kwargs["x"].shape[1], call.kwargs["ref"].shape[1]) for call in wrapper.call_args_list]
+        assert shapes == [(150, 0), (300, 0), (600, 0), (150, 150), (300, 150), (600, 150)]
+        for call in wrapper.call_args_list:
+            assert call.kwargs["cfg_strength"] == cfg
+            assert call.kwargs["new_request"] is True
+            assert call.kwargs["text"].shape == (1, 96, TEXT_HIDDEN_DIM)
+
+    @pytest.mark.parametrize("model_config, enabled", [({"auk_dit_warmup_frames": []}, True), ({}, False)])
+    def test_dit_warmup_can_be_skipped(self, build_pipeline, mocker, model_config, enabled):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        wrapper = mocker.Mock(enabled=enabled)
+        pipeline.cudagraph_wrapper = wrapper
+
+        pipeline._warmup_dit()
+
+        wrapper.assert_not_called()
 
     def test_latent_output_type_skips_the_decoder(self, build_pipeline):
         pipeline, _ = build_pipeline()

vllm_omni/deploy/auk.yaml

@@ -43,16 +43,22 @@ stages:
   # 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. Each DiT denoise step replays a lazily
-  # captured per-shape CUDA graph. Buckets, tile size and the
+  # decoded in overlapping tiles of it. Each DiT denoise step replays a
+  # per-shape CUDA graph of the regionally compiled blocks; the common 3/6/12 s
+  # shapes are compiled and captured at startup (auk_dit_warmup_frames), other
+  # shapes on first use. 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}
+  #   model_config: {auk_vae_compile_shapes: [...], auk_vae_tile_frames: <frames>, auk_vae_max_graphs: 32}
   - stage_id: 1
     devices: "0"
     max_num_seqs: 1
     gpu_memory_utilization: 0.35
     enforce_eager: false
     model_config:
+      # Codec decode buckets in latent frames (3.2 / 6.4 / 12.8 s). Tighter
+      # than the 128/256/512 defaults for the common 3-12 s requests: a
+      # 3 s clip decodes in the 160 bucket instead of padding to 256.
+      auk_vae_compile_shapes: [160, 320, 640]
       # Maximum number of cached DiT graphs. When the cache reaches this limit,
       # capturing a new shape retires all existing graphs before capture.
       max_dit_graphs: 32

vllm_omni/diffusion/models/auk/auk_transformer.py

@@ -29,25 +29,50 @@ Deliberate differences from the reference, all inference-only:
 
 import math
 from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass
 
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
-__all__ = ["AuKTransformer", "build_time_grid", "dit_state_dict", "sample_latents"]
+__all__ = ["AuKStepContext", "AuKTransformer", "build_time_grid", "dit_state_dict", "sample_latents"]
 
 
-def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor:
     """The single attention call site, so the backend can be swapped in one place.
 
-    ``q``/``k``/``v`` are ``[B, H, N, D]``. ``mask`` is a boolean key-padding
-    mask ``[B, K]`` in which ``True`` marks a valid key; it is broadcast to
-    ``[B, H, Q, K]`` only when given.
+    ``q``/``k``/``v`` are ``[B, H, N, D]``. ``bias`` is an additive key-padding
+    bias ``[B, 1, 1, K]`` (see :func:`_key_padding_bias`) or ``None``. Passing
+    it unexpanded lets SDPA broadcast it inside a fused kernel (cuDNN on
+    Hopper) instead of materializing a ``[B, H, Q, K]`` mask per layer.
     """
-    attn_mask = None
-    if mask is not None:
-        attn_mask = mask[:, None, None, :].expand(q.shape[0], q.shape[1], q.shape[-2], k.shape[-2])
-    return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+    if bias is None or not q.is_cuda:
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0, is_causal=False)
+    # With a bias, SDPA's default choice here is the memory-efficient kernel,
+    # about 3x slower than cuDNN's at these lengths on Hopper.
+    with sdpa_kernel(_MASKED_SDPA_PRIORITY, set_priority=True):
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=bias, dropout_p=0.0, is_causal=False)
+
+
+_MASKED_SDPA_PRIORITY = [
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+
+def _key_padding_bias(mask: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor | None:
+    """Additive attention bias ``[B, 1, 1, K]`` from a boolean key mask ``[B, K]``.
+
+    Valid keys add 0 and padded keys add ``-inf``, which is exactly what SDPA
+    does with a boolean mask internally.
+    """
+    if mask is None:
+        return None
+    bias = torch.zeros(mask.shape, dtype=dtype, device=mask.device).masked_fill(~mask, float("-inf"))
+    return bias[:, None, None, :]
 
 
 def _rotate_half(x: torch.Tensor) -> torch.Tensor:
@@ -57,14 +82,26 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     return torch.stack((-x2, x1), dim=-1).flatten(-2)
 
 
-def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
-    """Apply rotary frequencies ``[1, 1, N, D]`` to ``x`` ``[B, H, N, D]`` in fp32."""
+def _apply_rope_cos_sin(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Rotate ``x`` ``[B, H, N, D]`` in fp32 with precomputed ``cos``/``sin`` of the frequencies."""
     out_dtype = x.dtype
     x32 = x.float()
-    out = x32 * freqs.cos() + _rotate_half(x32) * freqs.sin()
+    out = x32 * cos + _rotate_half(x32) * sin
     return out.to(out_dtype)
 
 
+def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """Apply rotary frequencies ``[1, 1, N, D]`` to ``x`` ``[B, H, N, D]`` in fp32."""
+    return _apply_rope_cos_sin(x, freqs.cos(), freqs.sin())
+
+
+RopeCache = tuple[torch.Tensor, torch.Tensor]
+
+
+def _rope_cos_sin(freqs: torch.Tensor) -> RopeCache:
+    return freqs.cos(), freqs.sin()
+
+
 class Rotary(nn.Module):
     """Rotary position frequencies for one stream, evaluated in fp32.
 
@@ -230,32 +267,50 @@ class Attention(nn.Module):
         packed: torch.Tensor,
         q_norm: nn.Module,
         k_norm: nn.Module,
-        rope: torch.Tensor | None,
+        rope: RopeCache | None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Split a packed QKV projection into heads, RMSNorm Q/K, then rotate."""
+        """Split a packed QKV projection into heads, RMSNorm Q/K, then rotate.
+
+        ``rope`` is the ``(cos, sin)`` pair of the rotary frequencies.
+        """
         q, k, v = (t.view(t.shape[0], -1, self.heads, self.dim_head).transpose(1, 2) for t in packed.chunk(3, dim=-1))
         q = q_norm(q)
         k = k_norm(k)
         if rope is not None:
-            q = _apply_rope(q, rope)
-            k = _apply_rope(k, rope)
+            q = _apply_rope_cos_sin(q, *rope)
+            k = _apply_rope_cos_sin(k, *rope)
         return q, k, v
 
-    def _project(self, x: torch.Tensor, rope: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def _project(self, x: torch.Tensor, rope: RopeCache | None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         return self._qkv(self.to_qkv(x), self.q_norm, self.k_norm, rope)
 
     def forward(
-        self, x: torch.Tensor, mask: torch.Tensor | None = None, rope: torch.Tensor | None = None
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        rope: RopeCache | None = None,
+        bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        """``bias`` is the precomputed key-padding bias for ``mask``; built here when absent."""
         q, k, v = self._project(x, rope)
-        attn_mask = mask if self.attn_mask_enabled else None
-        out = self._merge_heads(_sdpa(q, k, v, attn_mask).to(q.dtype))
+        if bias is None and self.attn_mask_enabled:
+            bias = _key_padding_bias(mask, q.dtype)
+        out = self._merge_heads(_sdpa(q, k, v, bias).to(q.dtype))
         out = self.to_out[0](out)
         if mask is not None:
             out = out.masked_fill(~mask.unsqueeze(-1), 0.0)
         return out
 
 
+def joint_key_mask(mask: torch.Tensor | None, c_mask: torch.Tensor | None, text_len: int) -> torch.Tensor | None:
+    """Key mask of joint attention over ``[audio | text]``; ``None`` when the audio has no mask."""
+    if mask is None:
+        return None
+    if c_mask is not None:
+        return torch.cat([mask, c_mask], dim=1)
+    return F.pad(mask, (0, text_len), value=True)
+
+
 class JointAttention(Attention):
     """Attention over the audio and text streams jointly, with separate projections."""
 
@@ -267,7 +322,7 @@ class JointAttention(Attention):
         self.c_k_norm = nn.RMSNorm(dim_head, elementwise_affine=True)
         self.to_out_c = nn.Linear(inner_dim, context_dim)
 
-    def _project_context(self, c: torch.Tensor, c_rope: torch.Tensor | None) -> tuple[torch.Tensor, ...]:
+    def _project_context(self, c: torch.Tensor, c_rope: RopeCache | None) -> tuple[torch.Tensor, ...]:
         return self._qkv(self.to_qkv_c(c), self.c_q_norm, self.c_k_norm, c_rope)
 
     def forward(  # type: ignore[override]  # joint attention takes the context stream too
@@ -275,10 +330,12 @@ class JointAttention(Attention):
         x: torch.Tensor,
         c: torch.Tensor,
         mask: torch.Tensor | None = None,
-        rope: torch.Tensor | None = None,
-        c_rope: torch.Tensor | None = None,
+        rope: RopeCache | None = None,
+        c_rope: RopeCache | None = None,
         c_mask: torch.Tensor | None = None,
+        bias: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        """``bias`` is the precomputed :func:`joint_key_mask` bias; built here when absent."""
         audio_len = x.shape[1]
         q, k, v = self._project(x, rope)
         c_q, c_k, c_v = self._project_context(c, c_rope)
@@ -287,14 +344,10 @@ class JointAttention(Attention):
         k = torch.cat([k, c_k], dim=2)
         v = torch.cat([v, c_v], dim=2)
 
-        joint_mask = None
-        if self.attn_mask_enabled and mask is not None:
-            if c_mask is not None:
-                joint_mask = torch.cat([mask, c_mask], dim=1)
-            else:
-                joint_mask = F.pad(mask, (0, c.shape[1]), value=True)
+        if bias is None and self.attn_mask_enabled:
+            bias = _key_padding_bias(joint_key_mask(mask, c_mask, c.shape[1]), q.dtype)
 
-        out = self._merge_heads(_sdpa(q, k, v, joint_mask).to(q.dtype))
+        out = self._merge_heads(_sdpa(q, k, v, bias).to(q.dtype))
         x_out = self.to_out[0](out[:, :audio_len])
         c_out = self.to_out_c(out[:, audio_len:])
 
@@ -326,14 +379,15 @@ class DoubleBlock(nn.Module):
         c: torch.Tensor,
         t: torch.Tensor,
         mask: torch.Tensor | None,
-        rope: torch.Tensor,
-        c_rope: torch.Tensor,
+        rope: RopeCache,
+        c_rope: RopeCache,
         c_mask: torch.Tensor | None,
+        bias: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t)
         norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t)
 
-        x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask)
+        x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask, bias=bias)
 
         c = c + c_gate_msa[:, None] * c_attn
         norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
@@ -355,9 +409,16 @@ class SingleBlock(nn.Module):
         self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
         self.ff = FeedForward(dim=dim, mult=ff_mult)
 
-    def forward(self, x: torch.Tensor, t: torch.Tensor, mask: torch.Tensor | None, rope: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        t: torch.Tensor,
+        mask: torch.Tensor | None,
+        rope: RopeCache,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
         norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t)
-        x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope)
+        x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope, bias=bias)
         norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
         return x + gate_mlp[:, None] * self.ff(norm)
 
@@ -365,6 +426,9 @@ class SingleBlock(nn.Module):
 class AuKTransformer(nn.Module):
     """Flow-matching velocity predictor for AuK audio latents.
 
+    The forward pass is :meth:`prepare` (per request) followed by
+    :meth:`step` (per timestep); an Euler loop calls the two separately.
+
     Args:
         dim: Model width.
         heads: Attention heads.
@@ -379,6 +443,9 @@ class AuKTransformer(nn.Module):
             which is what the reference does with the flag off.
     """
 
+    # Regional torch.compile targets (see vllm_omni.diffusion.compile).
+    _repeated_blocks = ["DoubleBlock", "SingleBlock"]
+
     def __init__(
         self,
         *,
@@ -395,6 +462,7 @@ class AuKTransformer(nn.Module):
         super().__init__()
         self.dim = dim
         self.latent_dim = latent_dim
+        self.attn_mask_enabled = attn_mask_enabled
         self.text_cond: torch.Tensor | None = None
         self.text_uncond: torch.Tensor | None = None
 
@@ -433,40 +501,134 @@ class AuKTransformer(nn.Module):
         self.text_cond = None
         self.text_uncond = None
 
-    def _embed_audio(
-        self,
-        x: torch.Tensor,
-        ref: torch.Tensor | None,
-        drop_audio_cond: bool,
-        mask: torch.Tensor | None,
-        ref_mask: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, int]:
-        """Embed the target latents, prepending the reference prompt when present.
+    def _embed_prompt(self, ref: torch.Tensor, ref_mask: torch.Tensor | None, drop: bool) -> torch.Tensor:
+        """Embed the reference prompt, zeroed for the audio-dropped CFG branch."""
+        if drop:
+            ref = torch.zeros_like(ref)
+        return self.audio_embed(ref, mask=ref_mask)
 
-        Returns the ``[ref | target]`` sequence, its padding mask and the
-        reference length, which the caller uses to slice the target back out.
+    def prepare(
+        self,
+        text: torch.Tensor,
+        *,
+        target_len: int,
+        mask: torch.Tensor | None = None,
+        c_mask: torch.Tensor | None = None,
+        ref: torch.Tensor | None = None,
+        ref_mask: torch.Tensor | None = None,
+        drop_audio_cond: bool = False,
+        drop_text: bool = False,
+        cfg_infer: bool = False,
+        cache: bool = False,
+    ) -> "AuKStepContext":
+        """Everything a denoise step needs that does not depend on ``x`` or the time.
+
+        An Euler loop over fixed conditioning calls this once and then
+        :meth:`step` per timestep, so the text projection, the reference
+        prompt embedding, the padding masks and biases and the rotary tables
+        are not recomputed on every step. Arguments are those of
+        :meth:`forward`; ``target_len`` is the target frame count ``n``.
         """
+        batch = text.shape[0]
+        if c_mask is None:
+            # Padding rows are all-zero in the text encoder's output.
+            c_mask = text.abs().sum(-1) > 0
         if ref is not None and ref.shape[1] == 0:
             ref = None
 
-        target = self.audio_embed(x, mask=mask)
-        if ref is None:
-            return target, mask, 0
+        if cfg_infer:
+            if cache and self.text_cond is not None:
+                c_cond, c_uncond = self.text_cond, self.text_uncond
+            else:
+                c_cond = self.project_text(text)
+                # The uncond branch drops the text, which is exactly a zeroed projection.
+                c_uncond = torch.zeros_like(c_cond)
+                if cache:
+                    self.text_cond, self.text_uncond = c_cond, c_uncond
+            c = torch.cat((c_cond, c_uncond), dim=0)
+            prompt = None
+            if ref is not None:
+                prompt = torch.cat(
+                    (self._embed_prompt(ref, ref_mask, False), self._embed_prompt(ref, ref_mask, True)), dim=0
+                )
+        else:
+            c = self.project_text(text)
+            if drop_text:
+                c = torch.zeros_like(c)
+            prompt = None if ref is None else self._embed_prompt(ref, ref_mask, drop_audio_cond)
 
-        if drop_audio_cond:
-            ref = torch.zeros_like(ref)
-        prompt = self.audio_embed(ref, mask=ref_mask)
-        audio = torch.cat([prompt, target], dim=1)
-        prompt_len = prompt.shape[1]
-
-        audio_mask = None
-        if mask is not None or ref_mask is not None:
+        audio_mask = mask
+        if ref is not None and (mask is not None or ref_mask is not None):
             if mask is None:
-                mask = x.new_ones(x.shape[:2], dtype=torch.bool)
+                mask = torch.ones((batch, target_len), dtype=torch.bool, device=ref.device)
             if ref_mask is None:
-                ref_mask = ref.new_ones(ref.shape[:2], dtype=torch.bool)
+                ref_mask = torch.ones(ref.shape[:2], dtype=torch.bool, device=ref.device)
             audio_mask = torch.cat([ref_mask, mask], dim=1)
-        return audio, audio_mask, prompt_len
+
+        branches = 2 if cfg_infer else 1
+        if branches == 2:
+            c_mask = torch.cat((c_mask, c_mask), dim=0)
+            if audio_mask is not None:
+                audio_mask = torch.cat((audio_mask, audio_mask), dim=0)
+
+        prompt_len = 0 if prompt is None else prompt.shape[1]
+        text_len = c.shape[1]
+        audio_len = prompt_len + target_len
+        single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1)
+        joint_bias = single_bias = None
+        if self.attn_mask_enabled:
+            joint_bias = _key_padding_bias(joint_key_mask(audio_mask, c_mask, text_len), c.dtype)
+            single_bias = _key_padding_bias(single_mask, c.dtype)
+
+        return AuKStepContext(
+            c=c,
+            prompt=prompt,
+            target_mask=mask,
+            c_mask=c_mask,
+            audio_mask=audio_mask,
+            single_mask=single_mask,
+            joint_bias=joint_bias,
+            single_bias=single_bias,
+            rope_audio=_rope_cos_sin(self.rotary_embed(audio_len, audio_mask)),
+            rope_text=_rope_cos_sin(self.rotary_embed(text_len, c_mask)),
+            rope_single=_rope_cos_sin(self.rotary_embed(text_len + audio_len, single_mask)),
+            branches=branches,
+        )
+
+    def step(self, x: torch.Tensor, time: torch.Tensor, ctx: "AuKStepContext") -> torch.Tensor:
+        """Predict the velocity of ``x`` at ``time`` under a :meth:`prepare` context.
+
+        The target embedding does not depend on the CFG branch, so it is
+        computed once and shared by both branches.
+        """
+        if time.ndim == 0:
+            time = time.repeat(x.shape[0])
+        t = self.time_embed(time)
+        target = self.audio_embed(x, mask=ctx.target_mask)
+        if ctx.branches == 2:
+            t = torch.cat((t, t), dim=0)
+            target = torch.cat((target, target), dim=0)
+        audio = target if ctx.prompt is None else torch.cat([ctx.prompt, target], dim=1)
+        c = ctx.c
+
+        for block in self.transformer_blocks:
+            c, audio = block(
+                audio,
+                c,
+                t,
+                mask=ctx.audio_mask,
+                rope=ctx.rope_audio,
+                c_rope=ctx.rope_text,
+                c_mask=ctx.c_mask,
+                bias=ctx.joint_bias,
+            )
+
+        h = torch.cat([c, audio], dim=1)
+        for block in self.single_transformer_blocks:
+            h = block(h, t, mask=ctx.single_mask, rope=ctx.rope_single, bias=ctx.single_bias)
+
+        h = h[:, h.shape[1] - x.shape[1] :]
+        return self.proj_out(self.norm_out(h, t))
 
     def forward(
         self,
@@ -508,57 +670,89 @@ class AuKTransformer(nn.Module):
             Velocity ``[B, n, latent_dim]``, or ``[2B, n, latent_dim]`` under
             ``cfg_infer``.
         """
-        batch = x.shape[0]
-        if time.ndim == 0:
-            time = time.repeat(batch)
-        t = self.time_embed(time)
+        ctx = self.prepare(
+            text,
+            target_len=x.shape[1],
+            mask=mask,
+            c_mask=c_mask,
+            ref=ref,
+            ref_mask=ref_mask,
+            drop_audio_cond=drop_audio_cond,
+            drop_text=drop_text,
+            cfg_infer=cfg_infer,
+            cache=cache and cfg_infer,
+        )
+        return self.step(x, time, ctx)
 
-        if c_mask is None:
-            # Padding rows are all-zero in the text encoder's output.
-            c_mask = text.abs().sum(-1) > 0
 
-        if cfg_infer:
-            if cache and self.text_cond is not None:
-                c_cond, c_uncond = self.text_cond, self.text_uncond
-            else:
-                c_cond = self.project_text(text)
-                # The uncond branch drops the text, which is exactly a zeroed projection.
-                c_uncond = torch.zeros_like(c_cond)
-                if cache:
-                    self.text_cond, self.text_uncond = c_cond, c_uncond
-            x_cond, mask_cond, prompt_len = self._embed_audio(x, ref, False, mask, ref_mask)
-            x_uncond, mask_uncond, _ = self._embed_audio(x, ref, True, mask, ref_mask)
+@dataclass
+class AuKStepContext:
+    """Per-request state of the denoise step, built by :meth:`AuKTransformer.prepare`.
 
-            x = torch.cat((x_cond, x_uncond), dim=0)
-            c = torch.cat((c_cond, c_uncond), dim=0)
-            t = torch.cat((t, t), dim=0)
-            c_mask = torch.cat((c_mask, c_mask), dim=0)
-            audio_mask = None
-            if mask_cond is not None and mask_uncond is not None:
-                audio_mask = torch.cat((mask_cond, mask_uncond), dim=0)
-        else:
-            c = self.project_text(text)
-            if drop_text:
-                c = torch.zeros_like(c)
-            x, audio_mask, prompt_len = self._embed_audio(x, ref, drop_audio_cond, mask, ref_mask)
+    Batch rows are ``branches`` copies of the request (cond first under CFG).
+    ``target_mask`` covers the target frames only and has the request's batch
+    size; the other masks cover their full sequences and every branch.
+    """
 
-        seq_len = x.shape[1]
-        text_len = c.shape[1]
-        rope_audio = self.rotary_embed(seq_len, audio_mask)
-        rope_text = self.rotary_embed(text_len, c_mask)
+    c: torch.Tensor
+    prompt: torch.Tensor | None
+    target_mask: torch.Tensor | None
+    c_mask: torch.Tensor | None
+    audio_mask: torch.Tensor | None
+    single_mask: torch.Tensor | None
+    joint_bias: torch.Tensor | None
+    single_bias: torch.Tensor | None
+    rope_audio: RopeCache
+    rope_text: RopeCache
+    rope_single: RopeCache
+    branches: int
 
-        for block in self.transformer_blocks:
-            c, x = block(x, c, t, mask=audio_mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask)
+    def tensors(self) -> list[torch.Tensor | None]:
+        """Every tensor field in a fixed order, rope pairs flattened."""
+        return [
+            self.c,
+            self.prompt,
+            self.target_mask,
+            self.c_mask,
+            self.audio_mask,
+            self.single_mask,
+            self.joint_bias,
+            self.single_bias,
+            *self.rope_audio,
+            *self.rope_text,
+            *self.rope_single,
+        ]
 
-        x = torch.cat([c, x], dim=1)
-        single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1)
-        rope = self.rotary_embed(text_len + seq_len, single_mask)
+    def copy_(self, other: "AuKStepContext") -> None:
+        """Overwrite these tensors in place with ``other``'s, for CUDA graph static inputs."""
+        if other.branches != self.branches:
+            raise ValueError(f"AuK step context mismatch: {other.branches} vs {self.branches} branches")
+        for dst, src in zip(self.tensors(), other.tensors(), strict=True):
+            if (dst is None) != (src is None):
+                raise ValueError("AuK step context mismatch: a mask is present in only one context")
+            if dst is not None:
+                dst.copy_(src)
 
-        for block in self.single_transformer_blocks:
-            x = block(x, t, mask=single_mask, rope=rope)
+    def clone(self) -> "AuKStepContext":
+        """A copy whose tensors own fresh storage."""
 
-        x = x[:, text_len + prompt_len :]
-        return self.proj_out(self.norm_out(x, t))
+        def _clone(value: torch.Tensor | None) -> torch.Tensor | None:
+            return None if value is None else value.clone()
+
+        return AuKStepContext(
+            c=self.c.clone(),
+            prompt=_clone(self.prompt),
+            target_mask=_clone(self.target_mask),
+            c_mask=_clone(self.c_mask),
+            audio_mask=_clone(self.audio_mask),
+            single_mask=_clone(self.single_mask),
+            joint_bias=_clone(self.joint_bias),
+            single_bias=_clone(self.single_bias),
+            rope_audio=(self.rope_audio[0].clone(), self.rope_audio[1].clone()),
+            rope_text=(self.rope_text[0].clone(), self.rope_text[1].clone()),
+            rope_single=(self.rope_single[0].clone(), self.rope_single[1].clone()),
+            branches=self.branches,
+        )
 
 
 def dit_state_dict(
@@ -686,6 +880,7 @@ def sample_latents(
                     ref_mask=ref_mask,
                     timestep=t[i],
                     cfg_strength=cfg_strength,
+                    new_request=i == 0,
                 )
                 x = x + (t[i + 1] - t[i]) * velocity
         finally:
@@ -718,22 +913,21 @@ def _sample_latents(
     x = initial_latents
     guided = cfg_strength >= 1e-5
     try:
+        # The conditioning is fixed across the loop: prepare it once.
+        ctx = dit.prepare(
+            text,
+            target_len=x.shape[1],
+            c_mask=c_mask,
+            ref=ref,
+            ref_mask=ref_mask,
+            cfg_infer=guided,
+            cache=guided,
+        )
         for i in range(timesteps.shape[0] - 1):
+            v = dit.step(x, timesteps[i], ctx)
             if guided:
-                pred = dit(
-                    x,
-                    text,
-                    timesteps[i],
-                    c_mask=c_mask,
-                    ref=ref,
-                    ref_mask=ref_mask,
-                    cfg_infer=True,
-                    cache=True,
-                )
-                v_cond, v_uncond = pred.chunk(2, dim=0)
+                v_cond, v_uncond = v.chunk(2, dim=0)
                 v = v_cond + (v_cond - v_uncond) * cfg_strength
-            else:
-                v = dit(x, text, timesteps[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask)
             x = x + (timesteps[i + 1] - timesteps[i]) * v
     finally:
         # The cached text projections belong to this request only; a failed

vllm_omni/diffusion/models/auk/cudagraph_wrapper.py

@@ -12,7 +12,7 @@ import torch.nn.functional as F
 from vllm.logger import init_logger
 from vllm.utils.math_utils import round_up
 
-from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKStepContext, AuKTransformer
 
 logger = init_logger(__name__)
 
@@ -21,11 +21,7 @@ logger = init_logger(__name__)
 class _GraphEntry:
     graph: torch.cuda.CUDAGraph
     static_x: torch.Tensor
-    static_x_mask: torch.Tensor
-    static_text: torch.Tensor
-    static_c_mask: torch.Tensor
-    static_ref: torch.Tensor
-    static_ref_mask: torch.Tensor
+    static_ctx: AuKStepContext
     static_timestep: torch.Tensor
     static_cfg: torch.Tensor | None
     static_out: torch.Tensor
@@ -38,10 +34,15 @@ class AuKCUDAGraphWrapper:
     whether the CFG branch is enabled. Timestep and the CFG strength are
     mutable scalar buffers, so all Euler steps and CFG values within one path
     reuse one graph.
+
+    Only :meth:`AuKTransformer.step` is captured. The per-request context
+    (text projection, reference embedding, padding biases, rotary tables) is
+    built by :meth:`AuKTransformer.prepare` on the first step of a request and
+    copied into the graph's static context, so the replayed steps skip it.
     """
 
-    _TARGET_ALIGNMENT = 64
-    _TEXT_ALIGNMENT = 64
+    _TARGET_ALIGNMENT = 32
+    _TEXT_ALIGNMENT = 32
     _REF_ALIGNMENT = 50
 
     def __init__(self, dit: AuKTransformer, *, enabled: bool = True, max_graphs: int = 32) -> None:
@@ -50,6 +51,10 @@ class AuKCUDAGraphWrapper:
         self.max_graphs = max(1, int(max_graphs))
         self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict()
         self._pool_handle: int | None = None
+        # Key whose static context holds the current request, and the context
+        # the eager fallback reuses across the steps of one request.
+        self._active_key: tuple | None = None
+        self._eager_ctx: AuKStepContext | None = None
         logger.info("Initialized AuK DiT lazy CUDA graph cache: enabled=%s slots=%d", self.enabled, self.max_graphs)
 
     @classmethod
@@ -87,6 +92,41 @@ class AuKCUDAGraphWrapper:
         """Retire all graphs together so none outlive shared workspaces."""
         if len(self._cache) >= self.max_graphs:
             self._cache.clear()
+            self._active_key = None
+
+    def _prepare(
+        self,
+        x: torch.Tensor,
+        x_mask: torch.Tensor | None,
+        text: torch.Tensor,
+        c_mask: torch.Tensor,
+        ref: torch.Tensor,
+        ref_mask: torch.Tensor,
+        uses_cfg: bool,
+    ) -> AuKStepContext:
+        return self.dit.prepare(
+            text,
+            target_len=x.shape[1],
+            mask=x_mask,
+            c_mask=c_mask,
+            ref=ref,
+            ref_mask=ref_mask,
+            cfg_infer=uses_cfg,
+        )
+
+    def _step(
+        self,
+        x: torch.Tensor,
+        timestep: torch.Tensor,
+        ctx: AuKStepContext,
+        cfg_strength: torch.Tensor | float,
+    ) -> torch.Tensor:
+        """One guided (or unguided) velocity; the CFG combine is part of the graph."""
+        velocity = self.dit.step(x, timestep, ctx)
+        if ctx.branches == 1:
+            return velocity
+        conditional, unconditional = velocity.chunk(2, dim=0)
+        return conditional + (conditional - unconditional) * cfg_strength
 
     @torch.no_grad()
     def __call__(
@@ -99,25 +139,37 @@ class AuKCUDAGraphWrapper:
         ref_mask: torch.Tensor,
         timestep: torch.Tensor,
         cfg_strength: float,
+        new_request: bool = True,
     ) -> torch.Tensor:
+        """Velocity of ``x`` at ``timestep``.
+
+        ``new_request`` marks the first step of a request, whose conditioning
+        (``text``, ``c_mask``, ``ref``, ``ref_mask``) is then prepared anew; the
+        later steps of the request pass ``False`` and reuse it. The default
+        prepares on every call, which is always correct.
+        """
         uses_cfg = cfg_strength >= 1e-5
         if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
-            if uses_cfg:
-                return self._run_cfg(x, None, text, c_mask, ref, ref_mask, timestep, cfg_strength=cfg_strength)
-            return self._run(x, None, text, c_mask, ref, ref_mask, timestep)
+            if new_request or self._eager_ctx is None:
+                self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg)
+            return self._step(x, timestep, self._eager_ctx, cfg_strength)
 
         target_frames = x.shape[1]
         x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask)
-        inputs = (x, x_mask, text, c_mask, ref, ref_mask, timestep)
         key = self._key(x, text, ref, uses_cfg)
         entry = self._cache.get(key)
+        ctx = None
+        if entry is None or new_request or key != self._active_key:
+            ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg)
         if entry is None:
             # Same Retirement as #6587, details discussed in #7469
             self._retire_graph_generation_if_full()
             try:
-                entry = self._capture(*inputs, cfg_strength=cfg_strength, uses_cfg=uses_cfg)
+                assert ctx is not None
+                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength)
             except Exception:
                 self._cache.clear()
+                self._active_key = None
                 self.enabled = False
                 logger.exception(
                     "Disabling AuK DiT CUDA graphs after capture failure for key=%s; "
@@ -128,112 +180,55 @@ class AuKCUDAGraphWrapper:
             self._cache[key] = entry
         else:
             self._cache.move_to_end(key)
+            if ctx is not None:
+                entry.static_ctx.copy_(ctx)
+        self._active_key = key
 
         entry.static_x.copy_(x)
-        entry.static_x_mask.copy_(x_mask)
-        entry.static_text.copy_(text)
-        entry.static_c_mask.copy_(c_mask)
-        entry.static_ref.copy_(ref)
-        entry.static_ref_mask.copy_(ref_mask)
         entry.static_timestep.copy_(timestep)
         if entry.static_cfg is not None:
             entry.static_cfg.fill_(cfg_strength)
         entry.graph.replay()
         return entry.static_out[:, :target_frames].clone()
 
-    def _run(
-        self,
-        x: torch.Tensor,
-        x_mask: torch.Tensor | None,
-        text: torch.Tensor,
-        c_mask: torch.Tensor,
-        ref: torch.Tensor,
-        ref_mask: torch.Tensor,
-        timestep: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.dit(x, text, timestep, mask=x_mask, c_mask=c_mask, ref=ref, ref_mask=ref_mask)
-
-    def _run_cfg(
-        self,
-        x: torch.Tensor,
-        x_mask: torch.Tensor | None,
-        text: torch.Tensor,
-        c_mask: torch.Tensor,
-        ref: torch.Tensor,
-        ref_mask: torch.Tensor,
-        timestep: torch.Tensor,
-        *,
-        cfg_strength: torch.Tensor | float,
-    ) -> torch.Tensor:
-        pred = self.dit(
-            x,
-            text,
-            timestep,
-            mask=x_mask,
-            c_mask=c_mask,
-            ref=ref,
-            ref_mask=ref_mask,
-            cfg_infer=True,
-            cache=True,
-        )
-        conditional, unconditional = pred.chunk(2, dim=0)
-        return conditional + (conditional - unconditional) * cfg_strength
-
     def _capture(
         self,
-        *inputs: torch.Tensor,
+        x: torch.Tensor,
+        ctx: AuKStepContext,
+        timestep: torch.Tensor,
+        *,
         cfg_strength: float,
-        uses_cfg: bool,
     ) -> _GraphEntry:
-        static_inputs = tuple(value.clone() for value in inputs)
+        static_x = x.clone()
+        static_ctx = ctx.clone()
+        static_timestep = timestep.clone()
         static_cfg = None
-        if uses_cfg:
-            static_cfg = torch.empty((), device=static_inputs[0].device, dtype=torch.float32)
+        if ctx.branches == 2:
+            static_cfg = torch.empty((), device=static_x.device, dtype=torch.float32)
             static_cfg.fill_(cfg_strength)
-        try:
-            for _ in range(3):
-                if uses_cfg:
-                    assert static_cfg is not None
-                    self._run_cfg(*static_inputs, cfg_strength=static_cfg)
-                else:
-                    self._run(*static_inputs)
-            # CFG warm-up populates AuKTransformer's Python-side projected-text
-            # cache. Clear it before capture so the graph includes
-            # project_text(static_text), rather than closing over the first
-            # request's projection and ignoring later static_text updates.
-            self.dit.clear_cache()
-            if static_cfg is not None:
-                static_cfg.fill_(cfg_strength)
-            if self._pool_handle is None:
-                self._pool_handle = torch.cuda.graph_pool_handle()
-            graph = torch.cuda.CUDAGraph()
-            with torch.cuda.graph(graph, pool=self._pool_handle):
-                # Capture the actual CFG path. The scalar value remains a
-                # mutable buffer, so cfg=2 and cfg=3 share this graph.
-                if uses_cfg:
-                    assert static_cfg is not None
-                    static_out = self._run_cfg(*static_inputs, cfg_strength=static_cfg)
-                else:
-                    static_out = self._run(*static_inputs)
-        finally:
-            self.dit.clear_cache()
+        scale = static_cfg if static_cfg is not None else cfg_strength
+        # Warm-up runs any lazy torch.compile outside the capture.
+        for _ in range(3):
+            self._step(static_x, static_timestep, static_ctx, scale)
+        if self._pool_handle is None:
+            self._pool_handle = torch.cuda.graph_pool_handle()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, pool=self._pool_handle):
+            # The CFG scalar stays a mutable buffer, so cfg=2 and cfg=3 share this graph.
+            static_out = self._step(static_x, static_timestep, static_ctx, scale)
 
         logger.info(
             "Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s",
-            inputs[0].shape[1],
-            inputs[2].shape[1],
-            inputs[4].shape[1],
+            x.shape[1],
+            ctx.c.shape[1],
+            0 if ctx.prompt is None else ctx.prompt.shape[1],
             cfg_strength,
         )
         return _GraphEntry(
             graph=graph,
-            static_x=static_inputs[0],
-            static_x_mask=static_inputs[1],
-            static_text=static_inputs[2],
-            static_c_mask=static_inputs[3],
-            static_ref=static_inputs[4],
-            static_ref_mask=static_inputs[5],
-            static_timestep=static_inputs[6],
+            static_x=static_x,
+            static_ctx=static_ctx,
+            static_timestep=static_timestep,
             static_cfg=static_cfg,
             static_out=static_out,
         )

vllm_omni/diffusion/models/auk/pipeline_auk.py

@@ -13,6 +13,7 @@ from __future__ import annotations
 
 import json
 import os
+import time
 from collections.abc import Iterable
 from typing import Any, ClassVar
 
@@ -23,6 +24,7 @@ from safetensors import safe_open
 from torch import nn
 from vllm.logger import init_logger
 
+from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents
@@ -50,6 +52,13 @@ _FUSION_KEYS = frozenset({"layer_weights", "layer_scale"})
 # The distilled variant only reproduces its training recipe at these settings.
 _FLASH_NFE = 4
 _FLASH_CFG = 0.0
+# Startup DiT warm-up shapes: target frames of 3 / 6 / 12 s requests, with no
+# reference (instruct TTS) and a 3 s one (zero-shot). The first call compiles
+# the blocks; each shape then captures its graph bucket. Other shapes still
+# capture lazily, at about a second each.
+_WARMUP_TARGET_FRAMES = (150, 300, 600)
+_WARMUP_REF_FRAMES = (0, 150)
+_WARMUP_TEXT_TOKENS = 96
 # Decorrelates the VAE posterior draw from the initial latent for the same request seed.
 _VAE_SEED_OFFSET = 0x5EED_0A0C
 
@@ -182,6 +191,7 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
 
         # The transformer owns every key the checkpoint tool writes into `dit`,
         # `attn_mask_enabled` included, so the section passes straight through.
+        self.text_hidden_dim = int(config["dit"]["text_hidden_dim"])
         self.dit = AuKTransformer(latent_dim=self.latent_dim, **config["dit"])
         self.dit = self.dit.to(dtype=self.dtype)
         self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True)
@@ -218,22 +228,70 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
 
         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.
+        compiles the whole denoise step;
+        regional compiles the double- and single-stream blocks. Compilation is
+        lazy and happens in the DiT CUDA graph's warm-up, before capture, so
+        each graph replays the compiled kernels. 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
+        dynamic = self.od_config.diffusion_compile_dynamic
         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,
-            )
+            # The samplers drive prepare() and step() directly rather than
+            # __call__, so the per-step body is what gets compiled.
+            self.dit.step = torch.compile(self.dit.step, dynamic=dynamic)
         else:
-            logger.info("AuK DiT declares no repeated blocks; regional torch.compile is a no-op, running eager")
+            regionally_compile(self.dit, dynamic=dynamic)
+        logger.info("AuK DiT configured for lazy %s torch.compile with dynamic=%s", granularity, dynamic)
+        self._warmup_dit()
         self.vae_decode.warmup(self.device)
 
+    def _warmup_dit(self) -> None:
+        """Compile the DiT step and capture the common graph buckets before serving.
+
+        The engine's synthetic warm-up request cannot reach this pipeline
+        (``dummy_run_num_frames = 0``), so without this the first real request
+        pays the whole torch.compile. Runs the variant's default guidance path
+        on zero conditioning; ``model_config.auk_dit_warmup_frames`` overrides
+        the target lengths and ``[]`` disables the warm-up.
+        """
+        if not self.cudagraph_wrapper.enabled:
+            return
+        model_config = getattr(self.od_config, "model_config", None) or {}
+        frames = model_config.get("auk_dit_warmup_frames", _WARMUP_TARGET_FRAMES)
+        if not frames:
+            logger.info("AuK DiT startup warm-up disabled")
+            return
+        cfg = _FLASH_CFG if self.is_flash else self.default_cfg
+        text = torch.zeros(1, _WARMUP_TEXT_TOKENS, self.text_hidden_dim, device=self.device, dtype=self.dtype)
+        c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device)
+        timestep = torch.zeros((), device=self.device, dtype=torch.float32)
+        start = time.perf_counter()
+        with torch.inference_mode(), self._dit_autocast():
+            for ref_frames in _WARMUP_REF_FRAMES:
+                ref = torch.zeros(1, ref_frames, self.latent_dim, device=self.device, dtype=torch.float32)
+                ref_mask = torch.ones(ref.shape[:2], dtype=torch.bool, device=self.device)
+                for target_frames in frames:
+                    x = torch.zeros(1, int(target_frames), self.latent_dim, device=self.device, dtype=torch.float32)
+                    self.cudagraph_wrapper(
+                        x=x,
+                        text=text,
+                        c_mask=c_mask,
+                        ref=ref,
+                        ref_mask=ref_mask,
+                        timestep=timestep,
+                        cfg_strength=cfg,
+                        new_request=True,
+                    )
+        if self.device.type != "cpu":
+            torch.accelerator.synchronize(self.device)
+        logger.info(
+            "AuK DiT warm-up: compiled and captured %d shapes in %.1f s",
+            len(frames) * len(_WARMUP_REF_FRAMES),
+            time.perf_counter() - start,
+        )
+
     # 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.