[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>

benchmarks/tts/bench_tts.py

@@ -1,4 +1,7 @@
 #!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Universal TTS benchmark CLI for vllm-omni.
 
 Runs ``vllm bench serve --omni`` with model-aware defaults loaded from
@@ -87,6 +90,7 @@ def build_bench_args(
     served_model_name: str | None = None,
     num_warmups: int = 2,
     request_seed: int | None = None,
+    duration_seconds: float | None = None,
 ) -> list[str]:
     """Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
     dataset_name = _TASK_TO_DATASET[task]
@@ -95,7 +99,10 @@ def build_bench_args(
     task_extra_body: dict[str, Any] = dict((model_cfg.get("task_extra_body") or {}).get(task) or {})
     if request_seed is not None:
         task_extra_body["seed"] = request_seed
-
+    if duration_seconds is not None:
+        if not math.isfinite(duration_seconds) or duration_seconds <= 0:
+            raise ValueError("duration_seconds must be finite and positive")
+        task_extra_body["duration_seconds"] = duration_seconds
     # Resolve dataset path
     if dataset_path:
         resolved_dataset_path = dataset_path
@@ -194,7 +201,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None:
     if not results:
         return
     header = (
-        f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} "
+        f"{'Task':<16} {'Task Type':<12} {'Concurrency':>11} {'RTF mean':>10} "
         f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}"
     )
     print(f"\n{'=' * len(header)}")
@@ -204,6 +211,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None:
     print("-" * len(header))
     for r in results:
         task = r.get("_task", "?")
+        task_type = r.get("_task_type", "-")
         conc = r.get("_concurrency", "?")
         rtf = r.get("mean_audio_rtf", float("nan"))
         ttfp = r.get("mean_audio_ttfp_ms", float("nan"))
@@ -216,7 +224,7 @@ def print_summary_table(results: list[dict[str, Any]]) -> None:
             return f"{v:.{digits}f}" if not math.isnan(v) else "  n/a"
 
         print(
-            f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
+            f"{task:<16} {task_type:<12} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
             f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}"
         )
     print("=" * len(header))
@@ -249,6 +257,12 @@ def main() -> None:
         help="Number of prompts per run. If one value, applied to all concurrency levels.",
     )
     parser.add_argument("--num-warmups", type=int, default=2, help="Warmup requests before each measured run")
+    parser.add_argument(
+        "--duration-seconds",
+        type=float,
+        default=None,
+        help="Target audio duration per request (AuK defaults to 5 seconds). Independent of --output-len.",
+    )
     parser.add_argument(
         "--request-seed",
         type=int,
@@ -277,6 +291,8 @@ def main() -> None:
 
     if args.num_warmups < 0:
         parser.error("--num-warmups cannot be negative")
+    if args.duration_seconds is not None and (not math.isfinite(args.duration_seconds) or args.duration_seconds <= 0):
+        parser.error("--duration-seconds must be finite and positive")
 
     model_configs = load_model_configs(Path(args.model_configs))
     if args.model not in model_configs:
@@ -333,11 +349,16 @@ def main() -> None:
                 served_model_name=args.served_model_name,
                 num_warmups=args.num_warmups,
                 request_seed=args.request_seed,
+                duration_seconds=args.duration_seconds,
             )
             result = run_one_benchmark(cmd)
             if result is not None:
                 result["_task"] = task
                 result["_concurrency"] = concurrency
+                task_body = model_cfg.get("task_extra_body", {}).get(task, {})
+                task_type = task_body.get("task_type")
+                if task_type is not None:
+                    result["_task_type"] = task_type
                 all_results.append(result)
                 # Persist the metadata so plot_results.py can pick it up.
                 if args.output_dir and result_filename:

benchmarks/tts/model_configs.yaml

@@ -7,6 +7,22 @@
 # is tracked here.
 
 models:
+  tencent/AuK: &auk
+    supported_tasks: [default_voice, voice_clone]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    trust_remote_code: true
+    task_extra_body:
+      default_voice:
+        duration_seconds: 5.0
+        voice: default
+        task_type: CustomVoice
+      voice_clone:
+        duration_seconds: 5.0
+        voice: default
+        task_type: Base
+  tencent/AuK-Flash: *auk
+
   BreezeBlue/Breeze-TTS-2:
     supported_tasks: [default_voice, voice_design, voice_clone]
     backend: openai-audio-speech

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
 
 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
 | --- | --- | --- | --- | --- | --- |
+| AuK / AuK-Flash | assembled local bundle | ✓ (`ref_audio`) | HTTP audio after full generation | — | — |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
 | Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — |
@@ -99,6 +100,41 @@ For full request-shape documentation (all parameters, response formats, error co
 
 ---
 
+## AuK / AuK-Flash
+
+AuK serves speech generation and reference-voice synthesis at 24 kHz. Assemble a local bundle before serving:
+
+```bash
+python tools/prepare_auk_checkpoint.py \
+    --auk-dir ckpts/AuK --qwen-dir ckpts/Qwen2.5-Omni-3B --out ckpts/auk-omni
+MODEL=ckpts/auk-omni bash examples/online_serving/text_to_speech/auk/run_server.sh
+```
+
+AuK accepts the complete instruction after you fill in the [local recipe template](../../../../recipes/Tencent/AuK-H100.md#supported-model-contract) or the [official AuK cookbook](https://github.com/Tencent-Hunyuan/AuK/blob/main/docs/COOKBOOK.md). Pass it in `instructions` and leave `input` empty. For instruct TTS, fill in `Generate speech based on the following description: "{voice description}". The content to speak is: "{text}".`; for zero-shot TTS, editing, enhancement and separation, use the corresponding complete instruction from the cookbook. Set `duration_seconds` for text-only generation; with `ref_audio`, omitting it keeps the source length.
+
+For compatibility with clients that can only send `input`, AuK treats an input-only request as the complete instruction and logs a recommendation to move it to `instructions` with `input=""`.
+
+`extra_params` accepts `num_inference_steps`, `guidance_scale`, `sway`, `t_grid`, and `vae_sample`; `seed` is a top-level field.
+
+```bash
+python examples/online_serving/text_to_speech/auk/speech_client.py \
+    --model ckpts/auk-omni \
+    --instructions 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". The content to speak is: "Welcome back, how was your day?".' \
+    --duration-seconds 3 --seed 7
+```
+
+Complete-instruction editing also does not need a separate text argument:
+
+```bash
+python examples/online_serving/text_to_speech/auk/speech_client.py \
+    --model ckpts/auk-omni --instructions "Keep pure speech voice, remove noise and reverberation." \
+    --ref-audio noisy.wav --seed 7 --output clean.wav
+```
+
+The client accepts local reference files, URLs, and data URLs. `--stream` writes signed 16-bit 24 kHz PCM after the complete waveform is ready. See `speech_client.py --help` for editing and sampling options.
+
+---
+
 ## CosyVoice3
 
 2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.

examples/online_serving/text_to_speech/auk/run_server.sh

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# MODEL=/models/auk-omni bash examples/online_serving/text_to_speech/auk/run_server.sh
+set -euo pipefail
+
+: "${MODEL:?Set MODEL to the assembled AuK or AuK-Flash checkpoint directory}"
+exec vllm serve "$MODEL" --omni --host "${HOST:-127.0.0.1}" --port "${PORT:-8091}" "$@"

examples/online_serving/text_to_speech/auk/speech_client.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK speech generation and audio editing through /v1/audio/speech.
+
+See the online TTS user guide for checkpoint preparation and serving. Local
+reference audio is sent as a data URL; it need not exist on the server.
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import mimetypes
+import os
+from pathlib import Path
+
+import httpx
+
+
+def reference_audio(value: str) -> str:
+    if value.startswith(("http://", "https://", "data:")):
+        return value
+    path = Path(value).expanduser()
+    mime = mimetypes.guess_type(path.name)[0] or "audio/wav"
+    return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode('ascii')}"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--api-base", default="http://localhost:8091")
+    parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY", "EMPTY"))
+    parser.add_argument("--model", required=True, help="Model name or assembled bundle path used by the server")
+    parser.add_argument(
+        "--instructions",
+        required=True,
+        help=(
+            "Complete AuK instruction after filling the cookbook template "
+            "(https://github.com/Tencent-Hunyuan/AuK/blob/main/docs/COOKBOOK.md)"
+        ),
+    )
+    parser.add_argument("--ref-audio", help="Reference/source audio: local file, HTTP URL or data URL")
+    parser.add_argument("--duration-seconds", type=float, help="Target duration; defaults to source length with audio")
+    parser.add_argument("--seed", type=int)
+    parser.add_argument("--num-inference-steps", type=int)
+    parser.add_argument("--guidance-scale", type=float)
+    parser.add_argument("--sway", type=float)
+    parser.add_argument("--t-grid", type=float, nargs="+", help="Explicit increasing Euler time grid")
+    parser.add_argument("--vae-sample", action="store_true", help="Sample the VAE posterior instead of taking its mean")
+    parser.add_argument("--stream", action="store_true", help="Receive raw PCM over HTTP after full generation")
+    parser.add_argument("--response-format", choices=["wav", "pcm", "flac", "mp3", "opus"], default="wav")
+    parser.add_argument("--output", type=Path, help="Output file; defaults to output.wav or output.pcm")
+    parser.add_argument("--timeout", type=float, default=300)
+    args = parser.parse_args()
+    if args.duration_seconds is None and args.ref_audio is None:
+        parser.error("--duration-seconds is required without --ref-audio")
+
+    response_format = "pcm" if args.stream else args.response_format
+    payload = {
+        "model": args.model,
+        # OpenAI's shared Speech schema still declares input, while AuK's
+        # canonical request carries all content in the complete instruction.
+        "input": "",
+        "voice": "default",
+        "response_format": response_format,
+    }
+    for name in ("duration_seconds", "seed"):
+        value = getattr(args, name)
+        if value is not None:
+            payload[name] = value
+    payload["instructions"] = args.instructions
+    if args.ref_audio:
+        payload["ref_audio"] = reference_audio(args.ref_audio)
+    extra = {
+        name: getattr(args, name)
+        for name in ("num_inference_steps", "guidance_scale", "sway", "t_grid")
+        if getattr(args, name) is not None
+    }
+    if args.vae_sample:
+        extra["vae_sample"] = True
+    if extra:
+        payload["extra_params"] = extra
+    if args.stream:
+        payload.update(stream=True, stream_format="audio")
+
+    output = args.output or Path(f"output.{response_format}")
+    headers = {"Authorization": f"Bearer {args.api_key}"}
+    with httpx.Client(timeout=args.timeout) as client:
+        with client.stream(
+            "POST", f"{args.api_base.rstrip('/')}/v1/audio/speech", json=payload, headers=headers
+        ) as response:
+            if response.is_error:
+                response.read()
+                raise RuntimeError(f"Speech API returned {response.status_code}: {response.text}")
+            with output.open("wb") as handle:
+                for chunk in response.iter_bytes():
+                    handle.write(chunk)
+    print(f"Saved {output} ({output.stat().st_size} bytes). Native audio: 24000 Hz, mono.")
+    if response_format == "pcm":
+        print("PCM format: signed 16-bit little-endian, 24000 Hz, mono.")
+
+
+if __name__ == "__main__":
+    main()

tests/benchmarks/test_bench_tts_cli.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for the universal benchmarks/tts/bench_tts.py CLI."""
 
 from __future__ import annotations
@@ -51,6 +54,73 @@ def test_load_model_configs(model_configs_path: Path) -> None:
     assert configs["test/ModelA"]["supported_tasks"] == ["voice_clone", "default_voice"]
 
 
+@pytest.mark.parametrize("model", ["tencent/AuK", "tencent/AuK-Flash"])
+@pytest.mark.parametrize("task", ["default_voice", "voice_clone"])
+def test_auk_benchmark_duration_and_local_bundle(model, task):
+    config = bench_tts.load_model_configs(bench_tts._DEFAULT_MODEL_CONFIGS)[model]
+    cmd = bench_tts.build_bench_args(
+        host="localhost",
+        port=8000,
+        model=model,
+        task=task,
+        model_cfg=config,
+        locale="en",
+        num_prompts=20,
+        concurrency=8,
+        dataset_path="/data/seed-tts",
+        wer_eval=False,
+        output_dir=None,
+        result_filename=None,
+        extra_cli_args=[],
+        served_model_name="/models/auk",
+        duration_seconds=3.5,
+        request_seed=7,
+    )
+    assert cmd[cmd.index("--model") + 1] == "/models/auk"
+    body = json.loads(cmd[cmd.index("--extra-body") + 1])
+    assert body == {
+        "voice": "default",
+        "duration_seconds": 3.5,
+        "seed": 7,
+        "task_type": "CustomVoice" if task == "default_voice" else "Base",
+    }
+    assert config["task_extra_body"][task]["duration_seconds"] == 5.0
+
+
+@pytest.mark.parametrize(
+    ("task", "expected_task_type"),
+    [
+        ("default_voice", "CustomVoice"),
+        ("voice_clone", "Base"),
+    ],
+)
+def test_auk_task_type_survives_served_model_alias(task, expected_task_type):
+    model = "tencent/AuK-Flash"
+    served_model = "models/flash"
+    config = bench_tts.load_model_configs(bench_tts._DEFAULT_MODEL_CONFIGS)[model]
+    cmd = bench_tts.build_bench_args(
+        host="localhost",
+        port=8000,
+        model=model,
+        task=task,
+        model_cfg=config,
+        locale="en",
+        num_prompts=1,
+        concurrency=1,
+        dataset_path="/data/seed-tts",
+        wer_eval=False,
+        output_dir=None,
+        result_filename=None,
+        extra_cli_args=[],
+        served_model_name=served_model,
+    )
+
+    assert cmd[cmd.index("--model") + 1] == served_model
+    extra_body = json.loads(cmd[cmd.index("--extra-body") + 1])
+    assert extra_body["task_type"] == expected_task_type
+    assert "_vllm_omni_benchmark_model" not in extra_body
+
+
 def test_indextts25_is_registered_in_shared_model_configs() -> None:
     config_path = Path(bench_tts.__file__).with_name("model_configs.yaml")
     config = bench_tts.load_model_configs(config_path)["IndexTeam/IndexTTS-2.5"]
@@ -173,3 +243,20 @@ def test_unsupported_task_exits(model_configs_path: Path, capsys: pytest.Capture
     )
     with pytest.raises(SystemExit):
         bench_tts.main()
+
+
+def test_summary_table_includes_task_type(capsys: pytest.CaptureFixture) -> None:
+    bench_tts.print_summary_table(
+        [
+            {
+                "_task": "voice_clone",
+                "_task_type": "Base",
+                "_concurrency": 1,
+            }
+        ]
+    )
+
+    output = capsys.readouterr().out
+    assert "Task Type" in output
+    assert "voice_clone" in output
+    assert "Base" in output

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
 

tests/e2e/online_serving/test_auk.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK Speech API E2E coverage for the encoder -> diffusion deployment."""
+
+from __future__ import annotations
+
+import io
+import os
+
+import numpy as np
+import pytest
+import requests
+import soundfile as sf
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = os.environ.get("VLLM_OMNI_AUK_MODEL_DIR")
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.skipif(not MODEL, reason="set VLLM_OMNI_AUK_MODEL_DIR to an assembled AuK checkpoint"),
+]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("reference", [False, True], ids=["instruct", "voice-clone"])
+@pytest.mark.parametrize("stream", [False, True], ids=["wav", "streaming-pcm"])
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        pytest.param(
+            OmniServerParams(
+                model=MODEL or ".",
+                stage_config_path=get_deploy_config_path("auk.yaml"),
+                server_args=["--trust-remote-code", "--disable-log-stats"],
+            ),
+            id="auk",
+        )
+    ],
+    indirect=True,
+)
+def test_auk_speech_api_request(omni_server, reference, stream) -> None:
+    """The adapter accepts Speech API requests and returns bounded 24 kHz WAVs."""
+
+    def generate(duration: float) -> int:
+        extra = {}
+        if reference:
+            extra["ref_audio"] = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True)
+            # Match the shared seed-tts benchmark request envelope.
+            extra["ref_text"] = "Reference transcript is not required by AuK."
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech",
+            json={
+                "model": omni_server.model,
+                "input": "",
+                "instructions": (
+                    "Say the following with the same voice: 'Hello, this is an AuK speech API test.'"
+                    if reference
+                    else 'Generate speech based on the following description: "A clear, natural voice.". '
+                    'The content to speak is: "Hello, this is an AuK speech API test.".'
+                ),
+                "voice": "default",
+                "duration_seconds": duration,
+                "seed": 7,
+                "response_format": "pcm" if stream else "wav",
+                "stream": stream,
+                "stream_format": "audio" if stream else None,
+                **extra,
+            },
+            timeout=300,
+        )
+        assert response.status_code == 200, response.text
+        if stream:
+            waveform = np.frombuffer(response.content, dtype="<i2").astype(np.float32) / 32768
+        else:
+            waveform, sample_rate = sf.read(io.BytesIO(response.content), dtype="float32")
+            assert sample_rate == 24_000
+        assert np.isfinite(waveform).all()
+        assert np.max(np.abs(waveform)) > 0
+        return int(np.asarray(waveform).reshape(-1).shape[0])
+
+    assert generate(2.0) == 100 * 480

tests/entrypoints/openai_api/test_serving_speech.py

@@ -5367,6 +5367,9 @@ class TestTTSAsyncOffloading:
         adapter_model_type = "adapter_dummy_tts"
 
         class FakeAdapter:
+            def normalize(self, request):
+                return
+
             def validate(self, request):
                 return None
 

tests/entrypoints/openai_api/test_tts_adapter.py

@@ -8,11 +8,22 @@ Pure-Python registry/resolution logic; no model or GPU resources are loaded.
 import asyncio
 from types import SimpleNamespace
 from typing import Any
+from unittest.mock import AsyncMock
 
 import pytest
+import torch
 from vllm.sampling_params import SamplingParams
 
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.tts_adapters import (
     TTS_ADAPTER_REGISTRY,
     ARTTSAdapter,
@@ -22,6 +33,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (
     detect_tts_model_type,
     resolve_adapter,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.auk import AuKAdapter
 from vllm_omni.entrypoints.openai.tts_adapters.base import TTSCapabilities, resolve_stage_model_path
 from vllm_omni.entrypoints.openai.tts_adapters.covo_audio import CovoAudioAdapter
 from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v2 import HiggsAudioV2Adapter
@@ -42,10 +54,12 @@ from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import (
 )
 from vllm_omni.entrypoints.openai.tts_adapters.step_audio2 import StepAudio2Adapter
 from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.models.indextts2 import prompt_utils
 from vllm_omni.model_executor.models.indextts2.tokenizer_v2_5 import (
     INDEXTTS25_TOKENIZER_FILE,
 )
+from vllm_omni.model_executor.stage_input_processors.auk import encoder2dit
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -70,6 +84,7 @@ EXPECTED_MODEL_TYPES = {
     "step_audio2",
     "indextts2",
     "indextts2_5",
+    "auk",
     "gepard",
 }
 
@@ -99,6 +114,314 @@ def test_resolve_unknown_returns_none():
     assert resolve_adapter(None) is None
 
 
+@pytest.fixture
+def auk_adapter(mocker):
+    server = mocker.Mock(spec=OmniOpenAIServingSpeech)
+    server._max_instructions_length = 4096
+    server._validate_ref_audio_format.return_value = None
+    server._resolve_ref_audio = AsyncMock(return_value=([0.1] * 480, 24000, "key"))
+    return AuKAdapter(SpeechServingContext(server=server))
+
+
+def test_auk_adapter_detection(mocker):
+    assert detect_tts_model_type("encoder", "AuKForConditionalGeneration") == "auk"
+    # ``encoder`` is shared with MiniMax H3, so AuK must not claim the stage key.
+    assert AuKAdapter.stage_keys == frozenset()
+    assert AuKAdapter.arch_identifies_entry_stage
+    assert detect_tts_model_type("encoder", "MiniMaxH3Encoder") is None
+    assert detect_tts_model_type("encoder", None) is None
+
+    server = mocker.Mock(spec=OmniOpenAIServingSpeech)
+    server._diffusion_mode = False
+    server._tts_model_type = detect_tts_model_type("encoder", "AuKForConditionalGeneration")
+    server._adapter = None
+    server.engine_client = SimpleNamespace()
+
+    adapter = OmniOpenAIServingSpeech._get_tts_adapter(server)
+    assert isinstance(adapter, AuKAdapter)
+    server._drop_shadowing_uploads.assert_called_once_with()
+
+
+def test_auk_source_length_default_and_complete_instruction(auk_adapter):
+    request = OpenAICreateSpeechRequest(
+        input="",
+        ref_audio="reference.wav",
+        ref_text="Accepted transcript",
+        instructions="Change the pitch by one semitone.",
+        extra_params={"sway": -0.5, "t_grid": [0, 0.25, 1], "vae_sample": True},
+    )
+    assert auk_adapter.validate(request) is None
+    prepared = asyncio.run(auk_adapter.build(request, [], True))
+    assert "Change the pitch by one semitone." in prepared.prompt["prompt"]
+    assert "content to speak" not in prepared.prompt["prompt"]
+    knobs = prepared.prompt["additional_information"]["auk"]
+    assert knobs["gen_seconds"] is None
+    assert knobs["vae_sample"] is True
+    assert knobs["t_grid"] == [0, 0.25, 1]
+    auk_adapter.ctx.server._resolve_ref_audio.assert_awaited_once_with("reference.wav")
+
+
+@pytest.mark.parametrize(
+    "t_grid",
+    [
+        [],
+        [0.0],
+        [0.0, 0.0],
+        [0.5, 0.25],
+        [0.0, float("nan")],
+        [0.0, float("inf")],
+        [0.0, "invalid"],
+        "0,1",
+    ],
+)
+def test_auk_rejects_invalid_t_grid_before_dispatch(auk_adapter, t_grid):
+    request = OpenAICreateSpeechRequest(
+        input="Say the following: 'target text'",
+        duration_seconds=2,
+        extra_params={"t_grid": t_grid},
+    )
+
+    error = auk_adapter.validate(request)
+
+    assert error == "AuK extra_params.t_grid must contain at least two finite, strictly increasing values"
+    auk_adapter.ctx.server._resolve_ref_audio.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+    ("task_type", "expected_instruction"),
+    [
+        ("CustomVoice", 'Say the following: "target text"'),
+        ("Base", 'Say the following with the same voice: "target text"'),
+    ],
+)
+def test_auk_task_type_normalizes_benchmark_text(auk_adapter, task_type, expected_instruction, mocker):
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once")
+    request = OpenAICreateSpeechRequest(
+        input="target text",
+        task_type=task_type,
+        duration_seconds=2,
+        ref_audio="reference.wav" if task_type == "Base" else None,
+    )
+
+    server = mocker.Mock(spec=OmniOpenAIServingSpeech)
+    server._validate_speech_sample_rate.return_value = None
+    server._get_tts_adapter.return_value = auk_adapter
+    assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None
+
+    assert request.input == ""
+    assert request.instructions == expected_instruction
+    assert request.task_type is None
+    warning_once.assert_called_once()
+    assert "prefer a complete `instructions` prompt" in warning_once.call_args.args[0]
+
+    auk_adapter.normalize(request)
+
+    assert request.instructions == expected_instruction
+    warning_once.assert_called_once()
+
+
+def test_auk_task_type_quotes_embedded_delimiters(auk_adapter):
+    request = OpenAICreateSpeechRequest(
+        input='It\'s called "AuK"\\Flash',
+        task_type="CustomVoice",
+        duration_seconds=2,
+    )
+
+    auk_adapter.normalize(request)
+
+    assert request.instructions == 'Say the following: "It\'s called \\"AuK\\"\\\\Flash"'
+    assert request.task_type is None
+
+
+def test_auk_task_type_preserves_explicit_instructions(auk_adapter, mocker):
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once")
+    warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning")
+    request = OpenAICreateSpeechRequest(
+        input="target text",
+        task_type="Base",
+        instructions="Use this complete AuK instruction.",
+        duration_seconds=2,
+    )
+
+    server = mocker.Mock(spec=OmniOpenAIServingSpeech)
+    server._validate_speech_sample_rate.return_value = None
+    server._get_tts_adapter.return_value = auk_adapter
+    assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None
+
+    assert request.input == ""
+    assert request.instructions == "Use this complete AuK instruction."
+    assert request.task_type is None
+    warning_once.assert_called_once()
+    warning.assert_called_once()
+    assert "without applying another task template" in warning.call_args.args[0]
+
+
+def test_auk_task_type_does_not_double_wrap_preformatted_input(auk_adapter, mocker):
+    warning_once = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning_once")
+    warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning")
+    instruction = "Say the following with the same voice: 'target text'"
+    request = OpenAICreateSpeechRequest(
+        input=instruction,
+        task_type="Base",
+        duration_seconds=2,
+        ref_audio="reference.wav",
+    )
+    server = mocker.Mock(spec=OmniOpenAIServingSpeech)
+    server._validate_speech_sample_rate.return_value = None
+    server._get_tts_adapter.return_value = auk_adapter
+
+    assert OmniOpenAIServingSpeech._validate_tts_request(server, request) is None
+
+    assert request.input == ""
+    assert request.instructions == instruction
+    assert request.task_type is None
+    warning_once.assert_called_once()
+    warning.assert_called_once()
+    assert "already contains a complete task instruction" in warning.call_args.args[0]
+
+
+def test_auk_sampling_overrides_reach_stage1_pipeline(auk_adapter, mocker):
+    """Exercise one request from the Speech adapter through stage-1 admission."""
+
+    class RecordingDiffusionStage:
+        stage_type = "diffusion"
+        final_output = True
+        engine_input_source = [0]
+        requires_multimodal_data = False
+        custom_process_input_func = staticmethod(encoder2dit)
+
+        def __init__(self) -> None:
+            self.request: OmniDiffusionRequest | None = None
+
+        async def add_request_async(self, request_id, prompt, sampling_params, **_kwargs) -> None:
+            # Match the stage-client wire boundary and the receiving diffusion
+            # process, where params are serialized then reconstructed before
+            # creating the request consumed by the pipeline.
+            wire_params = StageDiffusionClient._sampling_params_to_dict(sampling_params)
+            self.request = OmniDiffusionRequest(
+                prompt=prompt,
+                sampling_params=StageDiffusionProc._reconstruct_sampling_params(
+                    object.__new__(StageDiffusionProc), wire_params
+                ),
+                request_id=request_id,
+            )
+
+    defaults = [SamplingParams(max_tokens=1), OmniDiffusionSamplingParams(seed=0, num_inference_steps=32)]
+    request = OpenAICreateSpeechRequest(
+        input="",
+        instructions='Generate speech based on the following description: "A calm voice". The content to speak is: "Hello".',
+        duration_seconds=2,
+        seed=9,
+        extra_params={
+            "num_inference_steps": 4,
+            "guidance_scale": 1.5,
+            "sway": -0.5,
+            "t_grid": [0.0, 0.25, 1.0],
+            "vae_sample": True,
+        },
+    )
+    prepared = asyncio.run(auk_adapter.build(request, defaults, False))
+    updated = auk_adapter.apply_sampling_overrides(defaults, request)
+    assert defaults[1].num_inference_steps == 32
+    assert defaults[1].seed == 0
+    assert updated[0].max_tokens == 1
+
+    source_stage = SimpleNamespace(
+        stage_type="llm",
+        final_output=False,
+        get_kv_sender_info=lambda: None,
+    )
+    diffusion_stage = RecordingDiffusionStage()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.stage_pools = [StagePool(0, source_stage), StagePool(1, diffusion_stage)]
+    orchestrator._cfg_tracker = CfgCompanionTracker()
+    orchestrator.duplex_control_plane = None
+    orchestrator._running_counter = None
+
+    req_state = OrchestratorRequestState(
+        request_id="auk-sampling-override",
+        prompt=prepared.prompt,
+        sampling_params_list=updated,
+        final_stage_id=1,
+    )
+    encoder_output = SimpleNamespace(
+        request_id=req_state.request_id,
+        finished=True,
+        multimodal_output={"hidden_states": {"output": torch.ones(3, 4)}},
+        outputs=[],
+    )
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, req_state.request_id, 0, encoder_output, req_state))
+
+    received = diffusion_stage.request
+    assert received is not None
+    assert received.prompt["prompt_embeds"].shape == (3, 4)
+    assert received.prompt["additional_information"]["auk"] == {
+        "gen_seconds": 2,
+        "sway": -0.5,
+        "t_grid": [0.0, 0.25, 1.0],
+        "vae_sample": True,
+        "has_audio": False,
+    }
+    assert received.sampling_params.num_inference_steps == 4
+    assert received.sampling_params.guidance_scale == 1.5
+    assert received.sampling_params.guidance_scale_provided is True
+    assert received.sampling_params.seed == 9
+
+    runner = mocker.Mock(spec=DiffusionModelRunner)
+    runner.device = torch.device("cpu")
+    DiffusionModelRunner._initialize_generator(runner, received.sampling_params)
+    assert received.sampling_params.generator.initial_seed() == 9
+
+    pipeline_params = DiffusionRequestBatch(requests=[received]).sampling_params
+    assert pipeline_params is received.sampling_params
+    assert pipeline_params.generator.initial_seed() == 9
+
+
+def test_auk_instructions_take_precedence_over_input(auk_adapter, mocker):
+    warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning")
+    request = OpenAICreateSpeechRequest(
+        input="ignored input",
+        instructions='Generate speech based on the following description: "Speak warmly". The content to speak is: "Hello".',
+        duration_seconds=2,
+    )
+    assert auk_adapter.validate(request) is None
+    prepared = asyncio.run(auk_adapter.build(request, [], False))
+    prompt = prepared.prompt
+    assert 'Generate speech based on the following description: "Speak warmly".' in prompt["prompt"]
+    assert 'The content to speak is: "Hello".' in prompt["prompt"]
+    assert "|<no_prompt_audio>|" in prompt["prompt"]
+    assert prompt["additional_information"]["auk"]["gen_seconds"] == 2
+    warning.assert_called_once()
+    assert "using the complete `instructions`" in warning.call_args.args[0]
+
+
+def test_auk_input_only_is_treated_as_complete_instruction(auk_adapter, mocker):
+    warning = mocker.patch("vllm_omni.entrypoints.openai.tts_adapters.auk.logger.warning")
+    request = OpenAICreateSpeechRequest(
+        input="Say the following with the same voice: 'Hello world.'",
+        ref_audio="reference.wav",
+        duration_seconds=2,
+    )
+    assert auk_adapter.validate(request) is None
+    prepared = asyncio.run(auk_adapter.build(request, [], True))
+    assert "Say the following with the same voice: 'Hello world.'" in prepared.prompt["prompt"]
+    assert warning.call_count == 1
+    assert 'Prefer `input=""`' in warning.call_args.args[0]
+
+
+def test_auk_complete_instruction_does_not_require_spoken_text(auk_adapter):
+    request = OpenAICreateSpeechRequest(
+        input="",
+        instructions="Replace 'morning' with 'evening' in the source recording.",
+        ref_audio="reference.wav",
+    )
+    assert auk_adapter.validate(request) is None
+    prepared = asyncio.run(auk_adapter.build(request, [], True))
+    assert prepared.prompt["prompt"].count("Replace 'morning'") == 1
+    assert "Generate speech based on" not in prepared.prompt["prompt"]
+    assert prepared.prompt["additional_information"]["auk"]["gen_seconds"] is None
+
+
 def test_stage_model_path_prefers_typed_override():
     engine_client = SimpleNamespace(
         stage_configs=[

tests/entrypoints/openai_api/test_tts_detection.py

@@ -182,6 +182,7 @@ _PIPELINE_STAGES = [
     "token2text",
     "token2wav",
     "tts",
+    "encoder",
 ]
 
 _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"]
@@ -283,14 +284,22 @@ def test_pipeline_stage_list_is_complete():
     assert not missing, f"stage keys claimed by adapters but absent from _PIPELINE_STAGES: {sorted(missing)}"
 
 
-def test_entry_stage_archs_is_ming_only():
-    """Only Ming dense identifies its entry stage by architecture.
+def test_entry_stage_archs_are_explicit():
+    """Only Ming dense and AuK identify their entry stage by architecture.
 
-    Widening this set would change which stage ``_find_tts_stage`` selects in
-    mixed deployments — notably VoxCPM2, which declares ``model_archs`` but is
-    found by its ``latent_generator`` stage key.
+    Both use a generic ``model_stage`` (``llm`` and ``encoder``) that other
+    pipelines share. Widening this set would change which stage
+    ``_find_tts_stage`` selects in mixed deployments — notably VoxCPM2, which
+    declares ``model_archs`` but is found by its ``latent_generator`` stage key.
     """
-    assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration"})
+    assert tts_entry_stage_archs() == frozenset({"MingTTSForConditionalGeneration", "AuKForConditionalGeneration"})
+
+
+def test_minimax_h3_encoder_is_not_a_tts_stage():
+    """MiniMax H3 shares AuK's ``encoder`` stage key but is not a speech model."""
+    assert "encoder" not in all_tts_stage_keys()
+    assert detect_tts_model_type("encoder", "MiniMaxH3Encoder") is None
+    assert detect_tts_model_type("encoder", "AuKForConditionalGeneration") == "auk"
 
 
 def test_every_detected_type_is_adapter_backed():

tests/helpers/client.py

@@ -1475,6 +1475,7 @@ class OnlineOmniClient:
             "max_new_tokens",
             "seed",
             "instructions",
+            "duration_seconds",
             "speed",
             "sample_rate",
             "extra_params",

tests/worker/test_gpu_ar_model_runner.py

@@ -214,6 +214,9 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
     )
 
     class Adapter:
