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

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

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