[Feat][Perf][AuK]Support Online Serving and DiT CUDAGraph (#7469)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>

tests/diffusion/models/auk/test_auk_cudagraph.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Single-request eager and CUDA-graph AuK DiT sampling equivalence."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, sample_latents
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_dit(device: str) -> AuKTransformer:
+    torch.manual_seed(12)
+    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,
+        )
+        .eval()
+        .to(device)
+    )
+
+
+def _sample_inputs(device: str, offset: float = 0.0) -> dict[str, torch.Tensor]:
+    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),
+    }
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+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")
+    capture_spy = mocker.spy(wrapper, "_capture")
+    inputs = _sample_inputs("cpu")
+    common = dict(
+        **inputs,
+        gen_frames=9,
+        t_grid=[0.0, 0.4, 1.0],
+        cfg_strength=cfg_strength,
+    )
+    eager = sample_latents(dit, **common, generator=torch.Generator().manual_seed(7))
+    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()
+
+    capture_spy.assert_not_called()
+    assert not wrapper._cache
+
+
+def test_graph_inputs_use_bounded_length_buckets() -> None:
+    wrapper = AuKCUDAGraphWrapper(_make_dit("cpu"))
+    x = torch.ones(1, 65, 4)
+    text = torch.ones(1, 65, 8)
+    c_mask = torch.ones(1, 65, dtype=torch.bool)
+    ref = torch.ones(1, 51, 4)
+    ref_mask = torch.ones(1, 51, dtype=torch.bool)
+
+    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[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.max_graphs == 32
+
+
+def test_full_graph_cache_retires_as_one_generation() -> None:
+    wrapper = AuKCUDAGraphWrapper(_make_dit("cpu"), max_graphs=2)
+    first = (64, 64, 50, False)
+    second = (128, 64, 50, False)
+    wrapper._cache[first] = object()
+
+    wrapper._retire_graph_generation_if_full()
+    assert list(wrapper._cache) == [first]
+
+    wrapper._cache[second] = object()
+    wrapper._retire_graph_generation_if_full()
+    assert not wrapper._cache
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_bucket_padding_preserves_real_frame_outputs(cfg_strength: float) -> None:
+    dit = _make_dit("cpu")
+    wrapper = AuKCUDAGraphWrapper(dit)
+    inputs = _sample_inputs("cpu")
+    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()
+
+    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()
+
+    torch.testing.assert_close(padded[:, : x.shape[1]], eager)
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+@pytest.mark.parametrize("cfg_strength", [0.0, 2.0])
+def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_strength: float) -> None:
+    dit = _make_dit("cuda")
+    wrapper = AuKCUDAGraphWrapper(dit)
+
+    for offset in (0.0, 0.25):
+        inputs = _sample_inputs("cuda", offset)
+        common = dict(
+            **inputs,
+            gen_frames=9,
+            t_grid=[0.0, 0.4, 1.0],
+            cfg_strength=cfg_strength,
+        )
+        eager = sample_latents(dit, **common, generator=torch.Generator(device="cuda").manual_seed(7))
+        graph = sample_latents(dit, **common, generator=torch.Generator(device="cuda").manual_seed(7), sampler=wrapper)
+        torch.testing.assert_close(graph, eager, atol=3e-6, rtol=3e-5)
+
+        bucketed = wrapper._bucket_inputs(
+            torch.empty(1, 9, 4, device="cuda"),
+            inputs["text"],
+            inputs["c_mask"],
+            inputs["ref"],
+            inputs["ref_mask"],
+        )
+        key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5)
+        assert key in wrapper._cache
+
+        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])
+        torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda"))
+
+    assert len(wrapper._cache) == 1
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA")
+@torch.inference_mode()
+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")
+    inputs = _sample_inputs("cuda")
+    common = dict(
+        **inputs,
+        gen_frames=9,
+        t_grid=[0.0, 0.4, 1.0],
+        cfg_strength=0.0,
+    )
+
+    with pytest.raises(RuntimeError, match="capture failed"):
+        sample_latents(
+            dit,
+            **common,
+            generator=torch.Generator(device="cuda").manual_seed(7),
+            sampler=wrapper,
+        )
+
+    result = sample_latents(
+        dit,
+        **common,
+        generator=torch.Generator(device="cuda").manual_seed(7),
+        sampler=wrapper,
+    )
+
+    assert result.shape == (1, 9, 4)
+    assert capture.call_count == 1
+    assert eager.call_count == 2
+    assert not wrapper.enabled
+    assert not wrapper._cache

tests/diffusion/models/auk/test_pipeline_auk.py

@@ -227,13 +227,16 @@ def build_pipeline(tmp_path, monkeypatch):
     monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer)
     monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE)
 
-    def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]:
+    def _build(
+        variant: str = "base", *, model_config: dict[str, Any] | None = None
+    ) -> tuple[AuKPipeline, list[dict[str, Any]]]:
         calls: list[dict[str, Any]] = []
         monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls))
         od_config = OmniDiffusionConfig(
             model=str(_write_checkpoint(tmp_path, variant)),
             dtype=torch.float32,
             model_class_name="AuKPipeline",
+            model_config=model_config or {},
         )
         return AuKPipeline(od_config=od_config), calls
 
@@ -274,6 +277,20 @@ def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray
 @pytest.mark.core_model
 @pytest.mark.cpu
 class TestRequestParsing:
+    @pytest.mark.parametrize(
+        "model_config, expected_slots",
+        [({}, 32), ({"max_dit_graphs": 1}, 1), ({"max_dit_graphs": 3}, 3), ({"max_dit_graphs": 64}, 64)],
+    )
+    def test_max_dit_graphs_reaches_lazy_cache(self, build_pipeline, model_config, expected_slots):
+        pipeline, _ = build_pipeline(model_config=model_config)
+        assert pipeline.cudagraph_wrapper.max_graphs == expected_slots
+        assert not pipeline.cudagraph_wrapper._cache
+
+    @pytest.mark.parametrize("max_dit_graphs", [0, -1, True, False, 1.5, "3", None])
+    def test_invalid_max_dit_graphs_is_rejected(self, build_pipeline, max_dit_graphs):
+        with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"):
+            build_pipeline(model_config={"max_dit_graphs": max_dit_graphs})
+
     def test_pipeline_declares_audio_output(self, build_pipeline):
         pipeline, _ = build_pipeline()
 
@@ -299,6 +316,7 @@ class TestRequestParsing:
         assert calls[0]["ref"].shape == (1, 100, LATENT_DIM)
         assert calls[0]["ref_mask"].shape == (1, 100)
         assert calls[0]["c_mask"].shape == (1, 8)
+        assert calls[0]["sampler"] is pipeline.cudagraph_wrapper
         assert outputs[0].output.shape == (300 * HOP,)
         assert outputs[0].output.dtype is torch.float32