+        def normalize(self, request):
+            return
+
         def validate(self, request):
             return None
 

tools/pre_commit/check_tts_adapter.py

@@ -29,7 +29,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2]
 SERVING_SPEECH = Path("vllm_omni/entrypoints/openai/serving_speech.py")
 
 # Budgets. These may only go DOWN. See the module docstring.
-MAX_MODEL_TYPE_BRANCHES = 6
+MAX_MODEL_TYPE_BRANCHES = 5
 
 _GUIDANCE = """
 Move the per-model behaviour into its adapter under

vllm_omni/deploy/auk.yaml

@@ -43,8 +43,8 @@ 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. The DiT stays eager unless
-  # diffusion_compile_granularity is set to full. Buckets, tile size and the
+  # decoded in overlapping tiles of it. Each DiT denoise step replays a lazily
+  # captured per-shape CUDA graph. 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}
   - stage_id: 1
@@ -52,6 +52,10 @@ stages:
     max_num_seqs: 1
     gpu_memory_utilization: 0.35
     enforce_eager: false
+    model_config:
+      # 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
     default_sampling_params:
       seed: 0
       num_inference_steps: 32

vllm_omni/diffusion/models/auk/auk_transformer.py

@@ -28,13 +28,13 @@ Deliberate differences from the reference, all inference-only:
 """
 
 import math
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-__all__ = ["AuKTransformer", "dit_state_dict", "sample_latents"]
+__all__ = ["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:
@@ -86,14 +86,17 @@ class Rotary(nn.Module):
         exponents = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim
         return 1.0 / (self.base**exponents)
 
-    def forward(self, seq_len: int) -> torch.Tensor:
-        """Return frequencies ``[1, 1, seq_len, dim]`` for positions ``0..seq_len-1``."""
+    def forward(self, seq_len: int, mask: torch.Tensor | None = None) -> torch.Tensor:
+        """Return rotary frequencies, compressing positions across padding."""
         inv_freq = self.inv_freq
         if inv_freq.dtype != torch.float32:
             inv_freq = self._frequencies(inv_freq.device)
-        pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32)
-        freqs = torch.outer(pos, inv_freq)
-        return torch.stack((freqs, freqs), dim=-1).flatten(-2)[None, None]
+        if mask is None:
+            pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32)[None]
+        else:
+            pos = (mask.to(torch.int32).cumsum(dim=1) - 1).clamp_min(0).to(torch.float32)
+        freqs = pos.unsqueeze(-1) * inv_freq
+        return torch.stack((freqs, freqs), dim=-1).flatten(-2).unsqueeze(1)
 
 
 class TimeEmbedding(nn.Module):
@@ -541,15 +544,15 @@ class AuKTransformer(nn.Module):
 
         seq_len = x.shape[1]
         text_len = c.shape[1]
-        rope_audio = self.rotary_embed(seq_len)
-        rope_text = self.rotary_embed(text_len)
+        rope_audio = self.rotary_embed(seq_len, audio_mask)
+        rope_text = self.rotary_embed(text_len, c_mask)
 
         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)
 
         x = torch.cat([c, x], dim=1)
-        rope = self.rotary_embed(text_len + seq_len)
         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)
 
         for block in self.single_transformer_blocks:
             x = block(x, t, mask=single_mask, rope=rope)
@@ -571,6 +574,31 @@ def dit_state_dict(
     return {name[len(prefix) :]: tensor for name, tensor in items if name.startswith(prefix)}
 
 
+def build_time_grid(
+    *,
+    nfe: int,
+    sway_sampling_coef: float | None,
+    t_grid: list[float] | None,
+    device: torch.device | str,
+) -> torch.Tensor:
+    """Build and validate the Euler schedule before a graph replay."""
+    if t_grid is not None:
+        timesteps = torch.tensor(t_grid, device="cpu", dtype=torch.float32)
+    else:
+        timesteps = torch.linspace(0, 1, nfe + 1, device="cpu", dtype=torch.float32)
+        if sway_sampling_coef is not None:
+            timesteps = timesteps + sway_sampling_coef * (torch.cos(math.pi / 2 * timesteps) - 1 + timesteps)
+    if (
+        timesteps.ndim != 1
+        or timesteps.numel() < 2
+        or not bool(torch.isfinite(timesteps).all() & torch.all(timesteps[1:] > timesteps[:-1]))
+    ):
+        raise ValueError(
+            f"AuK sampling needs a strictly increasing time grid with at least two points; got {timesteps.tolist()}"
+        )
+    return timesteps.to(device)
+
+
 @torch.no_grad()
 def sample_latents(
     dit: AuKTransformer,
@@ -589,6 +617,7 @@ def sample_latents(
     latent_dim: int | None = None,
     device: torch.device | str | None = None,
     dtype: torch.dtype | None = None,
+    sampler: Callable[..., torch.Tensor] | None = None,
 ) -> torch.Tensor:
     """Integrate the flow from noise to audio latents with explicit Euler steps.
 
@@ -640,25 +669,61 @@ def sample_latents(
             .unsqueeze(0)
         )
 
-    if t_grid is not None:
-        t = torch.tensor(t_grid, device=device, dtype=torch.float32)
-    else:
-        t = torch.linspace(0, 1, nfe + 1, device=device, dtype=torch.float32)
-        if sway_sampling_coef is not None:
-            t = t + sway_sampling_coef * (torch.cos(math.pi / 2 * t) - 1 + t)
+    t = build_time_grid(
+        nfe=nfe,
+        sway_sampling_coef=sway_sampling_coef,
+        t_grid=t_grid,
+        device=device,
+    )
+    if sampler is not None:
+        try:
+            for i in range(t.shape[0] - 1):
+                velocity = sampler(
+                    x=x,
+                    text=text,
+                    c_mask=c_mask,
+                    ref=ref,
+                    ref_mask=ref_mask,
+                    timestep=t[i],
+                    cfg_strength=cfg_strength,
+                )
+                x = x + (t[i + 1] - t[i]) * velocity
+        finally:
+            dit.clear_cache()
+        return x
+    return _sample_latents(
+        dit,
+        initial_latents=x,
+        text=text,
+        c_mask=c_mask,
+        ref=ref,
+        ref_mask=ref_mask,
+        timesteps=t,
+        cfg_strength=cfg_strength,
+    )
 
+
+def _sample_latents(
+    dit: AuKTransformer,
+    *,
+    initial_latents: torch.Tensor,
+    text: torch.Tensor,
+    c_mask: torch.Tensor | None,
+    ref: torch.Tensor,
+    ref_mask: torch.Tensor | None,
+    timesteps: torch.Tensor,
+    cfg_strength: float,
+) -> torch.Tensor:
+    """Euler integration shared by eager sampling and CUDA graph capture."""
+    x = initial_latents
     guided = cfg_strength >= 1e-5
-    if t.numel() < 2 or not bool(torch.all(t[1:] > t[:-1])):
-        raise ValueError(
-            f"AuK sampling needs a strictly increasing time grid with at least two points; got {t.tolist()}"
-        )
     try:
-        for i in range(t.shape[0] - 1):
+        for i in range(timesteps.shape[0] - 1):
             if guided:
                 pred = dit(
                     x,
                     text,
-                    t[i],
+                    timesteps[i],
                     c_mask=c_mask,
                     ref=ref,
                     ref_mask=ref_mask,
@@ -668,8 +733,8 @@ def sample_latents(
                 v_cond, v_uncond = pred.chunk(2, dim=0)
                 v = v_cond + (v_cond - v_uncond) * cfg_strength
             else:
-                v = dit(x, text, t[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask)
-            x = x + (t[i + 1] - t[i]) * v
+                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
         # step must not leak them into the next one.

vllm_omni/diffusion/models/auk/cudagraph_wrapper.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Per-shape CUDA graph for one AuK DiT denoise step."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from dataclasses import dataclass
+
+import torch
+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
+
+logger = init_logger(__name__)
+
+
+@dataclass
+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_timestep: torch.Tensor
+    static_cfg: torch.Tensor | None
+    static_out: torch.Tensor
+
+
+class AuKCUDAGraphWrapper:
+    """Replay one DiT denoise step and leave Euler scheduling to the caller.
+
+    The graph is keyed by the target, text and reference sequence lengths plus
+    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.
+    """
+
+    _TARGET_ALIGNMENT = 64
+    _TEXT_ALIGNMENT = 64
+    _REF_ALIGNMENT = 50
+
+    def __init__(self, dit: AuKTransformer, *, enabled: bool = True, max_graphs: int = 32) -> None:
+        self.dit = dit
+        self.enabled = bool(enabled)
+        self.max_graphs = max(1, int(max_graphs))
+        self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict()
+        self._pool_handle: int | None = None
+        logger.info("Initialized AuK DiT lazy CUDA graph cache: enabled=%s slots=%d", self.enabled, self.max_graphs)
+
+    @classmethod
+    def _bucket_inputs(
+        cls,
+        x: torch.Tensor,
+        text: torch.Tensor,
+        c_mask: torch.Tensor,
+        ref: torch.Tensor,
+        ref_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        target_bucket = round_up(x.shape[1], cls._TARGET_ALIGNMENT)
+        text_bucket = round_up(text.shape[1], cls._TEXT_ALIGNMENT)
+        ref_bucket = round_up(ref.shape[1], cls._REF_ALIGNMENT)
+        x_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device)
+        return (
+            F.pad(x, (0, 0, 0, target_bucket - x.shape[1])),
+            F.pad(x_mask, (0, target_bucket - x_mask.shape[1]), value=False),
+            F.pad(text, (0, 0, 0, text_bucket - text.shape[1])),
+            F.pad(c_mask, (0, text_bucket - c_mask.shape[1]), value=False),
+            F.pad(ref, (0, 0, 0, ref_bucket - ref.shape[1])),
+            F.pad(ref_mask, (0, ref_bucket - ref_mask.shape[1]), value=False),
+        )
+
+    @staticmethod
+    def _key(
+        x: torch.Tensor,
+        text: torch.Tensor,
+        ref: torch.Tensor,
+        uses_cfg: bool,
+    ) -> tuple[int, int, int, bool]:
+        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg)
+
+    def _retire_graph_generation_if_full(self) -> None:
+        """Retire all graphs together so none outlive shared workspaces."""
+        if len(self._cache) >= self.max_graphs:
+            self._cache.clear()
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        *,
+        x: torch.Tensor,
+        text: torch.Tensor,
+        c_mask: torch.Tensor,
+        ref: torch.Tensor,
+        ref_mask: torch.Tensor,
+        timestep: torch.Tensor,
+        cfg_strength: float,
+    ) -> torch.Tensor:
+        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)
+
+        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)
+        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)
+            except Exception:
+                self._cache.clear()
+                self.enabled = False
+                logger.exception(
+                    "Disabling AuK DiT CUDA graphs after capture failure for key=%s; "
+                    "subsequent requests will use eager execution.",
+                    key,
+                )
+                raise
+            self._cache[key] = entry
+        else:
+            self._cache.move_to_end(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,
+        cfg_strength: float,
+        uses_cfg: bool,
+    ) -> _GraphEntry:
+        static_inputs = tuple(value.clone() for value in inputs)
+        static_cfg = None
+        if uses_cfg:
+            static_cfg = torch.empty((), device=static_inputs[0].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()
+
+        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],
+            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_cfg=static_cfg,
+            static_out=static_out,
+        )
+
+
+__all__ = ["AuKCUDAGraphWrapper"]

vllm_omni/diffusion/models/auk/pipeline_auk.py

@@ -27,6 +27,7 @@ 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
 from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
+from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph
 from vllm_omni.diffusion.models.interface import (
     SupportAudioInput,
@@ -139,6 +140,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         super().__init__()
         del prefix  # Weights are not namespaced: one checkpoint, one pipeline.
         self.od_config = od_config
+        max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32)
+        if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1:
+            raise ValueError("AuK max_dit_graphs must be a positive integer")
         self.device = get_local_device()
         self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16
 
@@ -183,6 +187,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True)
         self.dit = self.dit.to(device=self.device).eval()
         self.dit.requires_grad_(False)
+        self.cudagraph_wrapper = AuKCUDAGraphWrapper(
+            self.dit, enabled=not od_config.enforce_eager, max_graphs=max_dit_graphs
+        )
         # The compiled decode buckets are warmed by setup_compile(), which the
         # model runner calls at startup unless the stage is enforce_eager. The
         # bucket list and the plain-graph cache size come from the stage's
@@ -462,6 +469,7 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
                     device=self.device,
                     dtype=torch.float32,
                     generator=generator,
+                    sampler=self.cudagraph_wrapper,
                 )
 
             latents = latents.float()

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1173,6 +1173,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         adapter = self._get_tts_adapter()
         if adapter is not None:
+            adapter.normalize(request)
             return adapter.validate(request)
 
         adapter_cls = resolve_adapter("qwen3_tts")
@@ -1183,7 +1184,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             server=self,
             engine_client=self.engine_client,
         )
-        return adapter_cls(ctx).validate(request)
+        adapter = adapter_cls(ctx)
+        adapter.normalize(request)
+        return adapter.validate(request)
 
     def _validate_speech_sample_rate(self, request: OpenAICreateSpeechRequest) -> str | None:
         if request.sample_rate is None:
@@ -2000,6 +2003,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         model_type: str | None = None
         has_inline_ref_audio = (request.ref_audio is not None) if has_inline_ref_audio is None else has_inline_ref_audio
         if (adapter := self._get_tts_adapter()) is not None:
+            adapter.normalize(request)
             validation_error = adapter.validate(request)
             if validation_error:
                 raise ValueError(validation_error)

vllm_omni/entrypoints/openai/tts_adapters/__init__.py

@@ -105,6 +105,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import (  # noqa: E402,F401
     audex,
     audex_tta,
     audio8_tts,
+    auk,
     breeze_tts_2,
     cosyvoice3,
     covo_audio,

vllm_omni/entrypoints/openai/tts_adapters/auk.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""AuK Speech API adapter for the current multi-stage serving contract."""
+
+from __future__ import annotations
+
+import copy
+import json
+import math
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from vllm.logger import init_logger
+
+from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter
+from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest
+from vllm_omni.model_extras.auk import DEFAULT_SWAY, auk_prompt
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
+
+
+@register_tts_adapter
+class AuKAdapter(ARTTSAdapter):
+    """Translate Speech API fields into the existing AuK encoder/DiT prompt."""
+
+    name = "auk"
+    # AuK's entry stage uses the generic ``model_stage="encoder"``, which other
+    # pipelines (MiniMax H3) also declare, so both stage discovery and
+    # model-type detection go through the architecture instead.
+    model_archs = frozenset({"AuKForConditionalGeneration"})
+    arch_identifies_entry_stage = True
+
+    def normalize(self, request: OpenAICreateSpeechRequest) -> None:
+        if request.task_type is None:
+            return
+        if request.task_type in ("CustomVoice", "Base"):
+            logger.warning_once(
+                "AuK `task_type` is a compatibility shortcut; prefer a complete `instructions` prompt. "
+                'CustomVoice applies `Say the following: "<input>"`; Base applies '
+                '`Say the following with the same voice: "<input>"`.'
+            )
+        if request.instructions and request.instructions.strip():
+            logger.warning(
+                "AuK request contains both `task_type` and a complete `instructions` prompt; "
+                "using `instructions` without applying another task template."
+            )
+            request.input = ""
+            request.task_type = None
+            return
+        instruction = request.input.strip()
+        if instruction.startswith(("Say the following:", "Say the following with the same voice:")):
+            logger.warning(
+                "AuK request `input` already contains a complete task instruction; "
+                "using it without applying another task template."
+            )
+            request.instructions = instruction
+            request.input = ""
+            request.task_type = None
+            return
+        if not instruction:
+            return
+        quoted_instruction = json.dumps(instruction, ensure_ascii=False)
+        if request.task_type == "Base":
+            request.instructions = f"Say the following with the same voice: {quoted_instruction}"
+        elif request.task_type == "CustomVoice":
+            request.instructions = f"Say the following: {quoted_instruction}"
+        else:
+            return
+        request.input = ""
+        # task_type is a compatibility shortcut for AuK, not a model-side
+        # parameter. Consuming it makes normalization idempotent when the
+        # shared batch path validates and later prepares the same request.
+        request.task_type = None
+
+    def validate(self, request: OpenAICreateSpeechRequest) -> str | None:
+        if not request.input.strip() and not (request.instructions and request.instructions.strip()):
+            return "AuK requires input text or a complete instructions prompt"
+        if request.duration_seconds is not None:
+            if not math.isfinite(request.duration_seconds) or request.duration_seconds <= 0:
+                return "AuK requires duration_seconds to be a positive finite value"
+        if request.duration_seconds is None and request.ref_audio is None:
+            return "AuK requires duration_seconds for requests without ref_audio"
+        if request.instructions is not None:
+            if len(request.instructions) > self.ctx.server._max_instructions_length:
+                return f"instructions exceeds max length {self.ctx.server._max_instructions_length}"
+        if request.ref_audio is not None:
+            if not isinstance(request.ref_audio, str):
+                return "AuK supports one ref_audio string per request"
+            error = self.ctx.server._validate_ref_audio_format(request.ref_audio)
+            if error:
+                return error
+        extra = request.extra_params or {}
+        for key in extra:
+            if key not in {"num_inference_steps", "guidance_scale", "sway", "t_grid", "vae_sample"}:
+                return f"AuK does not support extra_params.{key}"
+        t_grid = extra.get("t_grid")
+        if t_grid is not None:
+            t_grid_error = "AuK extra_params.t_grid must contain at least two finite, strictly increasing values"
+            if not isinstance(t_grid, (list, tuple)) or len(t_grid) < 2:
+                return t_grid_error
+            try:
+                values = [float(value) for value in t_grid]
+            except (TypeError, ValueError):
+                return t_grid_error
+            if not all(math.isfinite(value) for value in values) or any(
+                right <= left for left, right in zip(values, values[1:])
+            ):
+                return t_grid_error
+        return None
+
+    async def build(
+        self,
+        request: OpenAICreateSpeechRequest,
+        sampling_params_list: list,
+        has_inline_ref_audio: bool,
+    ) -> PreparedRequest:
+        del sampling_params_list, has_inline_ref_audio
+        reference: tuple[np.ndarray, int] | None = None
+        if request.ref_audio is not None:
+            samples, sample_rate, _ = await self.ctx.server._resolve_ref_audio(request.ref_audio)
+            reference = (np.asarray(samples, dtype=np.float32), int(sample_rate))
+
+        extra = request.extra_params or {}
+        # The caller supplies the complete cookbook prompt in instructions;
+        # this adapter deliberately does not invent or compose a template.
+        if request.instructions and request.instructions.strip():
+            if request.input.strip():
+                logger.warning(
+                    "AuK request contains both non-empty `input` and `instructions`; "
+                    "using the complete `instructions` prompt and ignoring `input`."
+                )
+            instruction = request.instructions.strip()
+        else:
+            logger.warning(
+                "AuK request only provides `input`; treating it as the complete AuK instruction. "
+                'Prefer `input=""` with the instruction in `instructions`.'
+            )
+            instruction = request.input.strip()
+        prompt = auk_prompt(
+            instruction,
+            reference,
+            gen_seconds=request.duration_seconds,
+            sway=float(extra.get("sway", DEFAULT_SWAY)),
+            t_grid=extra.get("t_grid"),
+            vae_sample=extra.get("vae_sample", False),
+        )
+        return PreparedRequest(prompt=prompt, tts_params={}, model_type=self.name)
+
+    def apply_sampling_overrides(
+        self,
+        sampling_params_list: list,
+        request: OpenAICreateSpeechRequest,
+        prompt: dict[str, Any] | None = None,
+        request_id: str | None = None,
+    ) -> list:
+        if len(sampling_params_list) < 2:
+            return sampling_params_list
+        updated = copy.deepcopy(sampling_params_list)
+        stage1 = updated[1]
+        extra = request.extra_params or {}
+        if "num_inference_steps" in extra:
+            stage1.num_inference_steps = int(extra["num_inference_steps"])
+        if "guidance_scale" in extra:
+            stage1.guidance_scale = float(extra["guidance_scale"])
+        if request.seed is not None:
+            stage1.seed = int(request.seed)
+        return updated
+
+
+__all__ = ["AuKAdapter"]